@genkit-ai/vertexai/lib/rerankers/reranker.mjs

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

import { RankedDocument, rerankerRef } from "genkit/reranker";
import { DEFAULT_MODEL, getRerankEndpoint } from "./constants.js";
import { VertexAIRerankerOptionsSchema } from "./types.js";
async function vertexAiRerankers(ai, options) {
  const rerankOptions = options.rerankOptions;
  if (rerankOptions.length === 0) {
    throw new Error("Provide at least one reranker configuration.");
  }
  const auth = options.authClient;
  const client = await auth.getClient();
  const projectId = options.projectId;
  for (const rerankOption of rerankOptions) {
    if (!rerankOption.name && !rerankOption.model) {
      throw new Error("At least one of name or model must be provided.");
    }
    ai.defineReranker(
      {
        name: `vertexai/${rerankOption.name || rerankOption.model}`,
        configSchema: VertexAIRerankerOptionsSchema.optional()
      },
      async (query, documents, _options) => {
        const response = await client.request({
          method: "POST",
          url: getRerankEndpoint(projectId, options.location ?? "us-central1"),
          data: {
            model: rerankOption.model || DEFAULT_MODEL,
            // Use model from config or default
            query: query.text,
            records: documents.map((doc, idx) => ({
              id: `${idx}`,
              content: doc.text
            }))
          }
        });
        const rankedDocuments = response.data.records.map((record) => {
          const doc = documents[record.id];
          return new RankedDocument({
            content: doc.content,
            metadata: {
              ...doc.metadata,
              score: record.score
            }
          });
        });
        return { documents: rankedDocuments };
      }
    );
  }
}
const vertexAiRerankerRef = (params) => {
  return rerankerRef({
    name: `vertexai/${params.rerankerName}`,
    info: {
      label: params.displayName ?? `Vertex AI Reranker`
    },
    configSchema: VertexAIRerankerOptionsSchema.optional()
  });
};
export {
  vertexAiRerankerRef,
  vertexAiRerankers
};
//# sourceMappingURL=reranker.mjs.map