diff --git a/src/api/providers/fetchers/modelCache.ts b/src/api/providers/fetchers/modelCache.ts index e48f6ec03b..4e296302d7 100644 --- a/src/api/providers/fetchers/modelCache.ts +++ b/src/api/providers/fetchers/modelCache.ts @@ -83,7 +83,7 @@ export const getModels = async (options: GetModelsOptions): Promise models = await getIOIntelligenceModels(options.apiKey) break case "watsonx": - models = await getWatsonxModels(options.apiKey, false) + models = await getWatsonxModels(options.apiKey) break default: { // Ensures router is exhaustively checked if RouterName is a strict union diff --git a/src/api/providers/fetchers/watsonx.ts b/src/api/providers/fetchers/watsonx.ts index e9255dbdbf..14e06b4828 100644 --- a/src/api/providers/fetchers/watsonx.ts +++ b/src/api/providers/fetchers/watsonx.ts @@ -16,8 +16,6 @@ import { WatsonXAI } from "@ibm-cloud/watsonx-ai" */ export async function getWatsonxModels( apiKey: string, - embedded: boolean, - projectId?: string, baseUrl?: string, platform: "ibmCloud" | "cloudPak" = "ibmCloud", username?: string, @@ -63,13 +61,7 @@ export async function getWatsonxModels( let knownModels: Record = {} try { - let response - if (embedded) { - response = await service.listFoundationModelSpecs({ filters: "function_embedding" }) - } else { - response = await service.listFoundationModelSpecs({ filters: "!function_embedding" }) - } - + const response = await service.listFoundationModelSpecs({ filters: "!function_embedding" }) if (response && response.result) { const result = response.result as any const modelsList = result.resources @@ -113,3 +105,87 @@ export async function getWatsonxModels( return {} } } + +/** + * Fetches available embedded watsonx models + * + * @param apiKey - The watsonx API key (for IBM Cloud or Cloud Pak with API key auth) + * @param baseUrl - Optional base URL for the watsonx API + * @param platform - Optional platform type (ibmCloud or cloudPak) + * @param username - Optional username for Cloud Pak for Data + * @param password - Optional password for Cloud Pak for Data (when using password auth) + * @returns A promise resolving to an object with model IDs as keys and model info as values + */ +export async function getEmbeddedWatsonxModels( + apiKey: string, + baseUrl?: string, + platform: "ibmCloud" | "cloudPak" = "ibmCloud", + username?: string, + password?: string, +): Promise> { + try { + let options: any = { + version: "2024-05-31", + } + + if (platform === "ibmCloud" || !platform) { + if (apiKey) { + options.authenticator = new IamAuthenticator({ + apikey: apiKey, + }) + } else { + return {} + } + } else if (platform === "cloudPak") { + if (!baseUrl) { + throw new Error("Base URL is required for IBM Cloud Pak for Data") + } + + if (username) { + if (password) { + options.authenticator = new CloudPakForDataAuthenticator({ + url: baseUrl, + username: username, + password: password, + }) + } else if (apiKey) { + options.authenticator = new CloudPakForDataAuthenticator({ + url: baseUrl, + username: username, + apikey: apiKey, + }) + } + } + } + + const service = WatsonXAI.newInstance(options) + + let knownModels: Record = {} + + try { + const response = await service.listFoundationModelSpecs({ filters: "function_embedding" }) + if (response && response.result) { + const result = response.result as any + + const modelsList = result.models || result.resources || result.foundation_models || [] + + if (Array.isArray(modelsList)) { + for (const model of modelsList) { + const modelId = model.id || model.name || model.model_id + if (modelId.startsWith("ibm")) { + const dimension = model.model_limits.embedding_dimension + knownModels[modelId] = { dimension } + } + } + } + } + } catch (error) { + console.warn("Error fetching embedded models from IBM watsonx API:", error) + return {} + } + return knownModels + } catch (apiError) { + console.error("Error fetching embedded IBM watsonx models:", apiError) + return {} + } +} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 2fcba11b38..906ad05c0e 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -55,7 +55,7 @@ const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"]) import { MarketplaceManager, MarketplaceItemType } from "../../services/marketplace" import { setPendingTodoList } from "../tools/updateTodoListTool" -import { getWatsonxModels } from "../../api/providers/fetchers/watsonx" +import { getEmbeddedWatsonxModels, getWatsonxModels } from "../../api/providers/fetchers/watsonx" export const webviewMessageHandler = async ( provider: ClineProvider, @@ -758,8 +758,6 @@ export const webviewMessageHandler = async ( const watsonxModels = await getWatsonxModels( apiKey, - false, - projectId, effectiveBaseUrl, platform, username, @@ -816,10 +814,8 @@ export const webviewMessageHandler = async ( effectiveBaseUrl = regionToUrl[region] || "https://us-south.ml.cloud.ibm.com" } - const watsonxModels = await getWatsonxModels( + const watsonxModels = await getEmbeddedWatsonxModels( apiKey, - true, - projectId, effectiveBaseUrl, platform as "ibmCloud" | "cloudPak", username, diff --git a/src/services/code-index/embedders/watsonx.ts b/src/services/code-index/embedders/watsonx.ts index fe9c34954f..3d8291790b 100644 --- a/src/services/code-index/embedders/watsonx.ts +++ b/src/services/code-index/embedders/watsonx.ts @@ -1,16 +1,12 @@ import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" import { MAX_ITEM_TOKENS } from "../constants" import { t } from "../../../i18n" -import { TelemetryEventName } from "@roo-code/types" -import { TelemetryService } from "@roo-code/telemetry" import { WatsonXAI } from "@ibm-cloud/watsonx-ai" import { IamAuthenticator, CloudPakForDataAuthenticator } from "ibm-cloud-sdk-core" /** * IBM watsonx embedder implementation using the native IBM Cloud watsonx.ai package. * - * Supported models: - * - ibm/slate-125m-english-rtrvr-v2 (dimension: 768) */ export class WatsonxEmbedder implements IEmbedder { private readonly watsonxClient: WatsonXAI @@ -91,101 +87,135 @@ export class WatsonxEmbedder implements IEmbedder { } /** - * Creates embeddings for the given texts using watsonx's embedding API - * @param texts Array of text strings to embed - * @param model Optional model identifier (uses constructor model if not provided) - * @returns Promise resolving to embedding response + * Gets the expected dimension for a given model ID + * @param modelId The model ID to get the dimension for + * @returns The expected dimension for the model, or 768 if unknown */ + private getExpectedDimension(modelId: string): number { + // Known dimensions for watsonx models + const knownDimensions: Record = { + "ibm/slate-125m-english-rtrvr-v2": 768, + "ibm/slate-125m-english-rtrvr": 768, + "ibm/slate-30m-english-rtrvr-v2": 384, + "ibm/slate-30m-english-rtrvr": 384, + "ibm/granite-embedding-107m-multilingual": 384, + "ibm/granite-embedding-278M-multilingual": 768, + } + return knownDimensions[modelId] || 768 + } + async createEmbeddings(texts: string[], model?: string): Promise { const MAX_RETRIES = 3 const INITIAL_DELAY_MS = 1000 - - try { - const modelToUse = model || this.modelId - - const embeddings: number[][] = [] - let promptTokens = 0 - let totalTokens = 0 - - for (const text of texts) { - if (!text.trim()) { - embeddings.push([]) - continue - } - - const estimatedTokens = Math.ceil(text.length / 4) - if (estimatedTokens > MAX_ITEM_TOKENS) { - console.warn( - t("embeddings:textExceedsTokenLimit", { - index: texts.indexOf(text), - itemTokens: estimatedTokens, - maxTokens: MAX_ITEM_TOKENS, - }), - ) - embeddings.push([]) - continue - } - - let lastError - for (let attempt = 0; attempt < MAX_RETRIES; attempt++) { - try { - const response = await this.watsonxClient.embedText({ - modelId: modelToUse, - inputs: [text], - projectId: this.projectId, - parameters: { - truncate_input_tokens: MAX_ITEM_TOKENS, - return_options: { - input_text: true, + const MAX_CONCURRENT_REQUESTS = 1 + const REQUEST_DELAY_MS = 500 + const modelToUse = model || this.modelId + const embeddings: number[][] = [] + let promptTokens = 0 + let totalTokens = 0 + for (let i = 0; i < texts.length; i += MAX_CONCURRENT_REQUESTS) { + const batch = texts.slice(i, i + MAX_CONCURRENT_REQUESTS) + const batchResults = await Promise.all( + batch.map(async (text, batchIndex) => { + const textIndex = i + batchIndex + if (!text.trim()) { + return { index: textIndex, embedding: [], tokens: 0 } + } + const estimatedTokens = Math.ceil(text.length / 4) + if (estimatedTokens > MAX_ITEM_TOKENS) { + console.warn( + t("embeddings:textExceedsTokenLimit", { + index: textIndex, + itemTokens: estimatedTokens, + maxTokens: MAX_ITEM_TOKENS, + }), + ) + return { index: textIndex, embedding: [], tokens: 0 } + } + let lastError + for (let attempt = 0; attempt < MAX_RETRIES; attempt++) { + try { + await delay(1000) + const response = await this.watsonxClient.embedText({ + modelId: modelToUse, + inputs: [text], + projectId: this.projectId, + parameters: { + truncate_input_tokens: MAX_ITEM_TOKENS, + return_options: { + input_text: true, + }, }, - }, - }) + }) + if (response.result && response.result.results && response.result.results.length > 0) { + let embedding = response.result.results[0].embedding + if (!embedding || embedding.length === 0) { + console.error(`Empty embedding returned for text at index ${textIndex}`) + const expectedDimension = this.getExpectedDimension(modelToUse) + if (expectedDimension > 0) { + embedding = new Array(expectedDimension).fill(0.0001) + } else { + throw new Error(`Cannot determine expected dimension for model ${modelToUse}`) + } + } + if (!embedding || embedding.length === 0) { + throw new Error("Failed to create valid embedding") + } - if (response.result && response.result.results && response.result.results.length > 0) { - embeddings.push(response.result.results[0].embedding) - - if (response.result.input_token_count) { - promptTokens += response.result.input_token_count - totalTokens += response.result.input_token_count + const tokens = response.result.input_token_count || 0 + return { index: textIndex, embedding, tokens } + } else { + console.warn(`No embedding results for text at index ${textIndex}`) + const expectedDimension = this.getExpectedDimension(modelToUse) + if (expectedDimension > 0) { + console.log(`Creating fallback embedding with dimension ${expectedDimension}`) + const fallbackEmbedding = new Array(expectedDimension).fill(0.0001) + return { index: textIndex, embedding: fallbackEmbedding, tokens: 0 } + } else { + return { index: textIndex, embedding: [], tokens: 0 } + } } - break - } else { - embeddings.push([]) - break - } - } catch (error) { - lastError = error + } catch (error) { + lastError = error - if (attempt < MAX_RETRIES - 1) { - const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempt) - console.warn( - `IBM watsonx API call failed, retrying in ${delayMs}ms (attempt ${attempt + 1}/${MAX_RETRIES})`, - ) - await new Promise((resolve) => setTimeout(resolve, delayMs)) + if (attempt < MAX_RETRIES - 1) { + const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempt) + console.warn( + `IBM watsonx API call failed, retrying in ${delayMs}ms (attempt ${attempt + 1}/${MAX_RETRIES})`, + ) + } } } - } - if (lastError && embeddings.length < texts.indexOf(text) + 1) { + console.error( + `Failed to embed text at index ${textIndex} after ${MAX_RETRIES} attempts:`, + lastError, + ) + return { index: textIndex, embedding: [], tokens: 0 } + }), + ) + + if (i + MAX_CONCURRENT_REQUESTS < texts.length) { + await new Promise((resolve) => setTimeout(resolve, REQUEST_DELAY_MS * 2)) + } + + // Process batch results + for (const result of batchResults) { + while (embeddings.length <= result.index) { embeddings.push([]) - console.error(`Failed to embed text after ${MAX_RETRIES} attempts:`, lastError) } - } - return { - embeddings, - usage: { - promptTokens, - totalTokens, - }, + embeddings[result.index] = result.embedding + promptTokens += result.tokens + totalTokens += result.tokens } - } catch (error) { - TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { - error: error instanceof Error ? error.message : String(error), - stack: error instanceof Error ? error.stack : undefined, - location: "WatsonxEmbedder:createEmbeddings", - }) - throw error + } + return { + embeddings, + usage: { + promptTokens, + totalTokens, + }, } } @@ -196,9 +226,6 @@ export class WatsonxEmbedder implements IEmbedder { async validateConfiguration(): Promise<{ valid: boolean; error?: string }> { try { const testText = "test" - - console.log("Testing IBM watsonx.ai configuration with model:", this.modelId) - const response = await this.watsonxClient.embedText({ modelId: this.modelId, inputs: [testText], @@ -218,17 +245,9 @@ export class WatsonxEmbedder implements IEmbedder { error: "embeddings:validation.invalidResponse", } } - - console.log("IBM watsonx configuration validated successfully") return { valid: true } } catch (error) { console.error("IBM watsonx validation error:", error) - TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { - error: error instanceof Error ? error.message : String(error), - stack: error instanceof Error ? error.stack : undefined, - location: "WatsonxEmbedder:validateConfiguration", - }) - let errorMessage = "embeddings:validation.unknownError" let errorDetails = "" @@ -246,7 +265,6 @@ export class WatsonxEmbedder implements IEmbedder { errorMessage = "embeddings:validation.invalidModelId" } } - return { valid: false, error: `${errorMessage} (${errorDetails})`, @@ -265,16 +283,8 @@ export class WatsonxEmbedder implements IEmbedder { const knownModels: Record = { "ibm/slate-125m-english-rtrvr-v2": { dimension: 768 }, } - try { - const response = await this.watsonxClient.listFoundationModelSpecs() - - console.log( - "IBM watsonx API response structure:", - Object.keys(response || {}).join(", "), - Object.keys(response?.result || {}).join(", "), - ) - + const response = await this.watsonxClient.listFoundationModelSpecs({ filters: "function_embedding" }) if (response && response.result) { const result = response.result as any @@ -283,24 +293,14 @@ export class WatsonxEmbedder implements IEmbedder { if (Array.isArray(modelsList)) { for (const model of modelsList) { const modelId = model.id || model.name || model.model_id - const modelInfo = JSON.stringify(model).toLowerCase() - if ( - modelId && - (modelInfo.includes("embed") || - modelInfo.includes("rtrvr") || - modelInfo.includes("retriev")) - ) { - const dimension = model.dimension || model.vector_size || model.embedding_size || 1536 - knownModels[modelId] = { dimension } - } + const dimension = model.model_limits.embedding_dimension + knownModels[modelId] = { dimension } } } } } catch (apiError) { console.warn("Error fetching models from IBM watsonx API:", apiError) } - - console.log(`Found ${Object.keys(knownModels).length} IBM watsonx embedding models`) return knownModels } catch (error) { console.error("Error in getAvailableModels:", error) @@ -319,3 +319,7 @@ export class WatsonxEmbedder implements IEmbedder { } } } + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)) +} diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index 5c8c6b93bb..a9d280a619 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -152,7 +152,7 @@ export interface ExtensionMessage { lmStudioModels?: ModelRecord vsCodeLmModels?: { vendor?: string; family?: string; version?: string; id?: string }[] watsonxModels?: Record - embeddedWatsonxModels?: Record + embeddedWatsonxModels?: Record huggingFaceModels?: Array<{ id: string object: string diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index b3172814c9..c2771a1fa3 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -54,7 +54,12 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { "codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 }, }, watsonx: { + "ibm/granite-embedding-107m-multilingual": { dimension: 384, scoreThreshold: 0.4 }, + "ibm/granite-embedding-278M-multilingual": { dimension: 768, scoreThreshold: 0.4 }, "ibm/slate-125m-english-rtrvr-v2": { dimension: 768, scoreThreshold: 0.4 }, + "ibm/slate-125m-english-rtrvr": { dimension: 768, scoreThreshold: 0.4 }, + "ibm/slate-30m-english-rtrvr-v2": { dimension: 384, scoreThreshold: 0.4 }, + "ibm/slate-30m-english-rtrvr": { dimension: 384, scoreThreshold: 0.4 }, }, } diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 173d182e88..fbb2c7a84a 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -325,7 +325,7 @@ export const CodeIndexPopover: React.FC = ({ } else if (event.data.type === "embeddedWatsonxModels" && event.data.embeddedWatsonxModels) { try { console.log("Received IBM Embeded watsonx models:", event.data.embeddedWatsonxModels) - const embeddedWatsonxModels: Record = {} + let embeddedWatsonxModels: Record = {} if ( !event.data.embeddedWatsonxModels || Object.keys(event.data.embeddedWatsonxModels).length === 0 @@ -333,11 +333,7 @@ export const CodeIndexPopover: React.FC = ({ console.warn("No models received from server, adding default model") embeddedWatsonxModels["ibm/slate-125m-english-rtrvr-v2"] = { dimension: 768 } } else { - Object.keys(event.data.embeddedWatsonxModels).forEach((modelId) => { - embeddedWatsonxModels[modelId] = { - dimension: 768, - } - }) + embeddedWatsonxModels = event.data.embeddedWatsonxModels } if (codebaseIndexModels) { codebaseIndexModels.watsonx = { ...embeddedWatsonxModels }