@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
140 lines • 4.72 kB
JavaScript
;
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,
RerankerInfoSchema: () => RerankerInfoSchema,
defineReranker: () => defineReranker,
rerank: () => rerank,
rerankerRef: () => rerankerRef
});
module.exports = __toCommonJS(reranker_exports);
var import_core = require("@genkit-ai/core");
var import_document = require("./document.js");
var import_retriever = require("./retriever.js");
const RankedDocumentDataSchema = import_core.z.object({
content: import_core.z.array(import_document.PartSchema),
metadata: import_core.z.object({
score: import_core.z.number()
// Enforces that 'score' must be a number
}).passthrough()
// Allows other properties in 'metadata' with any type
});
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(reranker, configSchema) {
const withMeta = reranker;
withMeta.__configSchema = configSchema;
return withMeta;
}
function defineReranker(registry, options, runner) {
const reranker = (0, import_core.defineAction)(
registry,
{
actionType: "reranker",
name: options.name,
inputSchema: options.configSchema ? RerankerRequestSchema.extend({
options: options.configSchema.optional()
}) : RerankerRequestSchema,
outputSchema: RerankerResponseSchema,
metadata: {
type: "reranker",
info: options.info
}
},
(i) => runner(
new import_retriever.Document(i.query),
i.documents.map((d) => new import_retriever.Document(d)),
i.options
)
);
const rwm = rerankerWithMetadata(
reranker,
options.configSchema
);
return rwm;
}
async function rerank(registry, params) {
let reranker;
if (typeof params.reranker === "string") {
reranker = await registry.lookupAction(`/reranker/${params.reranker}`);
} else if (Object.hasOwnProperty.call(params.reranker, "info")) {
reranker = await registry.lookupAction(`/reranker/${params.reranker.name}`);
} else {
reranker = params.reranker;
}
if (!reranker) {
throw new Error("Unable to resolve the reranker");
}
const response = await reranker({
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,
RerankerInfoSchema,
defineReranker,
rerank,
rerankerRef
});
//# sourceMappingURL=reranker.js.map