diff --git a/docs/rag.md b/docs/rag.md index 4d07708b62..02f55fde29 100644 --- a/docs/rag.md +++ b/docs/rag.md @@ -384,3 +384,89 @@ const docs = await retrieve({ options: { preRerankK: 7, k: 3 }, }); ``` + +### Rerankers and Two-Stage Retrieval + +A reranking model — also known as a cross-encoder — is a type of model that, given a query and document, will output a similarity score. We use this score to reorder the documents by relevance to our query. Reranker APIs take a list of documents (for example the output of a retriever) and reorders the documents based on their relevance to the query. This step can be useful for fine-tuning the results and ensuring that the most pertinent information is used in the prompt provided to a generative model. + + +#### Reranker Example + +A reranker in Genkit is defined in a similar syntax to retrievers and indexers. Here is an example using a reranker in Genkit. This flow reranks a set of documents based on their relevance to the provided query using a predefined Vertex AI reranker. + +```ts +import { rerank } from '@genkit-ai/ai/reranker'; +import { Document } from '@genkit-ai/ai/retriever'; +import { defineFlow } from '@genkit-ai/flow'; +import * as z from 'zod'; + +const FAKE_DOCUMENT_CONTENT = [ + 'pythagorean theorem', + 'e=mc^2', + 'pi', + 'dinosaurs', + 'quantum mechanics', + 'pizza', + 'harry potter', +]; + +export const rerankFlow = defineFlow( + { + name: 'rerankFlow', + inputSchema: z.object({ query: z.string() }), + outputSchema: z.array( + z.object({ + text: z.string(), + score: z.number(), + }) + ), + }, + async ({ query }) => { + const documents = FAKE_DOCUMENT_CONTENT.map((text) => + Document.fromText(text) + ); + + const rerankedDocuments = await rerank({ + reranker: 'vertexai/semantic-ranker-512', + query: Document.fromText(query), + documents, + }); + + return rerankedDocuments.map((doc) => ({ + text: doc.text(), + score: doc.metadata.score, + })); + } +); +``` +This reranker uses the Vertex AI genkit plugin with `semantic-ranker-512` to score and rank documents. The higher the score, the more relevant the document is to the query. + +#### Custom Rerankers + +You can also define custom rerankers to suit your specific use case. This is helpful when you need to rerank documents using your own custom logic or a custom model. Here’s a simple example of defining a custom reranker: +```typescript +import { defineReranker } from '@genkit-ai/ai/reranker'; +import * as z from 'zod'; + +export const customReranker = defineReranker( + { + name: 'custom/reranker', + configSchema: z.object({ + k: z.number().optional(), + }), + }, + async (query, documents, options) => { + // Your custom reranking logic here + const rerankedDocs = documents.map((doc) => { + const score = Math.random(); // Assign random scores for demonstration + return { + ...doc, + metadata: { ...doc.metadata, score }, + }; + }); + + return rerankedDocs.sort((a, b) => b.metadata.score - a.metadata.score).slice(0, options.k || 3); + } +); +``` +Once defined, this custom reranker can be used just like any other reranker in your RAG flows, giving you flexibility to implement advanced reranking strategies. \ No newline at end of file diff --git a/js/ai/package.json b/js/ai/package.json index 4b38e7fa72..c9980ae60a 100644 --- a/js/ai/package.json +++ b/js/ai/package.json @@ -90,6 +90,12 @@ "require": "./lib/tool.js", "import": "./lib/tool.mjs", "default": "./lib/tool.js" + }, + "./reranker": { + "types": "./lib/reranker.d.ts", + "require": "./lib/reranker.js", + "import": "./lib/reranker.mjs", + "default": "./lib/reranker.js" } }, "typesVersions": { @@ -114,6 +120,9 @@ ], "tool": [ "lib/tool" + ], + "reranker": [ + "lib/reranker" ] } } diff --git a/js/ai/src/reranker.ts b/js/ai/src/reranker.ts new file mode 100644 index 0000000000..9e6ea10fee --- /dev/null +++ b/js/ai/src/reranker.ts @@ -0,0 +1,205 @@ +/** + * 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 { Action, defineAction } from '@genkit-ai/core'; +import { lookupAction } from '@genkit-ai/core/registry'; +import * as z from 'zod'; +import { Part, PartSchema } from './document.js'; +import { Document, DocumentData, DocumentDataSchema } from './retriever.js'; + +type RerankerFn = ( + query: Document, + documents: Document[], + queryOpts: z.infer +) => Promise; + +export const RankedDocumentDataSchema = z.object({ + content: z.array(PartSchema), + metadata: z + .object({ + score: z.number(), // Enforces that 'score' must be a number + }) + .passthrough(), // Allows other properties in 'metadata' with any type +}); + +export type RankedDocumentData = z.infer; + +export class RankedDocument extends Document implements RankedDocumentData { + content: Part[]; + metadata: { score: number } & Record; + + constructor(data: RankedDocumentData) { + super(data); + this.content = data.content; + this.metadata = data.metadata; + } + /** + * Returns the score of the document. + * @returns The score of the document. + */ + score(): number { + return this.metadata.score; + } +} + +const RerankerRequestSchema = z.object({ + query: DocumentDataSchema, + documents: z.array(DocumentDataSchema), + options: z.any().optional(), +}); + +const RerankerResponseSchema = z.object({ + documents: z.array(RankedDocumentDataSchema), +}); +type RerankerResponse = z.infer; + +export const RerankerInfoSchema = z.object({ + label: z.string().optional(), + /** Supported model capabilities. */ + supports: z + .object({ + /** Model can process media as part of the prompt (multimodal input). */ + media: z.boolean().optional(), + }) + .optional(), +}); +export type RerankerInfo = z.infer; + +export type RerankerAction = + Action< + typeof RerankerRequestSchema, + typeof RerankerResponseSchema, + { model: RerankerInfo } + > & { + __configSchema?: CustomOptions; + }; + +function rerankerWithMetadata< + RerankerOptions extends z.ZodTypeAny = z.ZodTypeAny, +>( + reranker: Action, + configSchema?: RerankerOptions +): RerankerAction { + const withMeta = reranker as RerankerAction; + withMeta.__configSchema = configSchema; + return withMeta; +} + +/** + * Creates a reranker action for the provided {@link RerankerFn} implementation. + */ +export function defineReranker( + options: { + name: string; + configSchema?: OptionsType; + info?: RerankerInfo; + }, + runner: RerankerFn +) { + const reranker = defineAction( + { + 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 Document(i.query), + i.documents.map((d) => new Document(d)), + i.options + ) + ); + const rwm = rerankerWithMetadata( + reranker as Action< + typeof RerankerRequestSchema, + typeof RerankerResponseSchema + >, + options.configSchema + ); + return rwm; +} + +export interface RerankerParams< + CustomOptions extends z.ZodTypeAny = z.ZodTypeAny, +> { + reranker: RerankerArgument; + query: string | DocumentData; + documents: DocumentData[]; + options?: z.infer; +} + +export type RerankerArgument< + CustomOptions extends z.ZodTypeAny = z.ZodTypeAny, +> = RerankerAction | RerankerReference | string; + +/** + * Reranks documents from a {@link RerankerArgument} based on the provided query. + */ +export async function rerank( + params: RerankerParams +): Promise> { + let reranker: RerankerAction; + if (typeof params.reranker === 'string') { + reranker = await lookupAction(`/reranker/${params.reranker}`); + } else if (Object.hasOwnProperty.call(params.reranker, 'info')) { + reranker = await lookupAction(`/reranker/${params.reranker.name}`); + } else { + reranker = params.reranker as RerankerAction; + } + if (!reranker) { + throw new Error('Unable to resolve the reranker'); + } + const response = await reranker({ + query: + typeof params.query === 'string' + ? Document.fromText(params.query) + : params.query, + documents: params.documents, + options: params.options, + }); + + return response.documents.map((d) => new RankedDocument(d)); +} + +export const CommonRerankerOptionsSchema = z.object({ + k: z.number().describe('Number of documents to rerank').optional(), +}); + +export interface RerankerReference { + name: string; + configSchema?: CustomOptions; + info?: RerankerInfo; +} + +/** + * Helper method to configure a {@link RerankerReference} to a plugin. + */ +export function rerankerRef< + CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny, +>( + options: RerankerReference +): RerankerReference { + return { ...options }; +} diff --git a/js/ai/tests/reranker/reranker_test.ts b/js/ai/tests/reranker/reranker_test.ts new file mode 100644 index 0000000000..0226264e0f --- /dev/null +++ b/js/ai/tests/reranker/reranker_test.ts @@ -0,0 +1,221 @@ +/** + * 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 { GenkitError } from '@genkit-ai/core'; +import assert from 'node:assert'; +import { describe, it } from 'node:test'; +import * as z from 'zod'; +import { defineReranker, rerank } from '../../src/reranker'; +import { Document } from '../../src/retriever'; + +describe('reranker', () => { + describe('defineReranker()', () => { + it('reranks documents based on custom logic', async () => { + const customReranker = defineReranker( + { + name: 'custom/reranker', + configSchema: z.object({ + k: z.number().optional(), + }), + }, + async (query, documents, options) => { + // Custom reranking logic: score based on string length similarity to query + const queryLength = query.text().length; + const rerankedDocs = documents.map((doc) => { + const score = Math.abs(queryLength - doc.text().length); + return { + ...doc, + metadata: { ...doc.metadata, score }, + }; + }); + + return { + documents: rerankedDocs + .sort((a, b) => a.metadata.score - b.metadata.score) + .slice(0, options.k || 3), + }; + } + ); + + // Sample documents for testing + const documents = [ + Document.fromText('short'), + Document.fromText('a bit longer'), + Document.fromText('this is a very long document'), + ]; + + const query = Document.fromText('medium length'); + const rerankedDocuments = await rerank({ + reranker: customReranker, + query, + documents, + options: { k: 2 }, + }); + + // Validate the reranked results + assert.equal(rerankedDocuments.length, 2); + assert(rerankedDocuments[0].text().includes('a bit longer')); + assert(rerankedDocuments[1].text().includes('short')); + }); + + it('handles missing options gracefully', async () => { + const customReranker = defineReranker( + { + name: 'custom/reranker', + configSchema: z.object({ + k: z.number().optional(), + }), + }, + async (query, documents, options) => { + const rerankedDocs = documents.map((doc) => { + const score = Math.random(); // Simplified scoring for testing + return { + ...doc, + metadata: { ...doc.metadata, score }, + }; + }); + + return { + documents: rerankedDocs.sort( + (a, b) => b.metadata.score - a.metadata.score + ), + }; + } + ); + + const documents = [Document.fromText('doc1'), Document.fromText('doc2')]; + + const query = Document.fromText('test query'); + const rerankedDocuments = await rerank({ + reranker: customReranker, + query, + documents, + }); + + assert.equal(rerankedDocuments.length, 2); + assert(typeof rerankedDocuments[0].metadata.score === 'number'); + }); + + it('validates config schema and throws error on invalid input', async () => { + const customReranker = defineReranker( + { + name: 'custom/reranker', + configSchema: z.object({ + k: z.number().min(1), + }), + }, + async (query, documents, options) => { + // Simplified scoring for testing + const rerankedDocs = documents.map((doc) => ({ + ...doc, + metadata: { score: Math.random() }, + })); + return { + documents: rerankedDocs.sort( + (a, b) => b.metadata.score - a.metadata.score + ), + }; + } + ); + + const documents = [Document.fromText('doc1')]; + + const query = Document.fromText('test query'); + + try { + await rerank({ + reranker: customReranker, + query, + documents, + options: { k: 0 }, // Invalid input: k must be at least 1 + }); + assert.fail('Expected validation error'); + } catch (err) { + assert(err instanceof GenkitError); + assert.equal(err.status, 'INVALID_ARGUMENT'); + } + }); + + it('preserves document metadata after reranking', async () => { + const customReranker = defineReranker( + { + name: 'custom/reranker', + }, + async (query, documents) => { + const rerankedDocs = documents.map((doc, i) => ({ + ...doc, + metadata: { ...doc.metadata, score: 2 - i }, + })); + + return { + documents: rerankedDocs.sort( + (a, b) => b.metadata.score - a.metadata.score + ), + }; + } + ); + + const documents = [ + new Document({ content: [], metadata: { originalField: 'test1' } }), + new Document({ content: [], metadata: { originalField: 'test2' } }), + ]; + + const query = Document.fromText('test query'); + const rerankedDocuments = await rerank({ + reranker: customReranker, + query, + documents, + }); + + assert.equal(rerankedDocuments[0].metadata.originalField, 'test1'); + assert.equal(rerankedDocuments[1].metadata.originalField, 'test2'); + }); + + it('handles errors thrown by the reranker', async () => { + const customReranker = defineReranker( + { + name: 'custom/reranker', + }, + async (query, documents) => { + // Simulate an error in the reranker logic + throw new GenkitError({ + message: 'Something went wrong during reranking', + status: 'INTERNAL', + }); + } + ); + + const documents = [Document.fromText('doc1'), Document.fromText('doc2')]; + const query = Document.fromText('test query'); + + try { + await rerank({ + reranker: customReranker, + query, + documents, + }); + assert.fail('Expected an error to be thrown'); + } catch (err) { + assert(err instanceof GenkitError); + assert.equal(err.status, 'INTERNAL'); + assert.equal( + err.message, + 'INTERNAL: Something went wrong during reranking' + ); + } + }); + }); +}); diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts index 75daeedb74..e122be5f9f 100644 --- a/js/core/src/registry.ts +++ b/js/core/src/registry.ts @@ -40,7 +40,8 @@ export type ActionType = | 'model' | 'prompt' | 'util' - | 'tool'; + | 'tool' + | 'reranker'; /** * Looks up a registry key (action type and key) in the registry. diff --git a/js/plugins/vertexai/src/index.ts b/js/plugins/vertexai/src/index.ts index 2231e394a4..10c0dcaa4a 100644 --- a/js/plugins/vertexai/src/index.ts +++ b/js/plugins/vertexai/src/index.ts @@ -68,6 +68,7 @@ import { llama31, modelGardenOpenaiCompatibleModel, } from './model_garden.js'; +import { VertexRerankerConfig, vertexAiRerankers } from './reranker.js'; import { VectorSearchOptions, vertexAiIndexers, @@ -134,6 +135,8 @@ export interface PluginOptions { }; /** Configure Vertex AI vector search index options */ vectorSearchOptions?: VectorSearchOptions[]; + /** Configure reranker options */ + rerankOptions?: VertexRerankerConfig[]; } const CLOUD_PLATFROM_OAUTH_SCOPE = @@ -247,12 +250,21 @@ export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin( }); } + const rerankOptions = { + pluginOptions: options, + authClient, + projectId, + }; + + const rerankers = await vertexAiRerankers(rerankOptions); + return { models, embedders, evaluators: vertexEvaluators(authClient, metrics, projectId, location), retrievers, indexers, + rerankers, }; } ); diff --git a/js/plugins/vertexai/src/reranker.ts b/js/plugins/vertexai/src/reranker.ts new file mode 100644 index 0000000000..8361d13bd1 --- /dev/null +++ b/js/plugins/vertexai/src/reranker.ts @@ -0,0 +1,165 @@ +/** + * 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 { + defineReranker, + RankedDocument, + RerankerAction, + rerankerRef, +} from '@genkit-ai/ai/reranker'; +import { GoogleAuth } from 'google-auth-library'; +import z from 'zod'; +import { PluginOptions } from '.'; + +const DEFAULT_MODEL = 'semantic-ranker-512@latest'; + +const getRerankEndpoint = (projectId: string, location: string) => { + return `https://discoveryengine.googleapis.com/v1/projects/${projectId}/locations/${location}/rankingConfigs/default_ranking_config:rank`; +}; + +// Define the schema for the options used in the Vertex AI reranker +export const VertexAIRerankerOptionsSchema = z.object({ + k: z.number().optional().describe('Number of top documents to rerank'), // Optional: Number of documents to rerank + model: z.string().optional().describe('Model name for reranking'), // Optional: Model name, defaults to a pre-defined model + location: z + .string() + .optional() + .describe('Google Cloud location, e.g., "us-central1"'), // Optional: Location of the reranking model +}); + +// Type alias for the options schema +export type VertexAIRerankerOptions = z.infer< + typeof VertexAIRerankerOptionsSchema +>; + +// Define the structure for each individual reranker configuration +export const VertexRerankerConfigSchema = z.object({ + model: z.string().optional().describe('Model name for reranking'), // Optional: Model name, defaults to a pre-defined model +}); + +export interface VertexRerankerConfig { + name?: string; + model?: string; +} + +export interface VertexRerankPluginOptions { + rerankOptions: VertexRerankerConfig[]; + projectId: string; + location?: string; // Optional: Location of the reranker service +} + +export interface VertexRerankOptions { + authClient: GoogleAuth; + pluginOptions?: PluginOptions; +} + +/** + * Creates Vertex AI rerankers. + * + * This function returns a list of reranker actions for Vertex AI based on the provided + * rerank options and configuration. + * + * @param {VertexRerankOptions} params - The parameters for creating the rerankers. + * @returns {RerankerAction[]} - An array of reranker actions. + */ +export async function vertexAiRerankers( + params: VertexRerankOptions +): Promise[]> { + if (!params.pluginOptions) { + throw new Error( + 'Plugin options are required to create Vertex AI rerankers' + ); + } + const pluginOptions = params.pluginOptions; + if (!params.pluginOptions.rerankOptions) { + return []; + } + + const rerankOptions = params.pluginOptions.rerankOptions; + const rerankers: RerankerAction[] = []; + + if (!rerankOptions || rerankOptions.length === 0) { + return rerankers; + } + const auth = new GoogleAuth(); + const client = await auth.getClient(); + const projectId = await auth.getProjectId(); + + for (const rerankOption of rerankOptions) { + const reranker = defineReranker( + { + name: `vertexai/${rerankOption.name || rerankOption.model}`, + configSchema: VertexAIRerankerOptionsSchema.optional(), + }, + async (query, documents, _options) => { + const response = await client.request({ + method: 'POST', + url: getRerankEndpoint( + projectId, + pluginOptions.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 }; + } + ); + + rerankers.push(reranker); + } + + return rerankers; +} + +/** + * 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: { + name: string; + displayName?: string; +}) => { + return rerankerRef({ + name: `vertexai/${name}`, + info: { + label: params.displayName ?? `Vertex AI Reranker`, + }, + configSchema: VertexAIRerankerOptionsSchema.optional(), + }); +}; diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index 29d3f4c1f8..d3faa57295 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -1289,6 +1289,49 @@ importers: specifier: ^5.3.3 version: 5.4.5 + testapps/vertexai-reranker: + dependencies: + '@genkit-ai/ai': + specifier: workspace:* + version: link:../../ai + '@genkit-ai/core': + specifier: workspace:* + version: link:../../core + '@genkit-ai/dev-local-vectorstore': + specifier: workspace:* + version: link:../../plugins/dev-local-vectorstore + '@genkit-ai/dotprompt': + specifier: workspace:* + version: link:../../plugins/dotprompt + '@genkit-ai/evaluator': + specifier: workspace:* + version: link:../../plugins/evaluators + '@genkit-ai/firebase': + specifier: workspace:* + version: link:../../plugins/firebase + '@genkit-ai/flow': + specifier: workspace:* + version: link:../../flow + '@genkit-ai/vertexai': + specifier: workspace:* + version: link:../../plugins/vertexai + dotenv: + specifier: ^16.4.5 + version: 16.4.5 + express: + specifier: ^4.19.2 + version: 4.19.2 + google-auth-library: + specifier: ^9.11.0 + version: 9.11.0(encoding@0.1.13) + zod: + specifier: 3.22.4 + version: 3.22.4 + devDependencies: + typescript: + specifier: ^5.5.2 + version: 5.5.3 + testapps/vertexai-vector-search-bigquery: dependencies: '@genkit-ai/ai': diff --git a/js/testapps/vertexai-reranker/.env.example b/js/testapps/vertexai-reranker/.env.example new file mode 100644 index 0000000000..be43a54193 --- /dev/null +++ b/js/testapps/vertexai-reranker/.env.example @@ -0,0 +1,4 @@ +# .env.example + +PROJECT_ID=your_project_id_here +LOCATION=your_location_here diff --git a/js/testapps/vertexai-reranker/README.md b/js/testapps/vertexai-reranker/README.md new file mode 100644 index 0000000000..2d38a35051 --- /dev/null +++ b/js/testapps/vertexai-reranker/README.md @@ -0,0 +1,109 @@ +# Sample Vertex AI Plugin Reranker with Fake Document Content + +This sample app demonstrates the use of the Vertex AI plugin for reranking a set of documents based on a query using fake document content. This guide will walk you through setting up and running the sample. + +## Prerequisites + +Before running this sample, ensure you have the following: + +1. **Node.js** installed. +2. **PNPM** (Node Package Manager) installed. +3. A **Vertex AI** project with appropriate permissions for reranking models. + +## Getting Started + +### Step 1: Clone the Repository and Install Dependencies + +Clone this repository to your local machine and navigate to the project directory. Then, install the necessary dependencies: + +\`\`\`bash +pnpm install +\`\`\` + +### Step 2: Set Up Environment Variables + +Create a \`.env\` file in the root directory and set the following variables. You can use the provided \`.env.example\` as a reference. + +\`\`\`plaintext +PROJECT_ID=your_project_id_here +LOCATION=your_location_here +\`\`\` + +These variables are required to configure the Vertex AI project and location for reranking. + +### Step 3: Run the Sample + +Start the Genkit server: + +\`\`\`bash +genkit start +\`\`\` + +This will launch the server that hosts the reranking flow. + +## Sample Explanation + +### Overview + +This sample demonstrates how to use the Vertex AI plugin to rerank a predefined list of fake document content based on a query input. It utilizes a semantic reranker model from Vertex AI. + +### Key Components + +- **Fake Document Content**: A hardcoded array of strings representing document content. +- **Rerank Flow**: A flow that reranks the fake documents based on the provided query. +- **Genkit Configuration**: Configures Genkit with the Vertex AI plugin, setting up the project and reranking model. + +### Rerank Flow + +The \`rerankFlow\` function takes a query as input, reranks the predefined document content using the Vertex AI semantic reranker, and returns the documents sorted by relevance score. + +\`\`\`typescript +export const rerankFlow = defineFlow( +{ +name: 'rerankFlow', +inputSchema: z.object({ query: z.string() }), +outputSchema: z.array( +z.object({ +text: z.string(), +score: z.number(), +}) +), +}, +async ({ query }) => { +const documents = FAKE_DOCUMENT_CONTENT.map((text) => +Document.fromText(text) +); +const reranker = 'vertexai/reranker'; + + const rerankedDocuments = await rerank({ + reranker, + query: Document.fromText(query), + documents, + }); + + return rerankedDocuments.map((doc) => ({ + text: doc.text(), + score: doc.metadata.score, + })); + +} +); +\`\`\` + +### Running the Server + +The server is started using the \`startFlowsServer\` function, which sets up the Genkit server to handle flow requests. + +\`\`\`typescript +startFlowsServer(); +\`\`\` + +## License + +This project is licensed under the Apache License, Version 2.0. See the [LICENSE](LICENSE) file for details. + +## Conclusion + +This sample provides a basic demonstration of using the Vertex AI plugin with Genkit for reranking documents based on a query. It can be extended and adapted to suit more complex use cases and integrations with other data sources and services. + +For more information, please refer to the official [Firebase Genkit documentation](https://firebase.google.com/docs/genkit). diff --git a/js/testapps/vertexai-reranker/package.json b/js/testapps/vertexai-reranker/package.json new file mode 100644 index 0000000000..0837a56e91 --- /dev/null +++ b/js/testapps/vertexai-reranker/package.json @@ -0,0 +1,33 @@ +{ + "name": "vertexai-reranker", + "version": "1.0.0", + "description": "", + "main": "lib/index.js", + "scripts": { + "start": "node lib/index.js", + "compile": "tsc", + "build": "pnpm build:clean && pnpm compile", + "build:clean": "rm -rf ./lib", + "build:watch": "tsc --watch" + }, + "keywords": [], + "author": "", + "license": "ISC", + "dependencies": { + "@genkit-ai/ai": "workspace:*", + "@genkit-ai/core": "workspace:*", + "@genkit-ai/dev-local-vectorstore": "workspace:*", + "@genkit-ai/dotprompt": "workspace:*", + "@genkit-ai/evaluator": "workspace:*", + "@genkit-ai/firebase": "workspace:*", + "@genkit-ai/flow": "workspace:*", + "@genkit-ai/vertexai": "workspace:*", + "dotenv": "^16.4.5", + "express": "^4.19.2", + "google-auth-library": "^9.11.0", + "zod": "3.22.4" + }, + "devDependencies": { + "typescript": "^5.5.2" + } +} diff --git a/js/testapps/vertexai-reranker/src/config.ts b/js/testapps/vertexai-reranker/src/config.ts new file mode 100644 index 0000000000..17353dc818 --- /dev/null +++ b/js/testapps/vertexai-reranker/src/config.ts @@ -0,0 +1,19 @@ +/** + * 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 { config } from 'dotenv'; +config(); +export const PROJECT_ID = process.env.PROJECT_ID!; +export const LOCATION = process.env.LOCATION!; diff --git a/js/testapps/vertexai-reranker/src/index.ts b/js/testapps/vertexai-reranker/src/index.ts new file mode 100644 index 0000000000..6f6fcb496f --- /dev/null +++ b/js/testapps/vertexai-reranker/src/index.ts @@ -0,0 +1,96 @@ +/** + * 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. + */ + +// Sample app for using the proposed Vertex AI plugin retriever and indexer with a local file (just as a demo). + +import { configureGenkit } from '@genkit-ai/core'; +import { defineFlow, startFlowsServer } from '@genkit-ai/flow'; +// important imports for this sample: +import { Document } from '@genkit-ai/ai/retriever'; +import { vertexAI } from '@genkit-ai/vertexai'; +// // Environment variables set with dotenv for simplicity of sample +import { rerank } from '@genkit-ai/ai/reranker'; +import { z } from 'zod'; +import { LOCATION, PROJECT_ID } from './config'; + +// Configure Genkit with Vertex AI plugin +configureGenkit({ + plugins: [ + vertexAI({ + projectId: PROJECT_ID, + location: LOCATION, + googleAuth: { + scopes: ['https://www.googleapis.com/auth/cloud-platform'], + }, + rerankOptions: [ + { + model: 'vertexai/semantic-ranker-512', + }, + ], + }), + ], + logLevel: 'debug', + enableTracingAndMetrics: true, +}); +const FAKE_DOCUMENT_CONTENT = [ + 'pythagorean theorem', + 'e=mc^2', + 'pi', + 'dinosaurs', + "euler's identity", + 'prime numbers', + 'fourier transform', + 'ABC conjecture', + 'riemann hypothesis', + 'triangles', + "schrodinger's cat", + 'quantum mechanics', + 'the avengers', + "harry potter and the philosopher's stone", + 'movies', +]; + +export const rerankFlow = defineFlow( + { + name: 'rerankFlow', + inputSchema: z.object({ query: z.string() }), + outputSchema: z.array( + z.object({ + text: z.string(), + score: z.number(), + }) + ), + }, + async ({ query }) => { + const documents = FAKE_DOCUMENT_CONTENT.map((text) => + Document.fromText(text) + ); + const reranker = 'vertexai/reranker'; + + const rerankedDocuments = await rerank({ + reranker, + query: Document.fromText(query), + documents, + }); + + return rerankedDocuments.map((doc) => ({ + text: doc.text(), + score: doc.metadata.score, + })); + } +); + +startFlowsServer(); diff --git a/js/testapps/vertexai-reranker/tsconfig.json b/js/testapps/vertexai-reranker/tsconfig.json new file mode 100644 index 0000000000..efbb566bf7 --- /dev/null +++ b/js/testapps/vertexai-reranker/tsconfig.json @@ -0,0 +1,14 @@ +{ + "compileOnSave": true, + "include": ["src"], + "compilerOptions": { + "module": "commonjs", + "noImplicitReturns": true, + "outDir": "lib", + "sourceMap": true, + "strict": true, + "target": "es2017", + "skipLibCheck": true, + "esModuleInterop": true + } +}