Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: @invertase update Vertex AI reranker app with proper model and setup and docs #1259

Open
wants to merge 1 commit into
base: next
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions js/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

69 changes: 55 additions & 14 deletions js/testapps/vertexai-reranker/README.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# 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.
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. Follow this guide to set up and run the sample.

## Prerequisites

Expand All @@ -9,6 +9,7 @@ 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.
4. The **Discovery Engine API** enabled on your Google Cloud project. [Enable it here](https://console.cloud.google.com/apis/api/discoveryengine.googleapis.com/metrics).

## Getting Started

Expand All @@ -26,30 +27,48 @@ Create a \`.env\` file in the root directory and set the following variables. Yo

\`\`\`plaintext
PROJECT_ID=your_project_id_here
LOCATION=your_location_here
LOCATION=global
\`\`\`

These variables are required to configure the Vertex AI project and location for reranking.

### Step 3: Run the Sample
### Step 3: Build the Project

Start the Genkit server:
Run the following command to build the project:

\`\`\`bash
genkit start
pnpm build
\`\`\`

This will launch the server that hosts the reranking flow.
### Step 4: Run the Sample

To start the development environment and test the flow, run:

1. In the first terminal:

\`\`\`bash
pnpm genkit:dev
\`\`\`

This will start the Genkit development server.

2. In another terminal, use the UI or invoke the flow:

\`\`\`bash
genkit flow:run rerankFlow --input '{"query": "quantum mechanics"}'
\`\`\`

Alternatively, access the Genkit UI to interact with the 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.
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 leverages the \`semantic-ranker-512@latest\` model from Vertex AI.

### Key Components

- **Fake Document Content**: A hardcoded array of strings representing document content.
- **Fake Document Content**: A hardcoded array of strings simulating 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.

Expand All @@ -58,7 +77,7 @@ This sample demonstrates how to use the Vertex AI plugin to rerank a predefined
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(
export const rerankFlow = ai.defineFlow(
{
name: 'rerankFlow',
inputSchema: z.object({ query: z.string() }),
Expand All @@ -73,9 +92,9 @@ async ({ query }) => {
const documents = FAKE_DOCUMENT_CONTENT.map((text) =>
Document.fromText(text)
);
const reranker = 'vertexai/reranker';
const reranker = 'vertexai/semantic-ranker-512@latest';

const rerankedDocuments = await rerank({
const rerankedDocuments = await ai.rerank({
reranker,
query: Document.fromText(query),
documents,
Expand All @@ -92,12 +111,34 @@ const reranker = 'vertexai/reranker';

### Running the Server

The server is started using the \`startFlowsServer\` function, which sets up the Genkit server to handle flow requests.
The server is started using the \`startFlowServer\` function, which sets up the Genkit server to handle flow requests.

\`\`\`typescript
startFlowsServer();
ai.startFlowServer({
flows: [rerankFlow],
});
\`\`\`

### Example Input and Output

- **Input**:

\`\`\`json
{
"query": "quantum mechanics"
}
\`\`\`

- **Output**:

\`\`\`json
[
{ "text": "quantum mechanics", "score": 0.95 },
{ "text": "schrodinger's cat", "score": 0.85 },
{ "text": "e=mc^2", "score": 0.80 }
]
\`\`\`

## License

This project is licensed under the Apache License, Version 2.0. See the [LICENSE](LICENSE) file for details.
Expand All @@ -106,4 +147,4 @@ This project is licensed under the Apache License, Version 2.0. See the [LICENSE

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).
For more information, refer to the official [Firebase Genkit documentation](https://firebase.google.com/docs/genkit).
5 changes: 4 additions & 1 deletion js/testapps/vertexai-reranker/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@
"compile": "tsc",
"build": "pnpm build:clean && pnpm compile",
"build:clean": "rimraf ./lib",
"build:watch": "tsc --watch"
"build:watch": "tsc --watch",
"genkit:dev": "cross-env GENKIT_ENV=dev pnpm dev",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This probably needs to be applied globally in all testapps -- we no longer need to set GENKIT_ENV and rely on cross-env

Suggested change
"genkit:dev": "cross-env GENKIT_ENV=dev pnpm dev",
"genkit:dev": "genkit start -- pnpm dev",

"dev": "tsx --watch src/index.ts"
},
"keywords": [],
"author": "",
Expand All @@ -24,6 +26,7 @@
"google-auth-library": "^9.11.0"
},
"devDependencies": {
"cross-env": "^7.0.3",
"rimraf": "^6.0.1",
"typescript": "^5.5.2"
}
Expand Down
63 changes: 54 additions & 9 deletions js/testapps/vertexai-reranker/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,15 @@
* limitations under the License.
*/

// Sample app for using the proposed Vertex AI plugin retriever and indexer with a local file (just as a demo).
// Sample app for demonstrating the Vertex AI plugin retriever and reranker capabilities using a local file for demo purposes.

import { Document, genkit, z } from 'genkit';
// important imports for this sample:
// Import necessary Vertex AI plugins and configuration:
import { vertexAI } from '@genkit-ai/vertexai';
import { vertexAIRerankers } from '@genkit-ai/vertexai/rerankers';
import { LOCATION, PROJECT_ID } from './config';

// Configure Genkit with Vertex AI plugin
// Initialize Genkit with Vertex AI and Reranker plugins
const ai = genkit({
plugins: [
vertexAI({
Expand All @@ -37,13 +37,17 @@ const ai = genkit({
location: LOCATION,
rerankOptions: [
{
model: 'vertexai/reranker',
model: 'semantic-ranker-512@latest', // Semantic ranker model to be used
},
],
}),
],
});

/**
* Mock data to simulate document retrieval.
* Each item represents a document with simplified content.
*/
const FAKE_DOCUMENT_CONTENT = [
'pythagorean theorem',
'e=mc^2',
Expand All @@ -62,36 +66,77 @@ const FAKE_DOCUMENT_CONTENT = [
'movies',
];

/**
* @flow rerankFlow
* Defines a flow to rerank a set of documents based on a given query.
*
* @flowDescription
* This flow takes a query string, retrieves predefined documents, and ranks them
* based on their relevance to the query using Vertex AI's reranker model.
*
* @inputSchema
* The flow expects an object containing:
* - query: A string representing the user's search query.
*
* @outputSchema
* Returns an array of objects containing:
* - text: The content of the ranked document.
* - score: The relevance score assigned by the reranker model.
*
* @example
* Input:
* {
* "query": "quantum mechanics"
* }
*
* Output:
* [
* { "text": "quantum mechanics", "score": 0.95 },
* { "text": "schrodinger's cat", "score": 0.85 },
* { "text": "e=mc^2", "score": 0.80 }
* ]
*/
export const rerankFlow = ai.defineFlow(
{
name: 'rerankFlow',
inputSchema: z.object({ query: z.string() }),
inputSchema: z.object({ query: z.string() }), // Input must be an object with a 'query' string
outputSchema: z.array(
z.object({
text: z.string(),
score: z.number(),
text: z.string(), // Each result includes a document's text
score: z.number(), // And its relevance score
})
),
},
async ({ query }) => {
console.log('Received query:', query);

// Convert fake document content into Document objects
const documents = FAKE_DOCUMENT_CONTENT.map((text) =>
Document.fromText(text)
);
const reranker = 'vertexai/reranker';

// Specify the reranker to be used
const reranker = 'vertexai/semantic-ranker-512@latest';

// Call the reranker with the query and documents
const rerankedDocuments = await ai.rerank({
reranker,
query: Document.fromText(query),
documents,
});

// Return the reranked documents with text and score
return rerankedDocuments.map((doc) => ({
text: doc.text,
score: doc.metadata.score,
}));
}
);

/**
* Starts the Flow Server for testing and UI interaction.
* This allows the `rerankFlow` to be invoked from the Genkit UI or programmatically.
*/
ai.startFlowServer({
flows: [rerankFlow],
flows: [rerankFlow], // Registers the defined flow for the server
});
3 changes: 2 additions & 1 deletion js/testapps/vertexai-reranker/tsconfig.json
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@
"compileOnSave": true,
"include": ["src"],
"compilerOptions": {
"module": "commonjs",
"module": "Node16",
"moduleResolution": "node16",
"noImplicitReturns": true,
"outDir": "lib",
"sourceMap": true,
Expand Down
Loading