@genkit-ai/vertexai/lib/rerankers/v2/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, {
  DEFAULT_MODEL_NAME: () => DEFAULT_MODEL_NAME,
  GENERIC_MODEL: () => GENERIC_MODEL,
  KNOWN_MODELS: () => KNOWN_MODELS,
  VertexRerankerConfigSchema: () => VertexRerankerConfigSchema,
  defineReranker: () => defineReranker,
  isRerankerModelName: () => isRerankerModelName,
  listKnownRerankers: () => listKnownRerankers,
  reranker: () => reranker
});
module.exports = __toCommonJS(reranker_exports);
var import_genkit = require("genkit");
var import_plugin = require("genkit/plugin");
var import_reranker = require("genkit/reranker");
var import_utils = require("../../common/utils.js");
var import_client = require("./client.js");
const VertexRerankerConfigSchema = import_genkit.z.object({
  topN: import_genkit.z.number().optional().describe("Number of top documents to rerank"),
  // Optional: Number of documents to rerank
  ignoreRecordDetailsInResponse: import_genkit.z.boolean().optional().describe(
    "If true, the response will contain only record ID and score. By default, it is false, the response will contain record details."
  ),
  location: import_genkit.z.string().optional().describe('Google Cloud location, e.g., "us-central1"')
  // Optional: Location of the reranking model
}).passthrough();
function commonRef(name, info, configSchema = VertexRerankerConfigSchema) {
  return (0, import_reranker.rerankerRef)({
    name: `vertex-rerankers/${name}`,
    configSchema,
    info: info ?? {
      supports: {
        media: false
      }
    }
  });
}
const GENERIC_MODEL = commonRef("reranker");
const DEFAULT_MODEL_NAME = "semantic-ranker-default@latest";
const KNOWN_MODELS = {
  "semantic-ranker-default@latest": commonRef("semantic-ranker-default@latest"),
  "semantic-ranker-default-004": commonRef("semantic-ranker-default-004"),
  "semantic-ranker-fast-004": commonRef("semantic-ranker-fast-004"),
  "semantic-ranker-default-003": commonRef("semantic-ranker-default-003"),
  "semantic-ranker-default-002": commonRef("semantic-ranker-default-002")
};
function isRerankerModelName(value) {
  return !!value?.startsWith("semantic-ranker-");
}
function reranker(version, options = {}) {
  const name = (0, import_utils.checkModelName)(version);
  return (0, import_reranker.rerankerRef)({
    name: `vertex-rerankers/${name}`,
    configSchema: VertexRerankerConfigSchema,
    info: {
      ...GENERIC_MODEL.info
    }
  });
}
function listKnownRerankers(clientOptions) {
  return Object.keys(KNOWN_MODELS).map(
    (name) => defineReranker(name, clientOptions)
  );
}
function defineReranker(name, clientOptions) {
  const ref = reranker(name);
  return (0, import_plugin.reranker)(
    {
      name: ref.name,
      ...ref.info,
      configSchema: ref.configSchema
    },
    async (query, documents, options) => {
      const rerankRequest = {
        // Note that model silently falls back to default if it's not recognized
        // This happens in the vertexai reranker service backend.
        model: (0, import_utils.checkModelName)(ref.name),
        query: query.text,
        records: documents.map(toRerankerDoc),
        ...options
      };
      const response = await (0, import_client.rerankerRank)(
        ref.name,
        rerankRequest,
        clientOptions
      );
      return { documents: fromRerankResponse(response, documents) };
    }
  );
}
function toRerankerDoc(doc, idx) {
  return {
    id: `${idx}`,
    content: doc.text
  };
}
function fromRerankResponse(response, documents) {
  return response.records.map((record) => {
    const doc = documents[record.id];
    return new import_reranker.RankedDocument({
      content: doc.content,
      metadata: {
        ...doc.metadata,
        score: record.score
      }
    });
  });
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  DEFAULT_MODEL_NAME,
  GENERIC_MODEL,
  KNOWN_MODELS,
  VertexRerankerConfigSchema,
  defineReranker,
  isRerankerModelName,
  listKnownRerankers,
  reranker
});
//# sourceMappingURL=reranker.js.map