From e33d66fc91111a780b0bf82a31f5f6627cbca764 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Tue, 20 Aug 2024 13:32:11 +0100
Subject: [PATCH 1/8] feat(js): add reranker actions

---
 js/ai/package.json      |   9 ++
 js/ai/src/reranker.ts   | 244 ++++++++++++++++++++++++++++++++++++++++
 js/core/src/registry.ts |   3 +-
 3 files changed, 255 insertions(+), 1 deletion(-)
 create mode 100644 js/ai/src/reranker.ts

diff --git a/js/ai/package.json b/js/ai/package.json
index 4b38e7fa7..c9980ae60 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 000000000..ee9050428
--- /dev/null
+++ b/js/ai/src/reranker.ts
@@ -0,0 +1,244 @@
+/**
+ * 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<RerankerOptions extends z.ZodTypeAny> = (
+  query: Document,
+  documents: Document[],
+  queryOpts: z.infer<RerankerOptions>
+) => Promise<RerankerResponse>;
+
+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<typeof RankedDocumentDataSchema>;
+
+export class RankedDocument extends Document implements RankedDocumentData {
+  content: Part[];
+  metadata: { score: number } & Record<string, any>;
+
+  constructor(data: RankedDocumentData) {
+    super(data);
+    this.content = data.content;
+    this.metadata = data.metadata;
+  }
+
+  static fromText(
+    text: string,
+    metadata: { score: number } & Record<string, any>
+  ) {
+    // TODO: do we need validation here? I'm thinking in javascript we won't have type errors
+    if (metadata.score === undefined) {
+      throw new Error('Score is required');
+    }
+    return new RankedDocument({
+      content: [{ text }],
+      metadata,
+    });
+  }
+
+  /**
+   * Concatenates all `text` parts present in the document with no delimiter.
+   * @returns A string of all concatenated text parts.
+   */
+  text(): string {
+    return this.content.map((part) => part.text || '').join('');
+  }
+
+  /**
+   * Returns the score of the document.
+   * @returns The score of the document.
+   */
+  score(): number {
+    return this.metadata.score;
+  }
+
+  /**
+   * Returns the first media part detected in the document. Useful for extracting
+   * (for example) an image.
+   * @returns The first detected `media` part in the document.
+   */
+  media(): { url: string; contentType?: string } | null {
+    return this.content.find((part) => part.media)?.media || null;
+  }
+
+  toJSON(): DocumentData {
+    return {
+      content: this.content,
+      metadata: this.metadata,
+    };
+  }
+}
+
+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<typeof RerankerResponseSchema>;
+
+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<typeof RerankerInfoSchema>;
+
+export type RerankerAction<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =
+  Action<
+    typeof RerankerRequestSchema,
+    typeof RerankerResponseSchema,
+    { model: RerankerInfo }
+  > & {
+    __configSchema?: CustomOptions;
+  };
+
+function rerankerWithMetadata<
+  RerankerOptions extends z.ZodTypeAny = z.ZodTypeAny,
+>(
+  reranker: Action<typeof RerankerRequestSchema, typeof RerankerResponseSchema>,
+  configSchema?: RerankerOptions
+): RerankerAction<RerankerOptions> {
+  const withMeta = reranker as RerankerAction<RerankerOptions>;
+  withMeta.__configSchema = configSchema;
+  return withMeta;
+}
+
+/**
+ *  Creates a reranker action for the provided {@link RerankerFn} implementation.
+ */
+export function defineReranker<OptionsType extends z.ZodTypeAny = z.ZodTypeAny>(
+  options: {
+    name: string;
+    configSchema?: OptionsType;
+    info?: RerankerInfo;
+  },
+  runner: RerankerFn<OptionsType>
+) {
+  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<CustomOptions>;
+  query: string | DocumentData;
+  documents: DocumentData[];
+  options?: z.infer<CustomOptions>;
+}
+
+export type RerankerArgument<
+  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+> = RerankerAction<CustomOptions> | RerankerReference<CustomOptions> | string;
+
+/**
+ * Reranks documents from a {@link RerankerArgument} based on the provided query.
+ */
+export async function rerank<CustomOptions extends z.ZodTypeAny>(
+  params: RerankerParams<CustomOptions>
+): Promise<Array<RankedDocument>> {
+  let reranker: RerankerAction<CustomOptions>;
+  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<CustomOptions>;
+  }
+  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<CustomOptions extends z.ZodTypeAny> {
+  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<CustomOptionsSchema>
+): RerankerReference<CustomOptionsSchema> {
+  return { ...options };
+}
diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts
index 75daeedb7..e122be5f9 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.

From ad56ef3e7785136ed78550ddcc9ad869ca6410dd Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Tue, 20 Aug 2024 13:33:47 +0100
Subject: [PATCH 2/8] feat(js/plugins/vertexai): add reranker to vertexai
 plugin

---
 js/plugins/vertexai/src/index.ts    |  12 ++
 js/plugins/vertexai/src/reranker.ts | 166 ++++++++++++++++++++++++++++
 js/pnpm-lock.yaml                   |  58 ++++++++++
 3 files changed, 236 insertions(+)
 create mode 100644 js/plugins/vertexai/src/reranker.ts

diff --git a/js/plugins/vertexai/src/index.ts b/js/plugins/vertexai/src/index.ts
index 2231e394a..fb7044b3d 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<z.ZodTypeAny, any, any>[];
+  /** 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 = 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 000000000..17ec9cc5b
--- /dev/null
+++ b/js/plugins/vertexai/src/reranker.ts
@@ -0,0 +1,166 @@
+/**
+ * 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<EmbedderCustomOptions>} params - The parameters for creating the rerankers.
+ * @returns {RerankerAction<z.ZodTypeAny>[]} - An array of reranker actions.
+ */
+export function vertexAiRerankers(
+  params: VertexRerankOptions
+): RerankerAction<z.ZodTypeAny>[] {
+  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<z.ZodTypeAny>[] = [];
+
+  if (!rerankOptions || rerankOptions.length === 0) {
+    return rerankers;
+  }
+
+  for (const rerankOption of rerankOptions) {
+    const reranker = defineReranker(
+      {
+        name: `vertexai/${rerankOption.name || rerankOption.model}`,
+        configSchema: VertexAIRerankerOptionsSchema.optional(),
+      },
+      async (query, documents, _options) => {
+        const auth = new GoogleAuth();
+        const client = await auth.getClient();
+        const projectId = await auth.getProjectId();
+
+        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 29d3f4c1f..8c3eeb978 100644
--- a/js/pnpm-lock.yaml
+++ b/js/pnpm-lock.yaml
@@ -1289,6 +1289,64 @@ 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/googleai':
+        specifier: workspace:*
+        version: link:../../plugins/googleai
+      '@genkit-ai/vertexai':
+        specifier: workspace:*
+        version: link:../../plugins/vertexai
+      '@google-cloud/bigquery':
+        specifier: ^7.8.0
+        version: 7.8.0(encoding@0.1.13)
+      dotenv:
+        specifier: ^16.4.5
+        version: 16.4.5
+      express:
+        specifier: ^4.19.2
+        version: 4.19.2
+      genkitx-chromadb:
+        specifier: workspace:*
+        version: link:../../plugins/chroma
+      genkitx-langchain:
+        specifier: workspace:*
+        version: link:../../plugins/langchain
+      genkitx-pinecone:
+        specifier: workspace:*
+        version: link:../../plugins/pinecone
+      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':

From 7673c91022ddb4800b8b73551c96ce1477c60563 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Tue, 20 Aug 2024 13:34:14 +0100
Subject: [PATCH 3/8] feat(js/testapps): add testapp for vertexai reranker

---
 js/testapps/vertexai-reranker/.env.example  |  13 ++
 js/testapps/vertexai-reranker/README.md     | 137 ++++++++++++++++++++
 js/testapps/vertexai-reranker/package.json  |  38 ++++++
 js/testapps/vertexai-reranker/src/config.ts |  19 +++
 js/testapps/vertexai-reranker/src/index.ts  |  96 ++++++++++++++
 js/testapps/vertexai-reranker/tsconfig.json |  14 ++
 6 files changed, 317 insertions(+)
 create mode 100644 js/testapps/vertexai-reranker/.env.example
 create mode 100644 js/testapps/vertexai-reranker/README.md
 create mode 100644 js/testapps/vertexai-reranker/package.json
 create mode 100644 js/testapps/vertexai-reranker/src/config.ts
 create mode 100644 js/testapps/vertexai-reranker/src/index.ts
 create mode 100644 js/testapps/vertexai-reranker/tsconfig.json

diff --git a/js/testapps/vertexai-reranker/.env.example b/js/testapps/vertexai-reranker/.env.example
new file mode 100644
index 000000000..642ad56f1
--- /dev/null
+++ b/js/testapps/vertexai-reranker/.env.example
@@ -0,0 +1,13 @@
+# .env.example
+
+# Firebase project configuration
+PROJECT_ID=your_project_id_here
+LOCATION=your_location_here
+
+LOCAL_DIR=your_local_dir_here
+
+# Vector Search configuration
+VECTOR_SEARCH_PUBLIC_DOMAIN_NAME=your_vector_search_public_endpoint_here
+VECTOR_SEARCH_INDEX_ENDPOINT_ID=your_vector_search_index_endpoint_id_here
+VECTOR_SEARCH_INDEX_ID=your_vector_search_index_id_here
+VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_vector_search_deployed_index_id_here
diff --git a/js/testapps/vertexai-reranker/README.md b/js/testapps/vertexai-reranker/README.md
new file mode 100644
index 000000000..f0f84f9e1
--- /dev/null
+++ b/js/testapps/vertexai-reranker/README.md
@@ -0,0 +1,137 @@
+# Sample Vertex AI Plugin Retriever and Indexer with Local File
+
+This sample app demonstrates the use of the Vertex AI plugin retriever and indexer with a local file for demonstration purposes. 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 deployed index to an index endpoint in **Vertex AI Vector Search**.
+
+## Getting Started
+
+### Step 1: Clone the Repository and Install Dependencies
+
+Clone this repository to your local machine, and follow the instructions in the root README.md to build
+the core packages. This sample uses `workspace:*` dependencies, so they will need to be accessible.
+
+Then
+
+```bash
+cd js/testapps/vertex-vector-search-custom && pnpm i
+```
+
+### Step 3: Set Up Environment Variables
+
+Ensure you have a deployed index in Vertex AI Vector Search.
+
+Create a `.env` file in the root directory and set the following variables (see the .env.example as well if needed)
+
+```plaintext
+PROJECT_ID=your-google-cloud-project-id
+LOCATION=your-vertex-ai-location
+LOCAL_DIR=./data
+VECTOR_SEARCH_PUBLIC_DOMAIN_NAME=your-vector-search-public-domain-name
+VECTOR_SEARCH_INDEX_ENDPOINT_ID=your-index-endpoint-id
+VECTOR_SEARCH_INDEX_ID=your-index-id
+VECTOR_SEARCH_DEPLOYED_INDEX_ID=your-deployed-index-id
+GOOGLE_APPLICATION_CREDENTIALS=path-to-your-service-account-key.json
+```
+
+### Step 4: Run the Sample
+
+Start the Genkit server:
+
+```bash
+genkit start
+```
+
+## Sample Explanation
+
+### Overview
+
+This sample demonstrates how to define a custom document indexer and retriever using local JSON files. It integrates with Vertex AI for indexing and retrieval of documents.
+
+### Key Components
+
+- **Custom Document Indexer**: Stores documents in a local JSON file.
+- **Custom Document Retriever**: Retrieves documents from the local JSON file based on neighbor IDs.
+- **Genkit Configuration**: Configures Genkit with the Vertex AI plugin, setting up the project, location, and vector search index options.
+- **Indexing Flow**: Defines a flow for indexing documents.
+- **Query Flow**: Defines a flow for querying indexed documents.
+
+### Custom Document Indexer
+
+The `localDocumentIndexer` function reads existing documents from a local file, adds new documents, and writes them back to the file:
+
+```typescript
+const localDocumentIndexer: DocumentIndexer = async (documents: Document[]) => {
+  const content = await fs.promises.readFile(localFilePath, 'utf-8');
+  const currentLocalFile = JSON.parse(content);
+  const docsWithIds = Object.fromEntries(
+    documents.map((doc) => [
+      generateRandomId(),
+      { content: JSON.stringify(doc.content) },
+    ])
+  );
+  const newLocalFile = { ...currentLocalFile, ...docsWithIds };
+  await fs.promises.writeFile(
+    localFilePath,
+    JSON.stringify(newLocalFile, null, 2)
+  );
+  return Object.keys(docsWithIds);
+};
+```
+
+### Custom Document Retriever
+
+The `localDocumentRetriever` function reads the local file and retrieves documents based on neighbor IDs:
+
+```typescript
+const localDocumentRetriever: DocumentRetriever = async (
+  neighbors: Neighbor[]
+) => {
+  const content = await fs.promises.readFile(localFilePath, 'utf-8');
+  const currentLocalFile = JSON.parse(content);
+  const ids = neighbors
+    .map((neighbor) => neighbor.datapoint?.datapointId)
+    .filter(Boolean) as string[];
+  const docs = ids
+    .map((id) => {
+      const doc = currentLocalFile[id];
+      if (!doc || !doc.content) return null;
+      const parsedContent = JSON.parse(doc.content);
+      const text = parsedContent[0]?.text;
+      return text ? Document.fromText(text) : null;
+    })
+    .filter(Boolean) as Document[];
+  return docs;
+};
+```
+
+### Defining Flows
+
+Two flows are defined: `indexFlow` for indexing documents and `queryFlow` for querying documents.
+
+- **Index Flow**: Converts text inputs to documents and indexes them.
+- **Query Flow**: Retrieves documents based on a query and returns the results sorted by distance.
+
+### 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 Vertex AI plugins with Genkit for document indexing and retrieval. 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 000000000..9d1ab6997
--- /dev/null
+++ b/js/testapps/vertexai-reranker/package.json
@@ -0,0 +1,38 @@
+{
+  "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/googleai": "workspace:*",
+    "@genkit-ai/vertexai": "workspace:*",
+    "@google-cloud/bigquery": "^7.8.0",
+    "dotenv": "^16.4.5",
+    "express": "^4.19.2",
+    "genkitx-chromadb": "workspace:*",
+    "genkitx-langchain": "workspace:*",
+    "genkitx-pinecone": "workspace:*",
+    "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 000000000..17353dc81
--- /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 000000000..6f6fcb496
--- /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 000000000..efbb566bf
--- /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
+  }
+}

From 1403f1552c9b654272b25baec09c2ec84e21a924 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Tue, 20 Aug 2024 13:51:00 +0100
Subject: [PATCH 4/8] chore(js/ai/reranker): remove duplicate document method
 implementations

---
 js/ai/src/reranker.ts | 39 ---------------------------------------
 1 file changed, 39 deletions(-)

diff --git a/js/ai/src/reranker.ts b/js/ai/src/reranker.ts
index ee9050428..9e6ea10fe 100644
--- a/js/ai/src/reranker.ts
+++ b/js/ai/src/reranker.ts
@@ -46,29 +46,6 @@ export class RankedDocument extends Document implements RankedDocumentData {
     this.content = data.content;
     this.metadata = data.metadata;
   }
-
-  static fromText(
-    text: string,
-    metadata: { score: number } & Record<string, any>
-  ) {
-    // TODO: do we need validation here? I'm thinking in javascript we won't have type errors
-    if (metadata.score === undefined) {
-      throw new Error('Score is required');
-    }
-    return new RankedDocument({
-      content: [{ text }],
-      metadata,
-    });
-  }
-
-  /**
-   * Concatenates all `text` parts present in the document with no delimiter.
-   * @returns A string of all concatenated text parts.
-   */
-  text(): string {
-    return this.content.map((part) => part.text || '').join('');
-  }
-
   /**
    * Returns the score of the document.
    * @returns The score of the document.
@@ -76,22 +53,6 @@ export class RankedDocument extends Document implements RankedDocumentData {
   score(): number {
     return this.metadata.score;
   }
-
-  /**
-   * Returns the first media part detected in the document. Useful for extracting
-   * (for example) an image.
-   * @returns The first detected `media` part in the document.
-   */
-  media(): { url: string; contentType?: string } | null {
-    return this.content.find((part) => part.media)?.media || null;
-  }
-
-  toJSON(): DocumentData {
-    return {
-      content: this.content,
-      metadata: this.metadata,
-    };
-  }
 }
 
 const RerankerRequestSchema = z.object({

From 86dfa5fc6da50285a494f79c83be70ba4c98e11d Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Tue, 20 Aug 2024 13:56:05 +0100
Subject: [PATCH 5/8] chore(js/testapps): update README and deps of reranker
 testapp

---
 js/pnpm-lock.yaml                          |  15 --
 js/testapps/vertexai-reranker/.env.example |   9 --
 js/testapps/vertexai-reranker/README.md    | 156 +++++++++------------
 js/testapps/vertexai-reranker/package.json |   5 -
 4 files changed, 64 insertions(+), 121 deletions(-)

diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml
index 8c3eeb978..d3faa5729 100644
--- a/js/pnpm-lock.yaml
+++ b/js/pnpm-lock.yaml
@@ -1312,30 +1312,15 @@ importers:
       '@genkit-ai/flow':
         specifier: workspace:*
         version: link:../../flow
-      '@genkit-ai/googleai':
-        specifier: workspace:*
-        version: link:../../plugins/googleai
       '@genkit-ai/vertexai':
         specifier: workspace:*
         version: link:../../plugins/vertexai
-      '@google-cloud/bigquery':
-        specifier: ^7.8.0
-        version: 7.8.0(encoding@0.1.13)
       dotenv:
         specifier: ^16.4.5
         version: 16.4.5
       express:
         specifier: ^4.19.2
         version: 4.19.2
-      genkitx-chromadb:
-        specifier: workspace:*
-        version: link:../../plugins/chroma
-      genkitx-langchain:
-        specifier: workspace:*
-        version: link:../../plugins/langchain
-      genkitx-pinecone:
-        specifier: workspace:*
-        version: link:../../plugins/pinecone
       google-auth-library:
         specifier: ^9.11.0
         version: 9.11.0(encoding@0.1.13)
diff --git a/js/testapps/vertexai-reranker/.env.example b/js/testapps/vertexai-reranker/.env.example
index 642ad56f1..be43a5419 100644
--- a/js/testapps/vertexai-reranker/.env.example
+++ b/js/testapps/vertexai-reranker/.env.example
@@ -1,13 +1,4 @@
 # .env.example
 
-# Firebase project configuration
 PROJECT_ID=your_project_id_here
 LOCATION=your_location_here
-
-LOCAL_DIR=your_local_dir_here
-
-# Vector Search configuration
-VECTOR_SEARCH_PUBLIC_DOMAIN_NAME=your_vector_search_public_endpoint_here
-VECTOR_SEARCH_INDEX_ENDPOINT_ID=your_vector_search_index_endpoint_id_here
-VECTOR_SEARCH_INDEX_ID=your_vector_search_index_id_here
-VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_vector_search_deployed_index_id_here
diff --git a/js/testapps/vertexai-reranker/README.md b/js/testapps/vertexai-reranker/README.md
index f0f84f9e1..2d38a3505 100644
--- a/js/testapps/vertexai-reranker/README.md
+++ b/js/testapps/vertexai-reranker/README.md
@@ -1,6 +1,6 @@
-# Sample Vertex AI Plugin Retriever and Indexer with Local File
+# Sample Vertex AI Plugin Reranker with Fake Document Content
 
-This sample app demonstrates the use of the Vertex AI plugin retriever and indexer with a local file for demonstration purposes. 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. This guide will walk you through setting up and running the sample.
 
 ## Prerequisites
 
@@ -8,123 +8,95 @@ Before running this sample, ensure you have the following:
 
 1. **Node.js** installed.
 2. **PNPM** (Node Package Manager) installed.
-3. A deployed index to an index endpoint in **Vertex AI Vector Search**.
+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 follow the instructions in the root README.md to build
-the core packages. This sample uses `workspace:*` dependencies, so they will need to be accessible.
+Clone this repository to your local machine and navigate to the project directory. Then, install the necessary dependencies:
 
-Then
+\`\`\`bash
+pnpm install
+\`\`\`
 
-```bash
-cd js/testapps/vertex-vector-search-custom && pnpm i
-```
+### Step 2: Set Up Environment Variables
 
-### Step 3: 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.
 
-Ensure you have a deployed index in Vertex AI Vector Search.
+\`\`\`plaintext
+PROJECT_ID=your_project_id_here
+LOCATION=your_location_here
+\`\`\`
 
-Create a `.env` file in the root directory and set the following variables (see the .env.example as well if needed)
+These variables are required to configure the Vertex AI project and location for reranking.
 
-```plaintext
-PROJECT_ID=your-google-cloud-project-id
-LOCATION=your-vertex-ai-location
-LOCAL_DIR=./data
-VECTOR_SEARCH_PUBLIC_DOMAIN_NAME=your-vector-search-public-domain-name
-VECTOR_SEARCH_INDEX_ENDPOINT_ID=your-index-endpoint-id
-VECTOR_SEARCH_INDEX_ID=your-index-id
-VECTOR_SEARCH_DEPLOYED_INDEX_ID=your-deployed-index-id
-GOOGLE_APPLICATION_CREDENTIALS=path-to-your-service-account-key.json
-```
-
-### Step 4: Run the Sample
+### Step 3: Run the Sample
 
 Start the Genkit server:
 
-```bash
+\`\`\`bash
 genkit start
-```
+\`\`\`
+
+This will launch the server that hosts the reranking flow.
 
 ## Sample Explanation
 
 ### Overview
 
-This sample demonstrates how to define a custom document indexer and retriever using local JSON files. It integrates with Vertex AI for indexing and retrieval of documents.
+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
 
-- **Custom Document Indexer**: Stores documents in a local JSON file.
-- **Custom Document Retriever**: Retrieves documents from the local JSON file based on neighbor IDs.
-- **Genkit Configuration**: Configures Genkit with the Vertex AI plugin, setting up the project, location, and vector search index options.
-- **Indexing Flow**: Defines a flow for indexing documents.
-- **Query Flow**: Defines a flow for querying indexed documents.
-
-### Custom Document Indexer
-
-The `localDocumentIndexer` function reads existing documents from a local file, adds new documents, and writes them back to the file:
-
-```typescript
-const localDocumentIndexer: DocumentIndexer = async (documents: Document[]) => {
-  const content = await fs.promises.readFile(localFilePath, 'utf-8');
-  const currentLocalFile = JSON.parse(content);
-  const docsWithIds = Object.fromEntries(
-    documents.map((doc) => [
-      generateRandomId(),
-      { content: JSON.stringify(doc.content) },
-    ])
-  );
-  const newLocalFile = { ...currentLocalFile, ...docsWithIds };
-  await fs.promises.writeFile(
-    localFilePath,
-    JSON.stringify(newLocalFile, null, 2)
-  );
-  return Object.keys(docsWithIds);
-};
-```
-
-### Custom Document Retriever
-
-The `localDocumentRetriever` function reads the local file and retrieves documents based on neighbor IDs:
-
-```typescript
-const localDocumentRetriever: DocumentRetriever = async (
-  neighbors: Neighbor[]
-) => {
-  const content = await fs.promises.readFile(localFilePath, 'utf-8');
-  const currentLocalFile = JSON.parse(content);
-  const ids = neighbors
-    .map((neighbor) => neighbor.datapoint?.datapointId)
-    .filter(Boolean) as string[];
-  const docs = ids
-    .map((id) => {
-      const doc = currentLocalFile[id];
-      if (!doc || !doc.content) return null;
-      const parsedContent = JSON.parse(doc.content);
-      const text = parsedContent[0]?.text;
-      return text ? Document.fromText(text) : null;
-    })
-    .filter(Boolean) as Document[];
-  return docs;
-};
-```
-
-### Defining Flows
-
-Two flows are defined: `indexFlow` for indexing documents and `queryFlow` for querying documents.
-
-- **Index Flow**: Converts text inputs to documents and indexes them.
-- **Query Flow**: Retrieves documents based on a query and returns the results sorted by distance.
+- **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.
+The server is started using the \`startFlowsServer\` function, which sets up the Genkit server to handle flow requests.
 
-```typescript
+\`\`\`typescript
 startFlowsServer();
-```
+\`\`\`
 
 ## License
 
@@ -132,6 +104,6 @@ This project is licensed under the Apache License, Version 2.0. See the [LICENSE
 
 ## Conclusion
 
-This sample provides a basic demonstration of using Vertex AI plugins with Genkit for document indexing and retrieval. It can be extended and adapted to suit more complex use cases and integrations with other data sources and services.
+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
index 9d1ab6997..0837a56e9 100644
--- a/js/testapps/vertexai-reranker/package.json
+++ b/js/testapps/vertexai-reranker/package.json
@@ -21,14 +21,9 @@
     "@genkit-ai/evaluator": "workspace:*",
     "@genkit-ai/firebase": "workspace:*",
     "@genkit-ai/flow": "workspace:*",
-    "@genkit-ai/googleai": "workspace:*",
     "@genkit-ai/vertexai": "workspace:*",
-    "@google-cloud/bigquery": "^7.8.0",
     "dotenv": "^16.4.5",
     "express": "^4.19.2",
-    "genkitx-chromadb": "workspace:*",
-    "genkitx-langchain": "workspace:*",
-    "genkitx-pinecone": "workspace:*",
     "google-auth-library": "^9.11.0",
     "zod": "3.22.4"
   },

From ab4e127dbead2df4cd1f5c01a230a72ac4ee7eca Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Fri, 30 Aug 2024 15:23:36 +0100
Subject: [PATCH 6/8] fix: lift some reranker lines out of defineReranker

---
 js/plugins/vertexai/src/index.ts    |  2 +-
 js/plugins/vertexai/src/reranker.ts | 11 +++++------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/js/plugins/vertexai/src/index.ts b/js/plugins/vertexai/src/index.ts
index fb7044b3d..10c0dcaa4 100644
--- a/js/plugins/vertexai/src/index.ts
+++ b/js/plugins/vertexai/src/index.ts
@@ -256,7 +256,7 @@ export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin(
       projectId,
     };
 
-    const rerankers = vertexAiRerankers(rerankOptions);
+    const rerankers = await vertexAiRerankers(rerankOptions);
 
     return {
       models,
diff --git a/js/plugins/vertexai/src/reranker.ts b/js/plugins/vertexai/src/reranker.ts
index 17ec9cc5b..8361d13bd 100644
--- a/js/plugins/vertexai/src/reranker.ts
+++ b/js/plugins/vertexai/src/reranker.ts
@@ -75,9 +75,9 @@ export interface VertexRerankOptions {
  * @param {VertexRerankOptions<EmbedderCustomOptions>} params - The parameters for creating the rerankers.
  * @returns {RerankerAction<z.ZodTypeAny>[]} - An array of reranker actions.
  */
-export function vertexAiRerankers(
+export async function vertexAiRerankers(
   params: VertexRerankOptions
-): RerankerAction<z.ZodTypeAny>[] {
+): Promise<RerankerAction<z.ZodTypeAny>[]> {
   if (!params.pluginOptions) {
     throw new Error(
       'Plugin options are required to create Vertex AI rerankers'
@@ -94,6 +94,9 @@ export function vertexAiRerankers(
   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(
@@ -102,10 +105,6 @@ export function vertexAiRerankers(
         configSchema: VertexAIRerankerOptionsSchema.optional(),
       },
       async (query, documents, _options) => {
-        const auth = new GoogleAuth();
-        const client = await auth.getClient();
-        const projectId = await auth.getProjectId();
-
         const response = await client.request({
           method: 'POST',
           url: getRerankEndpoint(

From 29fcaba51928c9734e5a045bda23429c10f5b3d7 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Mon, 9 Sep 2024 09:18:00 +0100
Subject: [PATCH 7/8] docs(js): add reranker section

---
 docs/rag.md | 86 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/docs/rag.md b/docs/rag.md
index 4d07708b6..02f55fde2 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

From daef2e58715558bb487f709d80fb4b21a8634330 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Mon, 9 Sep 2024 09:56:18 +0100
Subject: [PATCH 8/8] test(js): add reranker tests

---
 js/ai/tests/reranker/reranker_test.ts | 221 ++++++++++++++++++++++++++
 1 file changed, 221 insertions(+)
 create mode 100644 js/ai/tests/reranker/reranker_test.ts

diff --git a/js/ai/tests/reranker/reranker_test.ts b/js/ai/tests/reranker/reranker_test.ts
new file mode 100644
index 000000000..0226264e0
--- /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'
+        );
+      }
+    });
+  });
+});