@genkit-ai/ai/lib/reranker.js

Genkit AI framework generative AI APIs.

"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, {
  CommonRerankerOptionsSchema: () => CommonRerankerOptionsSchema,
  RankedDocument: () => RankedDocument,
  RankedDocumentDataSchema: () => RankedDocumentDataSchema,
  RankedDocumentMetadataSchema: () => RankedDocumentMetadataSchema,
  RerankerInfoSchema: () => RerankerInfoSchema,
  defineReranker: () => defineReranker,
  rerank: () => rerank,
  reranker: () => reranker,
  rerankerRef: () => rerankerRef
});
module.exports = __toCommonJS(reranker_exports);
var import_core = require("@genkit-ai/core");
var import_schema = require("@genkit-ai/core/schema");
var import_parts = require("./parts.js");
var import_retriever = require("./retriever.js");
const RankedDocumentMetadataSchema = import_core.z.object({
  score: import_core.z.number()
  // Enforces that 'score' must be a number
}).passthrough();
const RankedDocumentDataSchema = import_core.z.object({
  content: import_core.z.array(import_parts.PartSchema),
  metadata: RankedDocumentMetadataSchema
});
class RankedDocument extends import_retriever.Document {
  content;
  metadata;
  constructor(data) {
    super(data);
    this.content = data.content;
    this.metadata = data.metadata;
  }
  /**
   * Returns the score of the document.
   * @returns The score of the document.
   */
  score() {
    return this.metadata.score;
  }
}
const RerankerRequestSchema = import_core.z.object({
  query: import_retriever.DocumentDataSchema,
  documents: import_core.z.array(import_retriever.DocumentDataSchema),
  options: import_core.z.any().optional()
});
const RerankerResponseSchema = import_core.z.object({
  documents: import_core.z.array(RankedDocumentDataSchema)
});
const RerankerInfoSchema = import_core.z.object({
  label: import_core.z.string().optional(),
  /** Supported model capabilities. */
  supports: import_core.z.object({
    /** Model can process media as part of the prompt (multimodal input). */
    media: import_core.z.boolean().optional()
  }).optional()
});
function rerankerWithMetadata(reranker2, configSchema) {
  const withMeta = reranker2;
  withMeta.__configSchema = configSchema;
  return withMeta;
}
function defineReranker(registry, options, runner) {
  const act = reranker(options, runner);
  registry.registerAction("reranker", act);
  return act;
}
function reranker(options, runner) {
  const reranker2 = (0, import_core.action)(
    {
      actionType: "reranker",
      name: options.name,
      inputSchema: options.configSchema ? RerankerRequestSchema.extend({
        options: options.configSchema.optional()
      }) : RerankerRequestSchema,
      outputSchema: RerankerResponseSchema,
      metadata: {
        type: "reranker",
        info: options.info,
        reranker: {
          customOptions: options.configSchema ? (0, import_schema.toJsonSchema)({ schema: options.configSchema }) : void 0
        }
      }
    },
    (i) => runner(
      new import_retriever.Document(i.query),
      i.documents.map((d) => new import_retriever.Document(d)),
      i.options
    )
  );
  const rwm = rerankerWithMetadata(
    reranker2,
    options.configSchema
  );
  return rwm;
}
async function rerank(registry, params) {
  let reranker2;
  if (typeof params.reranker === "string") {
    reranker2 = await registry.lookupAction(`/reranker/${params.reranker}`);
  } else if (Object.hasOwnProperty.call(params.reranker, "info")) {
    reranker2 = await registry.lookupAction(`/reranker/${params.reranker.name}`);
  } else {
    reranker2 = params.reranker;
  }
  if (!reranker2) {
    throw new Error("Unable to resolve the reranker");
  }
  const response = await reranker2({
    query: typeof params.query === "string" ? import_retriever.Document.fromText(params.query) : params.query,
    documents: params.documents,
    options: params.options
  });
  return response.documents.map((d) => new RankedDocument(d));
}
const CommonRerankerOptionsSchema = import_core.z.object({
  k: import_core.z.number().describe("Number of documents to rerank").optional()
});
function rerankerRef(options) {
  return { ...options };
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  CommonRerankerOptionsSchema,
  RankedDocument,
  RankedDocumentDataSchema,
  RankedDocumentMetadataSchema,
  RerankerInfoSchema,
  defineReranker,
  rerank,
  reranker,
  rerankerRef
});
//# sourceMappingURL=reranker.js.map