@genkit-ai/vertexai/lib/rerankers/reranker.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 reranker_exports = {};
__export(reranker_exports, {
  vertexAiRerankerRef: () => vertexAiRerankerRef,
  vertexAiRerankers: () => vertexAiRerankers
});
module.exports = __toCommonJS(reranker_exports);
var import_reranker = require("genkit/reranker");
var import_constants = require("./constants.js");
var import_types = require("./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: import_types.VertexAIRerankerOptionsSchema.optional()
      },
      async (query, documents, _options) => {
        const response = await client.request({
          method: "POST",
          url: (0, import_constants.getRerankEndpoint)(projectId, options.location ?? "us-central1"),
          data: {
            model: rerankOption.model || import_constants.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 import_reranker.RankedDocument({
            content: doc.content,
            metadata: {
              ...doc.metadata,
              score: record.score
            }
          });
        });
        return { documents: rankedDocuments };
      }
    );
  }
}
const vertexAiRerankerRef = (params) => {
  return (0, import_reranker.rerankerRef)({
    name: `vertexai/${params.rerankerName}`,
    info: {
      label: params.displayName ?? `Vertex AI Reranker`
    },
    configSchema: import_types.VertexAIRerankerOptionsSchema.optional()
  });
};
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  vertexAiRerankerRef,
  vertexAiRerankers
});
//# sourceMappingURL=reranker.js.map