Merge pull request #4419 from mq200/feat/watsonx_reranker_integration

add watsonx reranker integration

Author: sestinj

core/config/load.ts

@@ -505,7 +505,11 @@ async function intermediateToFinalConfig(
         model: "rerank-2",
         ...params,
       };
-      return new rerankerClass(llmOptions);
+      return new rerankerClass(llmOptions, (url: string | URL, init: any) =>
+        fetchwithRequestOptions(url, init, {
+          ...params?.requestOptions,
+        }),
+      );
     }
     return null;
   }
@@ -994,5 +998,6 @@ export {
   finalToBrowserConfig,
   intermediateToFinalConfig,
   loadContinueConfigFromJson,
-  type BrowserSerializedContinueConfig,
+  type BrowserSerializedContinueConfig
 };
+

core/context/allRerankers.ts

@@ -5,12 +5,14 @@ import HuggingFaceTEI from "../llm/llms/HuggingFaceTEI";
 import { LLMReranker } from "../llm/llms/llm";
 import ContinueProxy from "../llm/llms/stubs/ContinueProxy";
 import Voyage from "../llm/llms/Voyage";
+import WatsonX from "../llm/llms/WatsonX";
 
 export const AllRerankers: { [key: string]: any } = {
   cohere: Cohere,
   bedrock: Bedrock,
   llm: LLMReranker,
   voyage: Voyage,
+  watsonx: WatsonX,
   "free-trial": FreeTrial,
   "huggingface-tei": HuggingFaceTEI,
   "continue-proxy": ContinueProxy,

core/control-plane/schema.ts

@@ -19,6 +19,7 @@ const modelDescriptionSchema = z.object({
     "nebius",
     "siliconflow",
     "scaleway",
+    "watsonx"
   ]),
   model: z.string(),
   apiKey: z.string().optional(),
@@ -88,6 +89,7 @@ const embeddingsProviderSchema = z.object({
     "nebius",
     "siliconflow",
     "scaleway",
+    "watsonx"
   ]),
   apiBase: z.string().optional(),
   apiKey: z.string().optional(),
@@ -109,7 +111,7 @@ const embeddingsProviderSchema = z.object({
 });
 
 const rerankerSchema = z.object({
-  name: z.enum(["cohere", "voyage", "llm"]),
+  name: z.enum(["cohere", "voyage", "llm", "watsonx"]),
   params: z.record(z.any()).optional(),
 });
 

core/llm/llms/WatsonX.ts

@@ -1,4 +1,4 @@
-import { ChatMessage, CompletionOptions, LLMOptions } from "../../index.js";
+import { ChatMessage, Chunk, CompletionOptions, LLMOptions } from "../../index.js";
 import { renderChatMessage } from "../../util/messageContent.js";
 import { BaseLLM } from "../index.js";
 import { streamResponse } from "../stream.js";
@@ -138,6 +138,26 @@ class WatsonX extends BaseLLM {
     };
   }
 
+  protected async updateWatsonxToken() {
+        var now = new Date().getTime() / 1000;
+        if (
+          watsonxToken === undefined ||
+          now > watsonxToken.expiration ||
+          watsonxToken.token === undefined
+        ) {
+          watsonxToken = await this.getBearerToken();
+        } else {
+          console.log(
+            `Reusing token (expires in ${
+              (watsonxToken.expiration - now) / 60
+            } mins)`,
+          );
+        }
+        if (watsonxToken.token === undefined) {
+          throw new Error("Something went wrong. Check your credentials, please.");
+        }
+    }
+
   protected async _complete(
     prompt: string,
     signal: AbortSignal,
@@ -174,23 +194,8 @@ class WatsonX extends BaseLLM {
     signal: AbortSignal,
     options: CompletionOptions,
   ): AsyncGenerator<ChatMessage> {
-    var now = new Date().getTime() / 1000;
-    if (
-      watsonxToken === undefined ||
-      now > watsonxToken.expiration ||
-      watsonxToken.token === undefined
-    ) {
-      watsonxToken = await this.getBearerToken();
-    } else {
-      console.log(
-        `Reusing token (expires in ${
-          (watsonxToken.expiration - now) / 60
-        } mins)`,
-      );
-    }
-    if (watsonxToken.token === undefined) {
-      throw new Error("Something went wrong. Check your credentials, please.");
-    }
+    await this.updateWatsonxToken();
+
     const stopSequences =
       options.stop?.slice(0, 6) ??
       (options.model?.includes("granite") ? ["Question:"] : []);
@@ -267,20 +272,8 @@ class WatsonX extends BaseLLM {
   }
 
   protected async _embed(chunks: string[]): Promise<number[][]> {
-    var now = new Date().getTime() / 1000;
-    if (
-      watsonxToken === undefined ||
-      now > watsonxToken.expiration ||
-      watsonxToken.token === undefined
-    ) {
-      watsonxToken = await this.getBearerToken();
-    } else {
-      console.log(
-        `Reusing token (expires in ${
-          (watsonxToken.expiration - now) / 60
-        } mins)`,
-      );
-    }
+    await this.updateWatsonxToken();
+
     const payload: any = {
       inputs: chunks,
       parameters: {
@@ -318,6 +311,66 @@ class WatsonX extends BaseLLM {
     }
     return embeddings.map((e: any) => e.embedding);
   }
+
+
+  async rerank(query: string, chunks: Chunk[]): Promise<number[]> {
+      if (!query || !chunks.length) {
+        throw new Error("Query and chunks must not be empty");
+      }
+      try {
+        await this.updateWatsonxToken();
+        
+        const headers = {
+          "Content-Type": "application/json",
+          Authorization: `${
+            watsonxToken.expiration === -1 ? "ZenApiKey" : "Bearer"
+          } ${watsonxToken.token}`,
+        };
+    
+        const payload: any = {
+          inputs: chunks.map((chunk) => ({ text: chunk.content })),
+          query: query,
+          parameters: {
+            truncate_input_tokens: 500,
+            return_options: {
+                top_n: chunks.length
+            },
+          },
+          model_id: this.model,
+          project_id: this.projectId,
+        };
+
+        const resp = await this.fetch(
+          new URL(
+            `${this.apiBase}/ml/v1/text/rerank?version=${this.apiVersion}`,
+          ),
+          {
+            method: "POST",
+            headers: headers,
+            body: JSON.stringify(payload),
+          },
+        );
+
+        if (!resp.ok) {
+          throw new Error(`Failed to rerank chunks: ${await resp.text()}`);
+        }
+        const data = await resp.json();
+        const ranking = data.results;
+
+        if (!ranking) {
+          throw new Error("Empty response received from Watsonx");
+        }
+
+        // Sort results by index to maintain original order
+        return ranking
+          .sort((a: any, b: any) => a.index - b.index)
+          .map((result: any) => result.score);
+
+      } catch (error) {
+        console.error("Error in WatsonxReranker.rerank:", error);
+        throw error;
+      }
+  }
 }
 
 export default WatsonX;

docs/docs/customize/model-providers/more/watsonx.mdx

@@ -244,3 +244,37 @@ To view the list of available embeddings models, visit [this page](https://datap
   ```
 	</TabItem>
 </Tabs>
+
+
+## Reranker
+
+<Tabs groupId="config-example">
+	<TabItem value="yaml" label="YAML">
+  ```yaml title="config.yaml"
+  models:
+    - name: Watsonx Reranker
+      provider: watsonx
+      model: cross-encoder/ms-marco-minilm-l-12-v2
+      apiBase: https://us-south.ml.cloud.ibm.com
+      projectId: PROJECT_ID
+      apiKey: API_KEY/ZENAPI_KEY/USERNAME:PASSWORD
+      apiVersion: 2024-03-14
+  ```
+	</TabItem>
+	<TabItem value="json" label="JSON">
+  ```json title="config.json"
+  {
+    "reranker": {
+      "name": "watsonx",
+      "params": {
+        "model": "cross-encoder/ms-marco-minilm-l-12-v2",
+        "apiBase": "watsonx endpoint e.g. https://us-south.ml.cloud.ibm.com",
+        "projectId": "PROJECT_ID",
+        "apiKey": "API_KEY/ZENAPI_KEY/USERNAME:PASSWORD",
+        "apiVersion": "2024-03-14"
+      }
+    }
+  }
+  ```
+	</TabItem>
+</Tabs>

extensions/vscode/config_schema.json

@@ -2881,6 +2881,7 @@
                 "bedrock",
                 "cohere",
                 "voyage",
+                "watsonx",
                 "llm",
                 "free-trial",
                 "huggingface-tei"

packages/config-types/src/index.ts

@@ -58,6 +58,7 @@ export const modelDescriptionSchema = z.object({
     "continue-proxy",
     "nebius",
     "scaleway",
+    "watsonx"
   ]),
   model: z.string(),
   apiKey: z.string().optional(),
@@ -111,6 +112,7 @@ export const embeddingsProviderSchema = z.object({
     "continue-proxy",
     "nebius",
     "scaleway",
+    "watsonx"
   ]),
   apiBase: z.string().optional(),
   apiKey: z.string().optional(),
@@ -173,7 +175,7 @@ export const contextProviderSchema = z.object({
 export type ContextProvider = z.infer<typeof contextProviderSchema>;
 
 export const rerankerSchema = z.object({
-  name: z.enum(["cohere", "voyage", "llm", "continue-proxy"]),
+  name: z.enum(["cohere", "voyage", "watsonx", "llm", "continue-proxy"]),
   params: z.record(z.any()).optional(),
 });
 export type Reranker = z.infer<typeof rerankerSchema>;