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

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

import {
  z
} from "genkit";
import { reranker as pluginReranker } from "genkit/plugin";
import { RankedDocument, rerankerRef } from "genkit/reranker";
import { checkModelName } from "../../common/utils.mjs";
import { rerankerRank } from "./client.mjs";
const VertexRerankerConfigSchema = z.object({
  topN: z.number().optional().describe("Number of top documents to rerank"),
  // Optional: Number of documents to rerank
  ignoreRecordDetailsInResponse: 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: 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 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 = checkModelName(version);
  return 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 pluginReranker(
    {
      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: checkModelName(ref.name),
        query: query.text,
        records: documents.map(toRerankerDoc),
        ...options
      };
      const response = await 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 RankedDocument({
      content: doc.content,
      metadata: {
        ...doc.metadata,
        score: record.score
      }
    });
  });
}
export {
  DEFAULT_MODEL_NAME,
  GENERIC_MODEL,
  KNOWN_MODELS,
  VertexRerankerConfigSchema,
  defineReranker,
  isRerankerModelName,
  listKnownRerankers,
  reranker
};
//# sourceMappingURL=reranker.mjs.map