Merge pull request #3169 from AnnoyingTechnology/241203-bedrock-reran…

…king

Brings support for Bedrock's reranking models, Fixes #3152

core/context/rerankers/bedrock.ts

@@ -0,0 +1,99 @@
+import {
+  BedrockRuntimeClient,
+  InvokeModelCommand,
+} from "@aws-sdk/client-bedrock-runtime";
+import { fromIni } from "@aws-sdk/credential-providers";
+import { Chunk, Reranker, RerankerName } from "../../index.js";
+
+export class BedrockReranker implements Reranker {
+  name: RerankerName = "bedrock";
+
+  static defaultOptions = {
+    region: "us-east-1",
+    model: "amazon.rerank-v1:0",
+    profile: "bedrock",
+  };
+
+  private supportedModels = ["amazon.rerank-v1:0", "cohere.rerank-v3-5:0"];
+
+  constructor(
+    private readonly params: {
+      region?: string;
+      model?: string;
+      profile?: string;
+    } = {},
+  ) {
+    if (params.model && !this.supportedModels.includes(params.model)) {
+      throw new Error(
+        `Unsupported model: ${params.model}. Supported models are: ${this.supportedModels.join(", ")}`,
+      );
+    }
+  }
+
+  async rerank(query: string, chunks: Chunk[]): Promise<number[]> {
+    if (!query || !chunks.length) {
+      throw new Error("Query and chunks must not be empty");
+    }
+
+    try {
+      const credentials = await this._getCredentials();
+      const client = new BedrockRuntimeClient({
+        region: this.params.region ?? BedrockReranker.defaultOptions.region,
+        credentials,
+      });
+
+      const model = this.params.model ?? BedrockReranker.defaultOptions.model;
+
+      // Base payload for both models
+      const payload: any = {
+        query: query,
+        documents: chunks.map((chunk) => chunk.content),
+        top_n: chunks.length,
+      };
+
+      // Add api_version for Cohere model
+      if (model.startsWith("cohere.rerank")) {
+        payload.api_version = 2;
+      }
+
+      const input = {
+        body: JSON.stringify(payload),
+        modelId: model,
+        accept: "*/*",
+        contentType: "application/json",
+      };
+
+      const command = new InvokeModelCommand(input);
+      const response = await client.send(command);
+
+      if (!response.body) {
+        throw new Error("Empty response received from Bedrock");
+      }
+
+      const responseBody = JSON.parse(new TextDecoder().decode(response.body));
+
+      // Sort results by index to maintain original order
+      return responseBody.results
+        .sort((a: any, b: any) => a.index - b.index)
+        .map((result: any) => result.relevance_score);
+    } catch (error) {
+      console.error("Error in BedrockReranker.rerank:", error);
+      throw error;
+    }
+  }
+
+  private async _getCredentials() {
+    try {
+      const credentials = await fromIni({
+        profile: this.params.profile ?? BedrockReranker.defaultOptions.profile,
+        ignoreCache: true,
+      })();
+      return credentials;
+    } catch (e) {
+      console.warn(
+        `AWS profile with name ${this.params.profile ?? BedrockReranker.defaultOptions.profile} not found in ~/.aws/credentials, using default profile`,
+      );
+      return await fromIni()();
+    }
+  }
+}

core/context/rerankers/index.ts

@@ -1,5 +1,6 @@
 import { RerankerName } from "../../index.js";
 
+import { BedrockReranker } from "./bedrock.js";
 import { CohereReranker } from "./cohere.js";
 import { ContinueProxyReranker } from "./ContinueProxyReranker.js";
 import { FreeTrialReranker } from "./freeTrial.js";
@@ -9,9 +10,10 @@ import { VoyageReranker } from "./voyage.js";
 
 export const AllRerankers: { [key in RerankerName]: any } = {
   cohere: CohereReranker,
+  bedrock: BedrockReranker,
   llm: LLMReranker,
   voyage: VoyageReranker,
   "free-trial": FreeTrialReranker,
   "huggingface-tei": HuggingFaceTEIReranker,
   "continue-proxy": ContinueProxyReranker,
-};
+};

core/index.d.ts

@@ -1086,6 +1086,7 @@ export interface EmbeddingsProvider {
 
 export type RerankerName =
   | "cohere"
+  | "bedrock"
   | "voyage"
   | "llm"
   | "free-trial"

core/indexing/docs/article.ts

@@ -25,59 +25,75 @@ function breakdownArticleComponent(
   max_chunk_size: number,
 ): Chunk[] {
   const chunks: Chunk[] = [];
-
   const lines = article.body.split("\n");
   let startLine = 0;
   let endLine = 0;
   let content = "";
   let index = 0;
 
+  const createChunk = (
+    chunkContent: string,
+    chunkStartLine: number,
+    chunkEndLine: number,
+  ) => {
+    chunks.push({
+      content: chunkContent.trim(),
+      startLine: chunkStartLine,
+      endLine: chunkEndLine,
+      otherMetadata: {
+        title: cleanHeader(article.title),
+      },
+      index: index++,
+      filepath: new URL(
+        `${subpath}#${cleanFragment(article.title)}`,
+        url,
+      ).toString(),
+      digest: subpath,
+    });
+  };
+
   for (let i = 0; i < lines.length; i++) {
     const line = lines[i];
-    if (content.length + line.length <= max_chunk_size) {
+
+    // Handle oversized lines by splitting them
+    if (line.length > max_chunk_size) {
+      // First push any accumulated content
+      if (content.trim().length > 0) {
+        createChunk(content, startLine, endLine);
+        content = "";
+      }
+
+      // Split the long line into chunks
+      let remainingLine = line;
+      let subLineStart = i;
+      while (remainingLine.length > 0) {
+        const chunk = remainingLine.slice(0, max_chunk_size);
+        createChunk(chunk, subLineStart, i);
+        remainingLine = remainingLine.slice(max_chunk_size);
+      }
+      startLine = i + 1;
+      continue;
+    }
+
+    // Normal line handling
+    if (content.length + line.length + 1 <= max_chunk_size) {
       content += `${line}\n`;
       endLine = i;
     } else {
-      chunks.push({
-        content: content.trim(),
-        startLine: startLine,
-        endLine: endLine,
-        otherMetadata: {
-          title: cleanHeader(article.title),
-        },
-        index: index,
-        filepath: new URL(
-          `${subpath}#${cleanFragment(article.title)}`,
-          url,
-        ).toString(),
-        digest: subpath,
-      });
+      if (content.trim().length > 0) {
+        createChunk(content, startLine, endLine);
+      }
       content = `${line}\n`;
       startLine = i;
       endLine = i;
-      index += 1;
     }
   }
 
   // Push the last chunk
-  if (content) {
-    chunks.push({
-      content: content.trim(),
-      startLine: startLine,
-      endLine: endLine,
-      otherMetadata: {
-        title: cleanHeader(article.title),
-      },
-      index: index,
-      filepath: new URL(
-        `${subpath}#${cleanFragment(article.title)}`,
-        url,
-      ).toString(),
-      digest: subpath,
-    });
+  if (content.trim().length > 0) {
+    createChunk(content, startLine, endLine);
   }
 
-  // Don't use small chunks. Probably they're a mistake. Definitely they'll confuse the embeddings model.
   return chunks.filter((c) => c.content.trim().length > 20);
 }
 

docs/docs/customize/model-providers/top-level/bedrock.md

@@ -46,9 +46,19 @@ We recommend configuring [`amazon.titan-embed-text-v2:0`](https://docs.aws.amazo
 
 ## Reranking model
 
-Bedrock currently does not offer any reranking models.
+We recommend configuring `cohere.rerank-v3-5:0` as your reranking model, you may also use `amazon.rerank-v1:0`.
 
-[Click here](../../model-types/reranking.md) to see a list of reranking model providers.
+```json title="~/.continue/config.json"
+{
+  "reranker": {
+    "name": "bedrock",
+    "params": {
+      "model": "cohere.rerank-v3-5:0",
+      "region": "us-west-2"
+    }
+  }
+}
+```
 
 ## Authentication
 

docs/docs/reference.md

@@ -206,10 +206,11 @@ Configuration for the reranker model used in response ranking.
 
 **Properties:**
 
-- `name` (**required**): Reranker name, e.g., `cohere`, `voyage`, `llm`, `free-trial`, `huggingface-tei`
+- `name` (**required**): Reranker name, e.g., `cohere`, `voyage`, `llm`, `free-trial`, `huggingface-tei`, `bedrock`
 - `params`:
   - `model`: Model name
   - `apiKey`: Api key
+  - `region`: Region (for Bedrock only)
 
 Example
 

extensions/vscode/config_schema.json

@@ -2642,6 +2642,7 @@
           "properties": {
             "name": {
               "enum": [
+                "bedrock",
                 "cohere",
                 "voyage",
                 "llm",
@@ -2743,6 +2744,32 @@
                 }
               }
             },
+            {
+              "if": {
+                "properties": {
+                  "name": {
+                    "enum": ["bedrock"]
+                  }
+                },
+                "required": ["name"]
+              },
+              "then": {
+                "properties": {
+                  "params": {
+                    "type": "object",
+                    "properties": {
+                      "model": {
+                        "enum": [
+                          "cohere.rerank-v3-5:0",
+                          "amazon.rerank-v1:0"
+                        ]
+                      }
+                    },
+                    "required": ["model"]
+                  }
+                }
+              }
+            },
             {
               "if": {
                 "properties": {