scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(bedrock): add Cohere Embed v4 model and improve credential handling (Fixes #11823) (#11824)

Browse source 
feat(bedrock): add Cohere Embed v4 model and improve credential handling

- Add cohere.embed-v4:0 (1536-dim) to Bedrock embedding model profiles
- Add v4-specific request format (embedding_types: ["float"]) and response
  parsing (embeddings.float[0]) in BedrockEmbedder
- Replace fromEnv() with fromNodeProviderChain() for default credential
  chain when no AWS profile is specified, supporting SSO, IMDS, ECS, and
  other credential sources with built-in memoization
- Add unit tests for Cohere v4 request/response handling, credential
  provider selection, and v3 regression coverage

Fixes #11823
This commit is contained in:
cscvenkatmadurai 2026-03-05 17:45:17 -08:00 • committed by GitHub
parent 0e56afc764
commit 0892455db2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
src/services/code-index/embedders/bedrock.ts
View file

@ -1,5 +1,5 @@
import { BedrockRuntimeClient, InvokeModelCommand, InvokeModelCommandInput } from "@aws-sdk/client-bedrock-runtime"
import { fromEnv, fromIni } from "@aws-sdk/credential-providers"
import { fromIni, fromNodeProviderChain } from "@aws-sdk/credential-providers"
import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces"
import {
MAX_BATCH_TOKENS,
@ -38,7 +38,7 @@ export class BedrockEmbedder implements IEmbedder {
// Initialize the Bedrock client with credentials
// If profile is specified, use it; otherwise use default credential chain
const credentials = this.profile ? fromIni({ profile: this.profile }) : fromEnv()
const credentials = this.profile ? fromIni({ profile: this.profile }) : fromNodeProviderChain()
this.bedrockClient = new BedrockRuntimeClient({
userAgentAppId: `RooCode#${Package.version}`,
@ -209,10 +209,18 @@ export class BedrockEmbedder implements IEmbedder {
requestBody = {
inputText: text,
}
} else if (model.startsWith("cohere.embed")) {
} else if (model.startsWith("cohere.embed-v4")) {
// Cohere Embed v4 requires embedding_types parameter
requestBody = {
texts: [text],
input_type: "search_document", // or "search_query" depending on use case
input_type: "search_document",
embedding_types: ["float"],
}
} else if (model.startsWith("cohere.embed")) {
// Cohere Embed v3 format
requestBody = {
texts: [text],
input_type: "search_document",
}
} else {
// Default to Titan format
@ -248,10 +256,15 @@ export class BedrockEmbedder implements IEmbedder {
embedding: responseBody.embedding,
inputTextTokenCount: responseBody.inputTextTokenCount,
}
} else if (model.startsWith("cohere.embed-v4")) {
// Cohere Embed v4 returns { embeddings: { float: [[...]] } }
return {
embedding: responseBody.embeddings?.float?.[0] || responseBody.embeddings?.[0],
}
} else if (model.startsWith("cohere.embed")) {
// Cohere Embed v3 returns { embeddings: [[...]] }
return {
embedding: responseBody.embeddings[0],
// Cohere doesn't provide token count in response
}
} else {
// Default to Titan format

4
src/shared/embeddingModels.ts
View file

@ -66,7 +66,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
"amazon.titan-embed-image-v1": { dimension: 1024, scoreThreshold: 0.4 },
// Amazon Nova Embed models
"amazon.nova-2-multimodal-embeddings-v1:0": { dimension: 1024, scoreThreshold: 0.4 },
// Cohere models available through Bedrock
// Cohere Embed v4 (supports only text for now; multimodal image support planned)
"cohere.embed-v4:0": { dimension: 1536, scoreThreshold: 0.4 },
// Cohere Embed v3 models available through Bedrock
"cohere.embed-english-v3": { dimension: 1024, scoreThreshold: 0.4 },
"cohere.embed-multilingual-v3": { dimension: 1024, scoreThreshold: 0.4 },
},