diff --git a/core/context/rerankers/bedrock.ts b/core/context/rerankers/bedrock.ts new file mode 100644 index 0000000000..071910bd9c --- /dev/null +++ b/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 { + 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()(); + } + } +} diff --git a/core/context/rerankers/index.ts b/core/context/rerankers/index.ts index 9b1db230d8..ba6af849c3 100644 --- a/core/context/rerankers/index.ts +++ b/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, -}; +}; \ No newline at end of file diff --git a/core/index.d.ts b/core/index.d.ts index 06fbd7275a..6160f57c31 100644 --- a/core/index.d.ts +++ b/core/index.d.ts @@ -1086,6 +1086,7 @@ export interface EmbeddingsProvider { export type RerankerName = | "cohere" + | "bedrock" | "voyage" | "llm" | "free-trial" diff --git a/core/indexing/docs/article.ts b/core/indexing/docs/article.ts index 1bc74c62ec..2ae0e79953 100644 --- a/core/indexing/docs/article.ts +++ b/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); } diff --git a/docs/docs/customize/model-providers/top-level/bedrock.md b/docs/docs/customize/model-providers/top-level/bedrock.md index 8d46c5ad90..fb88d5ae02 100644 --- a/docs/docs/customize/model-providers/top-level/bedrock.md +++ b/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 diff --git a/docs/docs/reference.md b/docs/docs/reference.md index 0a5c02aad7..df3f08e75c 100644 --- a/docs/docs/reference.md +++ b/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 diff --git a/extensions/vscode/config_schema.json b/extensions/vscode/config_schema.json index 86a9d54a23..36fe771da7 100644 --- a/extensions/vscode/config_schema.json +++ b/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": {