Genkit MCP
Official
by firebase
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { Genkit } from 'genkit';
import { RankedDocument, rerankerRef } from 'genkit/reranker';
import { DEFAULT_MODEL, getRerankEndpoint } from './constants.js';
import { VertexAIRerankerOptionsSchema, VertexRerankOptions } from './types.js';
/**
* Creates Vertex AI rerankers.
*
* This function creates and registers rerankers for the specified models.
*
* @param {VertexRerankOptions<EmbedderCustomOptions>} options - The parameters for creating the rerankers.
* @returns {Promise<void>}
*/
export async function vertexAiRerankers(
ai: Genkit,
options: VertexRerankOptions
): Promise<void> {
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: VertexAIRerankerOptionsSchema.optional(),
},
async (query, documents, _options) => {
const response = await client.request({
method: 'POST',
url: getRerankEndpoint(projectId, options.location ?? 'us-central1'),
data: {
model: rerankOption.model || DEFAULT_MODEL, // Use model from config or default
query: query.text,
records: documents.map((doc, idx) => ({
id: `${idx}`,
content: doc.text,
})),
},
});
const rankedDocuments: RankedDocument[] = (
response.data as any
).records.map((record: any) => {
const doc = documents[record.id];
return new RankedDocument({
content: doc.content,
metadata: {
...doc.metadata,
score: record.score,
},
});
});
return { documents: rankedDocuments };
}
);
}
}
/**
* Creates a reference to a Vertex AI reranker.
*
* @param {Object} params - The parameters for the reranker reference.
* @param {string} [params.displayName] - An optional display name for the reranker.
* @returns {Object} - The reranker reference object.
*/
export const vertexAiRerankerRef = (params: {
rerankerName: string;
displayName?: string;
}) => {
return rerankerRef({
name: `vertexai/${params.rerankerName}`,
info: {
label: params.displayName ?? `Vertex AI Reranker`,
},
configSchema: VertexAIRerankerOptionsSchema.optional(),
});
};