mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
0892455db2
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
193 lines
7.7 KiB
TypeScript
193 lines
7.7 KiB
TypeScript
/**
|
|
* Defines profiles for different embedding models, including their dimensions.
|
|
*/
|
|
|
|
import type { EmbedderProvider, EmbeddingModelProfiles } from "@roo-code/types"
|
|
|
|
// Example profiles - expand this list as needed
|
|
export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
|
|
openai: {
|
|
"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
|
|
"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
|
|
},
|
|
ollama: {
|
|
"nomic-embed-text": { dimension: 768, scoreThreshold: 0.4 },
|
|
"nomic-embed-code": {
|
|
dimension: 3584,
|
|
scoreThreshold: 0.15,
|
|
queryPrefix: "Represent this query for searching relevant code: ",
|
|
},
|
|
"mxbai-embed-large": { dimension: 1024, scoreThreshold: 0.4 },
|
|
"all-minilm": { dimension: 384, scoreThreshold: 0.4 },
|
|
// Add default Ollama model if applicable, e.g.:
|
|
// 'default': { dimension: 768 } // Assuming a default dimension
|
|
},
|
|
"openai-compatible": {
|
|
"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
|
|
"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"nomic-embed-code": {
|
|
dimension: 3584,
|
|
scoreThreshold: 0.15,
|
|
queryPrefix: "Represent this query for searching relevant code: ",
|
|
},
|
|
},
|
|
gemini: {
|
|
"gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
|
|
// Deprecated: text-embedding-004 is migrated to gemini-embedding-001 in GeminiEmbedder
|
|
// Kept here for backward-compatible dimension lookup in createVectorStore()
|
|
"text-embedding-004": { dimension: 3072, scoreThreshold: 0.4 },
|
|
},
|
|
mistral: {
|
|
"codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
|
|
},
|
|
"vercel-ai-gateway": {
|
|
// OpenAI models
|
|
"openai/text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"openai/text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
|
|
"openai/text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
|
|
// Cohere models
|
|
"cohere/embed-v4.0": { dimension: 1024, scoreThreshold: 0.4 },
|
|
// Google models
|
|
"google/gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
|
|
"google/text-embedding-005": { dimension: 768, scoreThreshold: 0.4 },
|
|
"google/text-multilingual-embedding-002": { dimension: 768, scoreThreshold: 0.4 },
|
|
// Amazon models
|
|
"amazon/titan-embed-text-v2": { dimension: 1024, scoreThreshold: 0.4 },
|
|
// Mistral models
|
|
"mistral/codestral-embed": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"mistral/mistral-embed": { dimension: 1024, scoreThreshold: 0.4 },
|
|
},
|
|
bedrock: {
|
|
// Amazon Titan Embed models
|
|
"amazon.titan-embed-text-v1": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"amazon.titan-embed-text-v2:0": { dimension: 1024, scoreThreshold: 0.4 },
|
|
"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 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 },
|
|
},
|
|
openrouter: {
|
|
// OpenAI models via OpenRouter
|
|
"openai/text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
|
|
"openai/text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
|
|
"openai/text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
|
|
// Google models via OpenRouter
|
|
"google/gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
|
|
// Mistral models via OpenRouter
|
|
"mistralai/mistral-embed-2312": { dimension: 1024, scoreThreshold: 0.4 },
|
|
"mistralai/codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
|
|
// Qwen models via OpenRouter
|
|
"qwen/qwen3-embedding-0.6b": { dimension: 1024, scoreThreshold: 0.4 },
|
|
"qwen/qwen3-embedding-4b": { dimension: 2560, scoreThreshold: 0.4 },
|
|
"qwen/qwen3-embedding-8b": { dimension: 4096, scoreThreshold: 0.4 },
|
|
},
|
|
}
|
|
|
|
/**
|
|
* Retrieves the embedding dimension for a given provider and model ID.
|
|
* @param provider The embedder provider (e.g., "openai").
|
|
* @param modelId The specific model ID (e.g., "text-embedding-3-small").
|
|
* @returns The dimension size or undefined if the model is not found.
|
|
*/
|
|
export function getModelDimension(provider: EmbedderProvider, modelId: string): number | undefined {
|
|
const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
|
|
if (!providerProfiles) {
|
|
console.warn(`Provider not found in profiles: ${provider}`)
|
|
return undefined
|
|
}
|
|
|
|
const modelProfile = providerProfiles[modelId]
|
|
if (!modelProfile) {
|
|
// Don't warn here, as it might be a custom model ID not in our profiles
|
|
// console.warn(`Model not found for provider ${provider}: ${modelId}`)
|
|
return undefined // Or potentially return a default/fallback dimension?
|
|
}
|
|
|
|
return modelProfile.dimension
|
|
}
|
|
|
|
/**
|
|
* Retrieves the score threshold for a given provider and model ID.
|
|
* @param provider The embedder provider (e.g., "openai").
|
|
* @param modelId The specific model ID (e.g., "text-embedding-3-small").
|
|
* @returns The score threshold or undefined if the model is not found.
|
|
*/
|
|
export function getModelScoreThreshold(provider: EmbedderProvider, modelId: string): number | undefined {
|
|
const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
|
|
if (!providerProfiles) {
|
|
return undefined
|
|
}
|
|
|
|
const modelProfile = providerProfiles[modelId]
|
|
return modelProfile?.scoreThreshold
|
|
}
|
|
|
|
/**
|
|
* Retrieves the query prefix for a given provider and model ID.
|
|
* @param provider The embedder provider (e.g., "openai").
|
|
* @param modelId The specific model ID (e.g., "nomic-embed-code").
|
|
* @returns The query prefix or undefined if the model doesn't require one.
|
|
*/
|
|
export function getModelQueryPrefix(provider: EmbedderProvider, modelId: string): string | undefined {
|
|
const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
|
|
if (!providerProfiles) {
|
|
return undefined
|
|
}
|
|
|
|
const modelProfile = providerProfiles[modelId]
|
|
return modelProfile?.queryPrefix
|
|
}
|
|
|
|
/**
|
|
* Gets the default *specific* embedding model ID based on the provider.
|
|
* Does not include the provider prefix.
|
|
* Currently defaults to OpenAI's 'text-embedding-3-small'.
|
|
* TODO: Make this configurable or more sophisticated.
|
|
* @param provider The embedder provider.
|
|
* @returns The default specific model ID for the provider (e.g., "text-embedding-3-small").
|
|
*/
|
|
export function getDefaultModelId(provider: EmbedderProvider): string {
|
|
switch (provider) {
|
|
case "openai":
|
|
case "openai-compatible":
|
|
return "text-embedding-3-small"
|
|
|
|
case "ollama": {
|
|
// Choose a sensible default for Ollama, e.g., the first one listed or a specific one
|
|
const ollamaModels = EMBEDDING_MODEL_PROFILES.ollama
|
|
const defaultOllamaModel = ollamaModels && Object.keys(ollamaModels)[0]
|
|
if (defaultOllamaModel) {
|
|
return defaultOllamaModel
|
|
}
|
|
// Fallback if no Ollama models are defined (shouldn't happen with the constant)
|
|
console.warn("No default Ollama model found in profiles.")
|
|
// Return a placeholder or throw an error, depending on desired behavior
|
|
return "unknown-default" // Placeholder specific model ID
|
|
}
|
|
|
|
case "gemini":
|
|
return "gemini-embedding-001"
|
|
|
|
case "mistral":
|
|
return "codestral-embed-2505"
|
|
|
|
case "vercel-ai-gateway":
|
|
return "openai/text-embedding-3-large"
|
|
|
|
case "bedrock":
|
|
return "amazon.titan-embed-text-v2:0"
|
|
case "openrouter":
|
|
return "openai/text-embedding-3-large"
|
|
|
|
default:
|
|
// Fallback for unknown providers
|
|
console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)
|
|
return "text-embedding-3-small"
|
|
}
|
|
}
|