@genkit-ai/vertexai/lib/vectorsearch/vector_search/retrievers.js

Genkit AI framework plugin for Google Cloud Vertex AI APIs including Gemini APIs, Imagen, and more.

"use strict";
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
  for (var name in all)
    __defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
  if (from && typeof from === "object" || typeof from === "function") {
    for (let key of __getOwnPropNames(from))
      if (!__hasOwnProp.call(to, key) && key !== except)
        __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
  }
  return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var retrievers_exports = {};
__export(retrievers_exports, {
  vertexAiRetrieverRef: () => vertexAiRetrieverRef,
  vertexAiRetrievers: () => vertexAiRetrievers
});
module.exports = __toCommonJS(retrievers_exports);
var import_genkit = require("genkit");
var import_query_public_endpoint = require("./query_public_endpoint");
var import_types = require("./types");
var import_utils = require("./utils");
const DEFAULT_K = 10;
function vertexAiRetrievers(ai, params) {
  const vectorSearchOptions = params.pluginOptions.vectorSearchOptions;
  const defaultEmbedder = params.defaultEmbedder;
  const retrievers = [];
  if (!vectorSearchOptions || vectorSearchOptions.length === 0) {
    return retrievers;
  }
  for (const vectorSearchOption of vectorSearchOptions) {
    const { documentRetriever, indexId, publicDomainName } = vectorSearchOption;
    const embedderOptions = vectorSearchOption.embedderOptions;
    const retriever = ai.defineRetriever(
      {
        name: `vertexai/${indexId}`,
        configSchema: import_types.VertexAIVectorRetrieverOptionsSchema.optional()
      },
      async (content, options) => {
        const embedderReference = vectorSearchOption.embedder ?? defaultEmbedder;
        if (!embedderReference) {
          throw new Error(
            "Embedder reference is required to define Vertex AI retriever"
          );
        }
        const queryEmbedding = (await ai.embed({
          embedder: embedderReference,
          options: embedderOptions,
          content
        }))[0].embedding;
        const accessToken = await params.authClient.getAccessToken();
        if (!accessToken) {
          throw new Error(
            "Error generating access token when defining Vertex AI retriever"
          );
        }
        const projectId = params.pluginOptions.projectId;
        if (!projectId) {
          throw new Error(
            "Project ID is required to define Vertex AI retriever"
          );
        }
        const projectNumber = await (0, import_utils.getProjectNumber)(projectId);
        const location = params.pluginOptions.location;
        if (!location) {
          throw new Error("Location is required to define Vertex AI retriever");
        }
        let res = await (0, import_query_public_endpoint.queryPublicEndpoint)({
          featureVector: queryEmbedding,
          neighborCount: options?.k || DEFAULT_K,
          accessToken,
          projectId,
          location,
          publicDomainName,
          projectNumber,
          indexEndpointId: vectorSearchOption.indexEndpointId,
          deployedIndexId: vectorSearchOption.deployedIndexId,
          restricts: content.metadata?.restricts,
          numericRestricts: content.metadata?.numericRestricts
        });
        const nearestNeighbors = res.nearestNeighbors;
        const queryRes = nearestNeighbors ? nearestNeighbors[0] : null;
        const neighbors = queryRes ? queryRes.neighbors : null;
        if (!neighbors) {
          return { documents: [] };
        }
        const documents = await documentRetriever(neighbors, options);
        return { documents };
      }
    );
    retrievers.push(retriever);
  }
  return retrievers;
}
const vertexAiRetrieverRef = (params) => {
  return (0, import_genkit.retrieverRef)({
    name: `vertexai/${params.indexId}`,
    info: {
      label: params.displayName ?? `ertex AI - ${params.indexId}`
    },
    configSchema: import_types.VertexAIVectorRetrieverOptionsSchema.optional()
  });
};
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  vertexAiRetrieverRef,
  vertexAiRetrievers
});
//# sourceMappingURL=retrievers.js.map