@genkit-ai/vertexai/lib/vectorsearch/vector_search/indexers.mjs

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

import { indexerRef } from "genkit/retriever";
import {
  Datapoint,
  VertexAIVectorIndexerOptionsSchema
} from "./types";
import { upsertDatapoints } from "./upsert_datapoints";
const vertexAiIndexerRef = (params) => {
  return indexerRef({
    name: `vertexai/${params.indexId}`,
    info: {
      label: params.displayName ?? `Vertex AI - ${params.indexId}`
    },
    configSchema: VertexAIVectorIndexerOptionsSchema.optional()
  });
};
function vertexAiIndexers(ai, params) {
  const vectorSearchOptions = params.pluginOptions.vectorSearchOptions;
  const indexers = [];
  if (!vectorSearchOptions || vectorSearchOptions.length === 0) {
    return indexers;
  }
  for (const vectorSearchOption of vectorSearchOptions) {
    const { documentIndexer, indexId } = vectorSearchOption;
    const embedderReference = vectorSearchOption.embedder ?? params.defaultEmbedder;
    if (!embedderReference) {
      throw new Error(
        "Embedder reference is required to define Vertex AI retriever"
      );
    }
    const embedderOptions = vectorSearchOption.embedderOptions;
    const indexer = ai.defineIndexer(
      {
        name: `vertexai/${indexId}`,
        configSchema: VertexAIVectorIndexerOptionsSchema.optional()
      },
      async (docs, options) => {
        let docIds = [];
        try {
          docIds = await documentIndexer(docs, options);
        } catch (error) {
          throw new Error(
            `Error storing your document content/metadata: ${error}`
          );
        }
        const embeddings = await ai.embedMany({
          embedder: embedderReference,
          content: docs,
          options: embedderOptions
        });
        const datapoints = embeddings.map(({ embedding }, i) => {
          const dp = new Datapoint({
            datapointId: docIds[i],
            featureVector: embedding
          });
          if (docs[i].metadata?.restricts) {
            dp.restricts = docs[i].metadata?.restricts;
          }
          if (docs[i].metadata?.numericRestricts) {
            dp.numericRestricts = docs[i].metadata?.numericRestricts;
          }
          if (docs[i].metadata?.crowdingTag) {
            dp.crowdingTag = docs[i].metadata?.crowdingTag;
          }
          return dp;
        });
        try {
          await upsertDatapoints({
            datapoints,
            authClient: params.authClient,
            projectId: params.pluginOptions.projectId,
            location: params.pluginOptions.location,
            indexId
          });
        } catch (error) {
          throw error;
        }
      }
    );
    indexers.push(indexer);
  }
  return indexers;
}
export {
  vertexAiIndexerRef,
  vertexAiIndexers
};
//# sourceMappingURL=indexers.mjs.map