diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index be7778f538..5bca785386 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -22,7 +22,7 @@ export const codebaseIndexConfigSchema = z.object({ codebaseIndexEnabled: z.boolean().optional(), codebaseIndexQdrantUrl: z.string().optional(), codebaseIndexEmbedderProvider: z - .enum(["openai", "ollama", "openai-compatible", "gemini", "mistral", "vercel-ai-gateway"]) + .enum(["openai", "ollama", "openai-compatible", "gemini", "mistral", "vercel-ai-gateway", "watsonx"]) .optional(), codebaseIndexEmbedderBaseUrl: z.string().optional(), codebaseIndexEmbedderModelId: z.string().optional(), @@ -51,6 +51,7 @@ export const codebaseIndexModelsSchema = z.object({ gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(), mistral: z.record(z.string(), z.object({ dimension: z.number() })).optional(), "vercel-ai-gateway": z.record(z.string(), z.object({ dimension: z.number() })).optional(), + watsonx: z.record(z.string(), z.object({ dimension: z.number() })).optional(), }) export type CodebaseIndexModels = z.infer @@ -68,6 +69,8 @@ export const codebaseIndexProviderSchema = z.object({ codebaseIndexGeminiApiKey: z.string().optional(), codebaseIndexMistralApiKey: z.string().optional(), codebaseIndexVercelAiGatewayApiKey: z.string().optional(), + codebaseIndexWatsonxApiKey: z.string().optional(), + codebaseIndexWatsonxProjectId: z.string().optional(), }) export type CodebaseIndexProvider = z.infer diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index 7e79855f7e..d5f8261174 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -205,6 +205,9 @@ export const SECRET_STATE_KEYS = [ "featherlessApiKey", "ioIntelligenceApiKey", "vercelAiGatewayApiKey", + "watsonxApiKey", + "codebaseIndexWatsonxApiKey", + "codebaseIndexWatsonxProjectId", ] as const // Global secrets that are part of GlobalSettings (not ProviderSettings) diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 6d628ddfdf..de0ef1d54e 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -68,6 +68,7 @@ export const providerNames = [ "io-intelligence", "roo", "vercel-ai-gateway", + "watsonx", ] as const export const providerNamesSchema = z.enum(providerNames) @@ -343,6 +344,13 @@ const vercelAiGatewaySchema = baseProviderSettingsSchema.extend({ vercelAiGatewayModelId: z.string().optional(), }) +const watsonxSchema = baseProviderSettingsSchema.extend({ + watsonxBaseUrl: z.string().optional(), + watsonxApiKey: z.string().optional(), + watsonxProjectId: z.string().optional(), + watsonxModelId: z.string().optional(), +}) + const defaultSchema = z.object({ apiProvider: z.undefined(), }) @@ -384,6 +392,7 @@ export const providerSettingsSchemaDiscriminated = z.discriminatedUnion("apiProv qwenCodeSchema.merge(z.object({ apiProvider: z.literal("qwen-code") })), rooSchema.merge(z.object({ apiProvider: z.literal("roo") })), vercelAiGatewaySchema.merge(z.object({ apiProvider: z.literal("vercel-ai-gateway") })), + watsonxSchema.merge(z.object({ apiProvider: z.literal("watsonx") })), defaultSchema, ]) @@ -426,6 +435,7 @@ export const providerSettingsSchema = z.object({ ...rooSchema.shape, ...vercelAiGatewaySchema.shape, ...codebaseIndexProviderSchema.shape, + ...watsonxSchema.shape, }) export type ProviderSettings = z.infer diff --git a/packages/types/src/providers/index.ts b/packages/types/src/providers/index.ts index 21e43aaa99..c9b797757c 100644 --- a/packages/types/src/providers/index.ts +++ b/packages/types/src/providers/index.ts @@ -30,3 +30,4 @@ export * from "./xai.js" export * from "./vercel-ai-gateway.js" export * from "./zai.js" export * from "./deepinfra.js" +export * from "./watsonx.js" diff --git a/packages/types/src/providers/watsonx.ts b/packages/types/src/providers/watsonx.ts new file mode 100644 index 0000000000..1f78730128 --- /dev/null +++ b/packages/types/src/providers/watsonx.ts @@ -0,0 +1,97 @@ +import type { ModelInfo } from "../model.js" + +export type WatsonxAIModelId = keyof typeof watsonxAiModels +export const watsonxAiDefaultModelId: WatsonxAIModelId = "ibm/granite-3-3-8b-instruct" + +// Common model properties +const baseModelInfo: ModelInfo = { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: true, + supportsReasoningEffort: false, + supportsReasoningBudget: false, + requiredReasoningBudget: false, + inputPrice: 0, + outputPrice: 0, +} + +export const watsonxAiModels = { + // IBM Granite model + "ibm/granite-3-3-8b-instruct": { + ...baseModelInfo, + description: "Granite 3.3 8b Instruct - IBM-trained, dense decoder-only model", + }, + "ibm/granite-3-2-8b-instruct": { + ...baseModelInfo, + description: "Granite 3.2 8b Instruct - Text-only model capable of reasoning", + }, + "ibm/granite-3-2b-instruct": { + ...baseModelInfo, + description: "Granite 3 2b Instruct - IBM-trained, dense decoder-only model", + }, + "ibm/granite-3-8b-instruct": { + ...baseModelInfo, + description: "Granite 3 8b Instruct - IBM-trained, dense decoder-only model", + }, + "ibm/granite-guardian-3-2b": { + ...baseModelInfo, + description: "Granite Guardian 3 2b - IBM-trained, dense decoder-only model", + }, + "ibm/granite-guardian-3-8b": { + ...baseModelInfo, + description: "Granite Guardian 3 8b - IBM-trained, dense decoder-only model", + }, + "ibm/granite-vision-3-2-2b": { + ...baseModelInfo, + supportsImages: true, + description: "Granite 3 Vision - Image-text, text-out model capable of understanding images", + }, + // Meta Llama models + "meta-llama/llama-3-2-11b-vision-instruct": { + ...baseModelInfo, + supportsImages: true, + description: "Llama 3 2 11b Vision Instruct - Auto-regressive language model with transformer architecture", + }, + "meta-llama/llama-3-2-1b-instruct": { + ...baseModelInfo, + description: "Llama 3 2 1b Instruct - Auto-regressive language model with transformer architecture", + }, + "meta-llama/llama-3-2-3b-instruct": { + ...baseModelInfo, + description: "Llama 3 2 3b Instruct - Auto-regressive language model with transformer architecture", + }, + "meta-llama/llama-3-2-90b-vision-instruct": { + ...baseModelInfo, + supportsImages: true, + description: "Llama 3 2 90b Vision Instruct - Auto-regressive language model with transformer architecture", + }, + "meta-llama/llama-3-3-70b-instruct": { + ...baseModelInfo, + description: "Llama 3 3 70b Instruct - FP8 quantized version of the original FP16 weights", + }, + "meta-llama/llama-3-405b-instruct": { + ...baseModelInfo, + contextWindow: 128000, + description: "Llama 3 405b Instruct - Meta's largest open-source foundation model with 405 billion parameters", + }, + "meta-llama/llama-4-maverick-17b-1-0": { + ...baseModelInfo, + contextWindow: 128000, + description: "Llama 4 Maverick - 17 billion active parameter model with 128 experts", + }, + "meta-llama/llama-guard-3-11b-vision": { + ...baseModelInfo, + supportsImages: true, + description: "Llama Guard 3 11b Vision - Auto-regressive language model with transformer architecture", + }, + // Mistral AI models + "mistralai/mistral-medium-2505": { + ...baseModelInfo, + description: "Mistral Medium - Latest iteration of the Mistral Medium model family", + }, + "mistralai/mistral-small-3-1-24b-instruct-2503": { + ...baseModelInfo, + description: "Mistral Small 3.1 24B Base 2503 - Instruction-finetuned version of Mistral Small", + }, +} as const satisfies Record diff --git a/src/api/index.ts b/src/api/index.ts index ac00967676..1796b74d2a 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -40,6 +40,7 @@ import { FeatherlessHandler, VercelAiGatewayHandler, DeepInfraHandler, + WatsonxAIHandler, } from "./providers" import { NativeOllamaHandler } from "./providers/native-ollama" @@ -165,6 +166,8 @@ export function buildApiHandler(configuration: ProviderSettings): ApiHandler { return new FeatherlessHandler(options) case "vercel-ai-gateway": return new VercelAiGatewayHandler(options) + case "watsonx": + return new WatsonxAIHandler(options) default: apiProvider satisfies "gemini-cli" | undefined return new AnthropicHandler(options) diff --git a/src/api/providers/fetchers/watsonx.ts b/src/api/providers/fetchers/watsonx.ts new file mode 100644 index 0000000000..dcf049eb29 --- /dev/null +++ b/src/api/providers/fetchers/watsonx.ts @@ -0,0 +1,67 @@ +import { ModelInfo } from "@roo-code/types" +import { IamAuthenticator } from "ibm-cloud-sdk-core" +import { WatsonXAI } from "@ibm-cloud/watsonx-ai" + +/** + * Fetches available watsonx models + * + * @param apiKey - The watsonx API key + * @param projectId - Optional project ID for watsonx + * @param baseUrl - Optional base URL for the watsonx API + * @returns A promise resolving to an object with model IDs as keys and model info as values + */ +export async function getWatsonxModels( + apiKey: string, + projectId?: string, + baseUrl?: string, +): Promise> { + try { + const service = WatsonXAI.newInstance({ + version: "2024-05-31", + serviceUrl: baseUrl || "https://us-south.ml.cloud.ibm.com", + authenticator: new IamAuthenticator({ + apikey: apiKey, + }), + }) + + await service.getAuthenticator().authenticate() + let knownModels: Record = {} + + try { + const response = await service.listFoundationModelSpecs() + + 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 + const modelInfo = JSON.stringify(model).toLowerCase() + if ( + modelId && + !modelInfo.includes("embed") && + !modelInfo.includes("rtrvr") && + !modelInfo.includes("retriev") + ) { + const contextWindow = model.context_length || model.max_input_tokens || 8192 + const maxTokens = model.max_output_tokens || Math.floor(contextWindow / 2) + + knownModels[modelId] = { + contextWindow, + maxTokens, + supportsPromptCache: false, + } + } + } + } + } + } catch (apiError) { + console.warn("Error fetching models from IBM watsonx API:", apiError) + } + + return knownModels + } catch (error) { + console.error("Error fetching IBM watsonx models:", error) + return {} + } +} diff --git a/src/api/providers/index.ts b/src/api/providers/index.ts index 85d877b6bc..cc4627837e 100644 --- a/src/api/providers/index.ts +++ b/src/api/providers/index.ts @@ -34,3 +34,4 @@ export { RooHandler } from "./roo" export { FeatherlessHandler } from "./featherless" export { VercelAiGatewayHandler } from "./vercel-ai-gateway" export { DeepInfraHandler } from "./deepinfra" +export { WatsonxAIHandler } from "./watsonx" diff --git a/src/api/providers/watsonx.ts b/src/api/providers/watsonx.ts new file mode 100644 index 0000000000..94691da4b2 --- /dev/null +++ b/src/api/providers/watsonx.ts @@ -0,0 +1,165 @@ +import * as vscode from "vscode" +import { Anthropic } from "@anthropic-ai/sdk" +import { ModelInfo, watsonxAiDefaultModelId, watsonxAiModels, WatsonxAIModelId } from "@roo-code/types" +import type { ApiHandlerOptions } from "../../shared/api" +import { IamAuthenticator } from "ibm-cloud-sdk-core" +import { ApiStream } from "../transform/stream" +import { BaseProvider } from "./base-provider" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import { WatsonXAI } from "@ibm-cloud/watsonx-ai" +import { convertToWatsonxAiMessages } from "../transform/watsonxai-format" + +export class WatsonxAIHandler extends BaseProvider implements SingleCompletionHandler { + private options: ApiHandlerOptions + private projectId?: string + private service: WatsonXAI + + constructor(options: ApiHandlerOptions) { + super() + this.options = options + this.projectId = (this.options as any).watsonxProjectId + if (!this.projectId) { + throw new Error("You must provide a valid IBM watsonx project ID.") + } + const apiKey = (this.options as any).watsonxApiKey + if (!apiKey) { + throw new Error("You must provide a valid IBM watsonx API key.") + } + const serviceUrl = (this.options as any).watsonxBaseUrl || "https://us-south.ml.cloud.ibm.com" + + try { + const serviceOptions: any = { + version: "2024-05-31", + serviceUrl: serviceUrl, + authenticator: new IamAuthenticator({ + apikey: apiKey, + }), + } + this.service = WatsonXAI.newInstance(serviceOptions) + + this.service.getAuthenticator().authenticate() + } catch (error) { + throw new Error( + `IBM watsonx Authentication Error: ${error instanceof Error ? error.message : String(error)}`, + ) + } + } + + /** + * Creates parameters for WatsonX text chat API + * + * @param projectId - The IBM watsonx project ID + * @param modelId - The model ID to use + * @param messages - The messages to send + * @returns The parameters object for the API call + */ + private createTextChatParams(projectId: string, modelId: string, messages: any[]) { + const maxTokens = this.options.modelMaxTokens || 2048 + const temperature = this.options.modelTemperature || 0.7 + return { + projectId, + modelId, + messages, + maxTokens, + temperature, + } + } + + /** + * Creates a message using the IBM watsonx API directly + * + * @param systemPrompt - The system prompt to use + * @param messages - The conversation messages + * @param metadata - Optional metadata for the request + * @returns An async generator that yields the response + */ + async *createMessage( + systemPrompt: string, + messages: Anthropic.Messages.MessageParam[], + metadata?: ApiHandlerCreateMessageMetadata, + ): ApiStream { + const { id: modelId } = this.getModel() + + try { + // Convert messages to WatsonX format with system prompt + const watsonxMessages = [{ role: "system", content: systemPrompt }, ...convertToWatsonxAiMessages(messages)] + + const params = this.createTextChatParams(this.projectId!, modelId, watsonxMessages) + let responseText = "" + let usageInfo: any = null + + // Call the IBM watsonx API using textChat (non-streaming); can be changed to streaming.. + const response = await this.service.textChat(params) + + if (!response?.result?.choices?.[0]?.message?.content) { + throw new Error("Invalid or empty response from IBM watsonx API") + } + + responseText = response.result.choices[0].message.content + + yield { + type: "text", + text: responseText, + } + + usageInfo = response.result.usage || {} + const outputTokens = usageInfo.completion_tokens + + yield { + type: "usage", + inputTokens: usageInfo?.prompt_tokens, + outputTokens, + totalCost: 0, // Actual cost calculation could be added if available + } + } catch (error) { + await vscode.window.showErrorMessage(error.message) + yield { + type: "error", + error: error.type, + message: error.message, + } + } + } + + /** + * Completes a prompt using the IBM watsonx API directly with textChat + * + * @param prompt - The prompt to complete + * @returns The generated text + * @throws Error if the API call fails + */ + async completePrompt(prompt: string): Promise { + try { + const { id: modelId } = this.getModel() + const messages = [{ role: "user", content: prompt }] + const params = this.createTextChatParams(this.projectId!, modelId, messages) + const response = await this.service.textChat(params) + + if (!response?.result?.choices?.[0]?.message?.content) { + throw new Error("Invalid or empty response from IBM watsonx API") + } + + // Extract the message content directly + return response.result.choices[0].message.content + } catch (error) { + if (error instanceof Error) { + throw new Error(`IBM watsonx completion error: ${error.message}`) + } + throw new Error(`IBM watsonx completion error: ${error.message}`) + } + } + + /** + * Returns the model ID and model information for the current watsonx configuration + * + * @returns An object containing the model ID and model information + */ + override getModel(): { id: string; info: ModelInfo } { + return { + id: (this.options as any).watsonxModelId || watsonxAiDefaultModelId, + info: + watsonxAiModels[(this.options as any).watsonxModelId as WatsonxAIModelId] || + watsonxAiModels[watsonxAiDefaultModelId], + } + } +} diff --git a/src/api/transform/watsonxai-format.ts b/src/api/transform/watsonxai-format.ts new file mode 100644 index 0000000000..a81502a981 --- /dev/null +++ b/src/api/transform/watsonxai-format.ts @@ -0,0 +1,249 @@ +import { Anthropic } from "@anthropic-ai/sdk" +import OpenAI from "openai" + +/** + * Converts Anthropic message format to IBM watsonx.ai message format + * + * IBM watsonx.ai supports four message types: + * - TextChatMessageUser: Messages from the user + * - TextChatMessageAssistant: Messages from the assistant + * - TextChatMessageSystem: System instructions + * - TextChatMessageTool: Tool responses + * + * @param anthropicMessages - Messages in Anthropic format + * @returns Messages in IBM watsonx.ai format + */ +export function convertToWatsonxAiMessages( + anthropicMessages: Anthropic.Messages.MessageParam[], +): OpenAI.Chat.ChatCompletionMessageParam[] { + const watsonxAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [] + + for (const anthropicMessage of anthropicMessages) { + if ( + !anthropicMessage.content || + (Array.isArray(anthropicMessage.content) && anthropicMessage.content.length === 0) + ) { + continue + } + + switch (anthropicMessage.role) { + case "user": + // TextChatMessageUser + if (typeof anthropicMessage.content === "string") { + watsonxAiMessages.push({ + role: "user", + content: anthropicMessage.content, + }) + } else { + processUserMessage(anthropicMessage, watsonxAiMessages) + } + break + + case "assistant": + // TextChatMessageAssistant + if (typeof anthropicMessage.content === "string") { + watsonxAiMessages.push({ + role: "assistant", + content: anthropicMessage.content, + }) + } else { + processAssistantMessage(anthropicMessage, watsonxAiMessages) + } + break + + case "system" as any: + // TextChatMessageSystem + if (typeof anthropicMessage.content === "string") { + watsonxAiMessages.push({ + role: "system", + content: anthropicMessage.content, + }) + } else { + const textContent = anthropicMessage.content + .filter((block) => block.type === "text") + .map((block) => (block as any).text) + .join("\n") + + if (textContent) { + watsonxAiMessages.push({ + role: "system", + content: textContent, + }) + } + } + break + + default: + if (anthropicMessage.role === "tool") { + // TextChatMessageTool + const toolMessage = anthropicMessage as any + const toolCallId = toolMessage.tool_call_id + + if (typeof toolCallId === "string") { + const content = + typeof anthropicMessage.content === "string" + ? anthropicMessage.content + : anthropicMessage.content + .filter((block) => block.type === "text") + .map((block) => (block as any).text) + .join("\n") + + watsonxAiMessages.push({ + role: "tool", + tool_call_id: toolCallId, + content: content, + }) + } + } else if (typeof anthropicMessage.content === "string") { + watsonxAiMessages.push({ + role: anthropicMessage.role, + content: anthropicMessage.content, + }) + } + break + } + } + + return watsonxAiMessages +} + +function processUserMessage( + anthropicMessage: Anthropic.Messages.MessageParam, + watsonxAiMessages: OpenAI.Chat.ChatCompletionMessageParam[], +) { + const { contentBlocks, toolResultBlocks } = categorizeUserContent(anthropicMessage.content as any[]) + processToolResultBlocks(toolResultBlocks, watsonxAiMessages) + + if (contentBlocks.length > 0) { + const textBlocks = contentBlocks.filter((part) => part.type === "text") + + if (textBlocks.length === 1 && contentBlocks.length === 1) { + watsonxAiMessages.push({ + role: "user", + content: textBlocks[0].text, + }) + } else { + watsonxAiMessages.push({ + role: "user", + content: contentBlocks.map((part) => { + if (part.type === "image") { + return { + type: "image_url", + image_url: { url: `data:${part.source.media_type};base64,${part.source.data}` }, + } + } + return { type: "text", text: part.text } + }), + }) + } + } +} + +function processAssistantMessage( + anthropicMessage: Anthropic.Messages.MessageParam, + watsonxAiMessages: OpenAI.Chat.ChatCompletionMessageParam[], +) { + const { contentBlocks, toolUseBlocks } = categorizeAssistantContent(anthropicMessage.content as any[]) + + let content: string | undefined + if (contentBlocks.length > 0) { + content = contentBlocks.map((part) => (part.type === "text" ? part.text : "")).join("\n") + } + + const toolCalls = convertToolUseBlocksToToolCalls(toolUseBlocks) + + if (content || toolCalls.length > 0) { + watsonxAiMessages.push({ + role: "assistant", + content: content || "", + tool_calls: toolCalls.length > 0 ? toolCalls : undefined, + }) + } +} + +function categorizeUserContent(content: any[]) { + return content.reduce<{ + contentBlocks: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam)[] + toolResultBlocks: Anthropic.ToolResultBlockParam[] + }>( + (acc, part) => { + if (part.type === "tool_result") { + acc.toolResultBlocks.push(part) + } else if (part.type === "text" || part.type === "image") { + acc.contentBlocks.push(part) + } + return acc + }, + { contentBlocks: [], toolResultBlocks: [] }, + ) +} + +function categorizeAssistantContent(content: any[]) { + return content.reduce<{ + contentBlocks: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam)[] + toolUseBlocks: Anthropic.ToolUseBlockParam[] + }>( + (acc, part) => { + if (part.type === "tool_use") { + acc.toolUseBlocks.push(part) + } else if (part.type === "text" || part.type === "image") { + acc.contentBlocks.push(part) + } + return acc + }, + { contentBlocks: [], toolUseBlocks: [] }, + ) +} + +/** + * Process tool result blocks into IBM watsonx.ai TextChatMessageTool format + * + * @param toolResultBlocks - Tool result blocks from Anthropic + * @param watsonxAiMessages - Array to add the formatted messages to + */ +function processToolResultBlocks( + toolResultBlocks: Anthropic.ToolResultBlockParam[], + watsonxAiMessages: OpenAI.Chat.ChatCompletionMessageParam[], +) { + toolResultBlocks.forEach((toolResult) => { + if (!toolResult.tool_use_id) { + return + } + + let content: string + if (typeof toolResult.content === "string") { + content = toolResult.content + } else { + content = + toolResult.content + ?.map((part) => { + if (part.type === "image") { + return "(see following user message for image)" + } + return part.text + }) + .join("\n") ?? "" + } + + if (content.trim()) { + watsonxAiMessages.push({ + role: "tool", + tool_call_id: toolResult.tool_use_id, + content: content, + }) + } + }) +} + +function convertToolUseBlocksToToolCalls( + toolUseBlocks: Anthropic.ToolUseBlockParam[], +): OpenAI.Chat.ChatCompletionMessageToolCall[] { + return toolUseBlocks.map((toolUse) => ({ + id: toolUse.id, + type: "function", + function: { + name: toolUse.name, + arguments: JSON.stringify(toolUse.input), + }, + })) +} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index abdfae29fa..30c9b9a697 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -59,6 +59,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" export const webviewMessageHandler = async ( provider: ClineProvider, @@ -943,6 +944,29 @@ export const webviewMessageHandler = async ( // TODO: Cache like we do for OpenRouter, etc? provider.postMessageToWebview({ type: "vsCodeLmModels", vsCodeLmModels }) break + case "requestWatsonxModels": + if (message?.values?.apiKey) { + try { + const watsonxModels = await getWatsonxModels(message.values.apiKey, message.values.projectId) + const formattedModels: Record = {} + Object.entries(watsonxModels).forEach(([modelId]) => { + formattedModels[modelId] = { + dimension: 1536, + } + }) + provider.postMessageToWebview({ + type: "watsonxModels", + watsonxModels: formattedModels, + }) + } catch (error) { + console.error("Failed to fetch watsonx models:", error) + provider.postMessageToWebview({ + type: "watsonxModels", + watsonxModels: {}, + }) + } + } + break case "requestHuggingFaceModels": try { const { getHuggingFaceModelsWithMetadata } = await import("../../api/providers/fetchers/huggingface") @@ -2428,6 +2452,19 @@ export const webviewMessageHandler = async ( ) } + if (settings.codebaseIndexWatsonxApiKey !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexWatsonxApiKey", + settings.codebaseIndexWatsonxApiKey, + ) + } + if (settings.codebaseIndexWatsonxProjectId !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexWatsonxProjectId", + settings.codebaseIndexWatsonxProjectId, + ) + } + // Send success response first - settings are saved regardless of validation await provider.postMessageToWebview({ type: "codeIndexSettingsSaved", @@ -2564,6 +2601,7 @@ export const webviewMessageHandler = async ( const hasVercelAiGatewayApiKey = !!(await provider.context.secrets.get( "codebaseIndexVercelAiGatewayApiKey", )) + const hasWatsonxApiKey = !!(await provider.context.secrets.get("codebaseIndexWatsonxApiKey")) provider.postMessageToWebview({ type: "codeIndexSecretStatus", @@ -2574,6 +2612,7 @@ export const webviewMessageHandler = async ( hasGeminiApiKey, hasMistralApiKey, hasVercelAiGatewayApiKey, + hasWatsonxApiKey, }, }) break diff --git a/src/i18n/locales/en/embeddings.json b/src/i18n/locales/en/embeddings.json index fc902cadc1..b6b5fdebc8 100644 --- a/src/i18n/locales/en/embeddings.json +++ b/src/i18n/locales/en/embeddings.json @@ -48,6 +48,7 @@ "geminiConfigMissing": "Gemini configuration missing for embedder creation", "mistralConfigMissing": "Mistral configuration missing for embedder creation", "vercelAiGatewayConfigMissing": "Vercel AI Gateway configuration missing for embedder creation", + "watsonxConfigMissing": "IBM watsonx configuration missing for embedder creation", "invalidEmbedderType": "Invalid embedder type configured: {{embedderProvider}}", "vectorDimensionNotDeterminedOpenAiCompatible": "Could not determine vector dimension for model '{{modelId}}' with provider '{{provider}}'. Please ensure the 'Embedding Dimension' is correctly set in the OpenAI-Compatible provider settings.", "vectorDimensionNotDetermined": "Could not determine vector dimension for model '{{modelId}}' with provider '{{provider}}'. Check model profiles or configuration.", diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 2c0e8bb5c9..24c8577d97 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -20,6 +20,10 @@ export class CodeIndexConfigManager { private geminiOptions?: { apiKey: string } private mistralOptions?: { apiKey: string } private vercelAiGatewayOptions?: { apiKey: string } + private watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string private searchMinScore?: number @@ -71,6 +75,8 @@ export class CodeIndexConfigManager { const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? "" const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? "" const vercelAiGatewayApiKey = this.contextProxy?.getSecret("codebaseIndexVercelAiGatewayApiKey") ?? "" + const codebaseIndexWatsonxApiKey = this.contextProxy?.getSecret("codebaseIndexWatsonxApiKey") ?? "" + const codebaseIndexWatsonxProjectId = this.contextProxy?.getSecret("codebaseIndexWatsonxProjectId") ?? "" // Update instance variables with configuration this.codebaseIndexEnabled = codebaseIndexEnabled ?? true @@ -97,7 +103,6 @@ export class CodeIndexConfigManager { this.openAiOptions = { openAiNativeApiKey: openAiKey } - // Set embedder provider with support for openai-compatible if (codebaseIndexEmbedderProvider === "ollama") { this.embedderProvider = "ollama" } else if (codebaseIndexEmbedderProvider === "openai-compatible") { @@ -108,6 +113,8 @@ export class CodeIndexConfigManager { this.embedderProvider = "mistral" } else if (codebaseIndexEmbedderProvider === "vercel-ai-gateway") { this.embedderProvider = "vercel-ai-gateway" + } else if (codebaseIndexEmbedderProvider === "watsonx") { + this.embedderProvider = "watsonx" } else { this.embedderProvider = "openai" } @@ -129,6 +136,15 @@ export class CodeIndexConfigManager { this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined this.mistralOptions = mistralApiKey ? { apiKey: mistralApiKey } : undefined this.vercelAiGatewayOptions = vercelAiGatewayApiKey ? { apiKey: vercelAiGatewayApiKey } : undefined + if (codebaseIndexWatsonxApiKey) { + this.watsonxOptions = { + codebaseIndexWatsonxApiKey: codebaseIndexWatsonxApiKey, + codebaseIndexWatsonxProjectId: codebaseIndexWatsonxProjectId, + } + this.contextProxy.storeSecret("codebaseIndexWatsonxProjectId", codebaseIndexWatsonxProjectId) + } else { + this.watsonxOptions = undefined + } } /** @@ -147,6 +163,10 @@ export class CodeIndexConfigManager { geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } vercelAiGatewayOptions?: { apiKey: string } + watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -167,6 +187,8 @@ export class CodeIndexConfigManager { geminiApiKey: this.geminiOptions?.apiKey ?? "", mistralApiKey: this.mistralOptions?.apiKey ?? "", vercelAiGatewayApiKey: this.vercelAiGatewayOptions?.apiKey ?? "", + codebaseIndexWatsonxApiKey: this.watsonxOptions?.codebaseIndexWatsonxApiKey ?? "", + codebaseIndexWatsonxProjectId: this.watsonxOptions?.codebaseIndexWatsonxProjectId ?? "", qdrantUrl: this.qdrantUrl ?? "", qdrantApiKey: this.qdrantApiKey ?? "", } @@ -192,6 +214,7 @@ export class CodeIndexConfigManager { geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, vercelAiGatewayOptions: this.vercelAiGatewayOptions, + watsonxOptions: this.watsonxOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, @@ -231,6 +254,8 @@ export class CodeIndexConfigManager { return isConfigured } else if (this.embedderProvider === "vercel-ai-gateway") { const apiKey = this.vercelAiGatewayOptions?.apiKey + } else if (this.embedderProvider === "watsonx") { + const apiKey = this.watsonxOptions?.codebaseIndexWatsonxApiKey const qdrantUrl = this.qdrantUrl const isConfigured = !!(apiKey && qdrantUrl) return isConfigured @@ -269,6 +294,8 @@ export class CodeIndexConfigManager { const prevGeminiApiKey = prev?.geminiApiKey ?? "" const prevMistralApiKey = prev?.mistralApiKey ?? "" const prevVercelAiGatewayApiKey = prev?.vercelAiGatewayApiKey ?? "" + const prevWatsonxApiKey = prev?.codebaseIndexWatsonxApiKey ?? "" + const prevWatsonxProjectId = prev?.codebaseIndexWatsonxProjectId ?? "" const prevQdrantUrl = prev?.qdrantUrl ?? "" const prevQdrantApiKey = prev?.qdrantApiKey ?? "" @@ -307,6 +334,8 @@ export class CodeIndexConfigManager { const currentGeminiApiKey = this.geminiOptions?.apiKey ?? "" const currentMistralApiKey = this.mistralOptions?.apiKey ?? "" const currentVercelAiGatewayApiKey = this.vercelAiGatewayOptions?.apiKey ?? "" + const currentWatsonxApiKey = this.watsonxOptions?.codebaseIndexWatsonxApiKey ?? "" + const currentWatsonxProjectId = this.watsonxOptions?.codebaseIndexWatsonxProjectId ?? "" const currentQdrantUrl = this.qdrantUrl ?? "" const currentQdrantApiKey = this.qdrantApiKey ?? "" @@ -337,6 +366,10 @@ export class CodeIndexConfigManager { return true } + if (prevWatsonxApiKey !== currentWatsonxApiKey || prevWatsonxProjectId !== currentWatsonxProjectId) { + return true + } + // Check for model dimension changes (generic for all providers) if (prevModelDimension !== currentModelDimension) { return true @@ -395,6 +428,7 @@ export class CodeIndexConfigManager { geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, vercelAiGatewayOptions: this.vercelAiGatewayOptions, + watsonxOptions: this.watsonxOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, diff --git a/src/services/code-index/embedders/watsonx.ts b/src/services/code-index/embedders/watsonx.ts new file mode 100644 index 0000000000..a7f6deb108 --- /dev/null +++ b/src/services/code-index/embedders/watsonx.ts @@ -0,0 +1,283 @@ +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 } 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: 1536) + */ +export class WatsonxEmbedder implements IEmbedder { + private readonly watsonxClient: WatsonXAI + private static readonly WATSONX_VERSION = "2024-05-31" + private static readonly WATSONX_REGION = "us-south" + private static readonly DEFAULT_MODEL = "ibm/slate-125m-english-rtrvr-v2" + private readonly modelId: string + private readonly projectId?: string + + /** + * Creates a new watsonx embedder + * @param apiKey The watsonx API key for authentication + * @param modelId The model ID to use (defaults to ibm/slate-125m-english-rtrvr-v2) + * @param projectId Optional IBM Cloud project ID for watsonx + * @param proxyUrl Optional proxy URL for connecting through MCP servers + */ + constructor(apiKey: string, modelId?: string, projectId?: string) { + if (!apiKey) { + throw new Error(t("embeddings:validation.apiKeyRequired")) + } + this.modelId = modelId || WatsonxEmbedder.DEFAULT_MODEL + this.projectId = projectId + + const options: any = { + version: WatsonxEmbedder.WATSONX_VERSION, + authenticator: new IamAuthenticator({ + apikey: apiKey, + }), + serviceUrl: `https://${WatsonxEmbedder.WATSONX_REGION}.ml.cloud.ibm.com`, + } + + this.watsonxClient = new WatsonXAI(options) + + try { + this.watsonxClient.getAuthenticator().authenticate() + } catch (error) { + console.error("WatsonX authentication failed:", error) + throw new Error(t("embeddings:validation.authenticationFailed")) + } + } + + /** + * 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 + */ + 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, + }, + }, + }) + + 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 + } + break + } else { + embeddings.push([]) + break + } + } 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 (lastError && embeddings.length < texts.indexOf(text) + 1) { + embeddings.push([]) + console.error(`Failed to embed text after ${MAX_RETRIES} attempts:`, lastError) + } + } + + return { + embeddings, + usage: { + promptTokens, + totalTokens, + }, + } + } 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 + } + } + + /** + * Validates the watsonx embedder configuration by testing the API key and connection + * @returns Promise resolving to validation result with success status and optional error message + */ + 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], + projectId: this.projectId, + parameters: { + truncate_input_tokens: MAX_ITEM_TOKENS, + return_options: { + input_text: true, + }, + }, + }) + + if (!response?.result?.results || response.result.results.length === 0) { + console.error("IBM watsonx validation failed: Invalid response format", response) + return { + valid: false, + 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 = "" + + if (error instanceof Error) { + errorDetails = error.message + if (error.message.includes("401") || error.message.includes("unauthorized")) { + errorMessage = "embeddings:validation.invalidApiKey" + } else if (error.message.includes("404") || error.message.includes("not found")) { + errorMessage = "embeddings:validation.endpointNotFound" + } else if (error.message.includes("timeout") || error.message.includes("ECONNREFUSED")) { + errorMessage = "embeddings:validation.connectionTimeout" + } else if (error.message.includes("project")) { + errorMessage = "embeddings:validation.invalidProjectId" + } else if (error.message.includes("model")) { + errorMessage = "embeddings:validation.invalidModelId" + } + } + + return { + valid: false, + error: `${errorMessage} (${errorDetails})`, + } + } + } + + /** + * Fetches available embedding models from the IBM watsonx API + * @returns Promise resolving to an object with model IDs as keys and model info as values + */ + async getAvailableModels(): Promise> { + try { + console.log("Fetching available IBM watsonx embedding models...") + + const knownModels: Record = { + "ibm/slate-125m-english-rtrvr-v2": { dimension: 1536 }, + } + + try { + const response = await this.watsonxClient.listFoundationModelSpecs() + + console.log( + "IBM watsonx API response structure:", + Object.keys(response || {}).join(", "), + Object.keys(response?.result || {}).join(", "), + ) + + 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 + 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 } + } + } + } + } + } 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) + return { + "ibm/slate-125m-english-rtrvr-v2": { dimension: 768 }, + } + } + } + + /** + * Returns information about this embedder + */ + get embedderInfo(): EmbedderInfo { + return { + name: "watsonx", + } + } +} diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index f168e26869..7e1c94a841 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -15,6 +15,10 @@ export interface CodeIndexConfig { geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } vercelAiGatewayOptions?: { apiKey: string } + watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -37,6 +41,8 @@ export type PreviousConfigSnapshot = { geminiApiKey?: string mistralApiKey?: string vercelAiGatewayApiKey?: string + codebaseIndexWatsonxApiKey?: string + codebaseIndexWatsonxProjectId?: string qdrantUrl?: string qdrantApiKey?: string } diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts index 1fcda3aca3..b9770fd553 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -28,7 +28,14 @@ export interface EmbeddingResponse { } } -export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" +export type AvailableEmbedders = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "watsonx" export interface EmbedderInfo { name: AvailableEmbedders diff --git a/src/services/code-index/interfaces/manager.ts b/src/services/code-index/interfaces/manager.ts index 527900f6d1..e199047cd1 100644 --- a/src/services/code-index/interfaces/manager.ts +++ b/src/services/code-index/interfaces/manager.ts @@ -70,7 +70,14 @@ export interface ICodeIndexManager { } export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error" -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" +export type EmbedderProvider = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "watsonx" export interface IndexProgressUpdate { systemStatus: IndexingState diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index 6d69e1f0b6..34206f700d 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -5,6 +5,7 @@ import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible" import { GeminiEmbedder } from "./embedders/gemini" import { MistralEmbedder } from "./embedders/mistral" import { VercelAiGatewayEmbedder } from "./embedders/vercel-ai-gateway" +import { WatsonxEmbedder } from "./embedders/watsonx" import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels" import { QdrantVectorStore } from "./vector-store/qdrant-client" import { codeParser, DirectoryScanner, FileWatcher } from "./processors" @@ -79,6 +80,15 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.vercelAiGatewayConfigMissing")) } return new VercelAiGatewayEmbedder(config.vercelAiGatewayOptions.apiKey, config.modelId) + } else if (provider === "watsonx") { + if (!config.watsonxOptions?.codebaseIndexWatsonxApiKey) { + throw new Error(t("embeddings:serviceFactory.watsonxConfigMissing")) + } + return new WatsonxEmbedder( + config.watsonxOptions.codebaseIndexWatsonxApiKey, + config.modelId, + config.watsonxOptions.codebaseIndexWatsonxProjectId, + ) } throw new Error( diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index aaddc520cb..164da9e016 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -79,6 +79,7 @@ export interface ExtensionMessage { | "ollamaModels" | "lmStudioModels" | "vsCodeLmModels" + | "watsonxModels" | "huggingFaceModels" | "vsCodeLmApiAvailable" | "updatePrompt" @@ -152,6 +153,7 @@ export interface ExtensionMessage { ollamaModels?: ModelRecord lmStudioModels?: ModelRecord vsCodeLmModels?: { vendor?: string; family?: string; version?: string; id?: string }[] + watsonxModels?: Record huggingFaceModels?: Array<{ id: string object: string diff --git a/src/shared/ProfileValidator.ts b/src/shared/ProfileValidator.ts index 78ff6ed9fe..91f359f30e 100644 --- a/src/shared/ProfileValidator.ts +++ b/src/shared/ProfileValidator.ts @@ -92,6 +92,8 @@ export class ProfileValidator { return profile.ioIntelligenceModelId case "deepinfra": return profile.deepInfraModelId + case "watsonx": + return profile.watsonxModelId case "human-relay": case "fake-ai": default: diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index 93d0b9bc45..94013c46f8 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -70,6 +70,7 @@ export interface WebviewMessage { | "requestOllamaModels" | "requestLmStudioModels" | "requestVsCodeLmModels" + | "requestWatsonxModels" | "requestHuggingFaceModels" | "openImage" | "saveImage" @@ -297,6 +298,8 @@ export interface WebviewMessage { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVercelAiGatewayApiKey?: string + codebaseIndexWatsonxApiKey?: string + codebaseIndexWatsonxProjectId?: string } } diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index 80c51a6b45..62ffdc9653 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -2,7 +2,14 @@ * Defines profiles for different embedding models, including their dimensions. */ -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" // Add other providers as needed +export type EmbedderProvider = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "watsonx" // Add other providers as needed export interface EmbeddingModelProfile { dimension: number @@ -70,6 +77,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { "mistral/codestral-embed": { dimension: 1536, scoreThreshold: 0.4 }, "mistral/mistral-embed": { dimension: 1024, scoreThreshold: 0.4 }, }, + watsonx: { + "ibm/slate-125m-english-rtrvr-v2": { dimension: 768, scoreThreshold: 0.4 }, + }, } /** @@ -163,6 +173,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string { case "vercel-ai-gateway": return "openai/text-embedding-3-large" + case "watsonx": + return "ibm/slate-125m-english-rtrvr-v2" + default: // Fallback for unknown providers console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`) diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 45bf4224a1..8dda13fdbc 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -73,6 +73,8 @@ interface LocalCodeIndexSettings { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVercelAiGatewayApiKey?: string + codebaseIndexWatsonxApiKey?: string + codebaseIndexWatsonxProjectId?: string } // Validation schema for codebase index settings @@ -149,6 +151,15 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => { .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), }) + case "watsonx": + return baseSchema.extend({ + codebaseIndexWatsonxApiKey: z.string().min(1, t("settings:codeIndex.validation.watsonxApiKeyRequired")), + codebaseIndexWatsonxProjectId: z.string().optional(), + codebaseIndexEmbedderModelId: z + .string() + .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), + }) + default: return baseSchema } @@ -194,6 +205,8 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVercelAiGatewayApiKey: "", + codebaseIndexWatsonxApiKey: "", + codebaseIndexWatsonxProjectId: "", }) // Initial settings state - stores the settings when popover opens @@ -229,6 +242,8 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVercelAiGatewayApiKey: "", + codebaseIndexWatsonxApiKey: "", + codebaseIndexWatsonxProjectId: "", } setInitialSettings(settings) setCurrentSettings(settings) @@ -258,11 +273,32 @@ export const CodeIndexPopover: React.FC = ({ return () => window.removeEventListener("message", handleMessage) }, [open]) + // Request WatsonX models when provider is selected and API key is available + useEffect(() => { + if ( + currentSettings.codebaseIndexEmbedderProvider === "watsonx" && + currentSettings.codebaseIndexWatsonxApiKey && + currentSettings.codebaseIndexWatsonxApiKey !== SECRET_PLACEHOLDER + ) { + vscode.postMessage({ + type: "requestWatsonxModels", + values: { + apiKey: currentSettings.codebaseIndexWatsonxApiKey, + projectId: currentSettings.codebaseIndexWatsonxProjectId, + }, + }) + } + }, [ + currentSettings.codebaseIndexEmbedderProvider, + currentSettings.codebaseIndexWatsonxApiKey, + currentSettings.codebaseIndexWatsonxProjectId, + ]) + // Use a ref to capture current settings for the save handler const currentSettingsRef = useRef(currentSettings) currentSettingsRef.current = currentSettings - // Listen for indexing status updates and save responses + // Listen for indexing status updates, save responses, and watsonx models useEffect(() => { const handleMessage = (event: MessageEvent) => { if (event.data.type === "indexingStatusUpdate") { @@ -297,6 +333,10 @@ export const CodeIndexPopover: React.FC = ({ setSaveStatus("idle") setSaveError(null) } + } else if (event.data.type === "watsonxModels" && event.data.watsonxModels) { + // Update the extension state context with the watsonx models + // The models will be automatically available through the codebaseIndexModels context + console.log("Received WatsonX models:", event.data.watsonxModels) } } @@ -342,6 +382,16 @@ export const CodeIndexPopover: React.FC = ({ prev.codebaseIndexVercelAiGatewayApiKey === SECRET_PLACEHOLDER ) { updated.codebaseIndexVercelAiGatewayApiKey = secretStatus.hasVercelAiGatewayApiKey + } + + if (!prev.codebaseIndexWatsonxApiKey || prev.codebaseIndexWatsonxApiKey === SECRET_PLACEHOLDER) { + updated.codebaseIndexWatsonxApiKey = secretStatus.hasWatsonxApiKey ? SECRET_PLACEHOLDER : "" + } + if ( + !prev.codebaseIndexWatsonxProjectId || + prev.codebaseIndexWatsonxProjectId === SECRET_PLACEHOLDER + ) { + updated.codebaseIndexWatsonxProjectId = secretStatus.hasWatsonxProjectId ? SECRET_PLACEHOLDER : "" } @@ -418,7 +468,8 @@ export const CodeIndexPopover: React.FC = ({ key === "codebaseIndexOpenAiCompatibleApiKey" || key === "codebaseIndexGeminiApiKey" || key === "codebaseIndexMistralApiKey" || - key === "codebaseIndexVercelAiGatewayApiKey" + key === "codebaseIndexVercelAiGatewayApiKey" || + key === "codebaseIndexWatsonxApiKey" ) { dataToValidate[key] = "placeholder-valid" } @@ -528,6 +579,12 @@ export const CodeIndexPopover: React.FC = ({ const transformStyleString = `translateX(-${100 - progressPercentage}%)` + // Helper function to safely access models for any provider + const getProviderModels = (provider: EmbedderProvider) => { + if (!codebaseIndexModels) return {} + return (codebaseIndexModels as any)[provider] || {} + } + const getAvailableModels = () => { if (!codebaseIndexModels) return [] @@ -669,6 +726,9 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.vercelAiGatewayProvider")} + + {t("settings:codeIndex.watsonxProvider")} + @@ -714,10 +774,10 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.selectModel")} {getAvailableModels().map((modelId) => { - const model = - codebaseIndexModels?.[ - currentSettings.codebaseIndexEmbedderProvider - ]?.[modelId] + const providerModels = getProviderModels( + currentSettings.codebaseIndexEmbedderProvider, + ) + const model = providerModels[modelId] return ( {modelId}{" "} @@ -971,10 +1031,10 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.selectModel")} {getAvailableModels().map((modelId) => { - const model = - codebaseIndexModels?.[ - currentSettings.codebaseIndexEmbedderProvider - ]?.[modelId] + const providerModels = getProviderModels( + currentSettings.codebaseIndexEmbedderProvider, + ) + const model = providerModels[modelId] return ( {modelId}{" "} @@ -1036,10 +1096,99 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.selectModel")} {getAvailableModels().map((modelId) => { - const model = - codebaseIndexModels?.[ - currentSettings.codebaseIndexEmbedderProvider - ]?.[modelId] + const providerModels = getProviderModels( + currentSettings.codebaseIndexEmbedderProvider, + ) + const model = providerModels[modelId] + return ( + + {modelId}{" "} + {model + ? t("settings:codeIndex.modelDimensions", { + dimension: model.dimension, + }) + : ""} + + ) + })} + + {formErrors.codebaseIndexEmbedderModelId && ( +

+ {formErrors.codebaseIndexEmbedderModelId} +

+ )} + + + )} + + {currentSettings.codebaseIndexEmbedderProvider === "watsonx" && ( + <> +
+ + + updateSetting("codebaseIndexWatsonxApiKey", e.target.value) + } + placeholder={t("settings:codeIndex.watsonxApiKeyPlaceholder")} + className={cn("w-full", { + "border-red-500": formErrors.watsonxApiKey, + })} + /> + {formErrors.watsonxApiKey && ( +

+ {formErrors.watsonxApiKey} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexWatsonxProjectId", e.target.value) + } + placeholder={ + t("settings:codeIndex.watsonxProjectIdPlaceholder") || + "Optional IBM Cloud project ID" + } + className={cn("w-full", { + "border-red-500": formErrors.watsonxProjectId, + })} + /> + {formErrors.watsonxProjectId && ( +

+ {formErrors.watsonxProjectId} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexEmbedderModelId", e.target.value) + } + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexEmbedderModelId, + })}> + + {t("settings:codeIndex.selectModel")} + + {getAvailableModels().map((modelId) => { + const providerModels = getProviderModels( + currentSettings.codebaseIndexEmbedderProvider, + ) + const model = providerModels[modelId] return ( {modelId}{" "} diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index 7f2ac4ed7a..9dcb6cb4e7 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -37,6 +37,7 @@ import { rooDefaultModelId, vercelAiGatewayDefaultModelId, deepInfraDefaultModelId, + watsonxAiDefaultModelId, } from "@roo-code/types" import { vscode } from "@src/utils/vscode" @@ -89,6 +90,7 @@ import { Unbound, Vertex, VSCodeLM, + WatsonxAI, XAI, ZAi, Fireworks, @@ -348,6 +350,7 @@ const ApiOptions = ({ openai: { field: "openAiModelId" }, ollama: { field: "ollamaModelId" }, lmstudio: { field: "lmStudioModelId" }, + watsonx: { field: "apiModelId", default: watsonxAiDefaultModelId }, } const config = PROVIDER_MODEL_CONFIG[value] @@ -642,6 +645,10 @@ const ApiOptions = ({ /> )} + {selectedProvider === "watsonx" && ( + + )} + {selectedProvider === "human-relay" && ( <>
diff --git a/webview-ui/src/components/settings/constants.ts b/webview-ui/src/components/settings/constants.ts index ae336730ff..f8e7716a0c 100644 --- a/webview-ui/src/components/settings/constants.ts +++ b/webview-ui/src/components/settings/constants.ts @@ -21,6 +21,7 @@ import { fireworksModels, rooModels, featherlessModels, + watsonxAiModels, } from "@roo-code/types" export const MODELS_BY_PROVIDER: Partial>> = { @@ -44,6 +45,7 @@ export const MODELS_BY_PROVIDER: Partial a.label.localeCompare(b.label)) diff --git a/webview-ui/src/components/settings/providers/WatsonxAI.tsx b/webview-ui/src/components/settings/providers/WatsonxAI.tsx new file mode 100644 index 0000000000..5c44223c77 --- /dev/null +++ b/webview-ui/src/components/settings/providers/WatsonxAI.tsx @@ -0,0 +1,86 @@ +import { useCallback } from "react" +import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react" + +import type { ProviderSettings } from "@roo-code/types" +import { watsonxAiDefaultModelId, watsonxAiModels } from "@roo-code/types" + +import { useAppTranslation } from "@src/i18n/TranslationContext" +import { VSCodeButtonLink } from "@src/components/common/VSCodeButtonLink" + +import { inputEventTransform } from "../transforms" + +type WatsonxAIProps = { + apiConfiguration: ProviderSettings + setApiConfigurationField: (field: K, value: ProviderSettings[K]) => void +} + +export const WatsonxAI = ({ apiConfiguration, setApiConfigurationField }: WatsonxAIProps) => { + const { t } = useAppTranslation() + + const handleInputChange = useCallback( + (field: keyof ProviderSettings, transform: (event: E) => any = inputEventTransform) => + (event: E | Event) => { + setApiConfigurationField(field, transform(event as E)) + }, + [setApiConfigurationField], + ) + + const defaultModel = watsonxAiDefaultModelId + const modelInfo = watsonxAiModels[defaultModel] || {} + const defaultModelDescription = + typeof modelInfo === "object" && "contextWindow" in modelInfo + ? `Context window: ${modelInfo.contextWindow} tokens` + : "IBM watsonx model" + + return ( + <> + + + +
+ {t("settings:providers.apiKeyStorageNotice")} +
+ {!apiConfiguration?.watsonxApiKey && ( + + Get WatsonX API Key + + )} + + + + +
+ Project ID is required for IBM watsonx integration +
+ + + + +
+ Default: https://us-south.ml.cloud.ibm.com +

Default Model Information

+
+
+ Model ID: {defaultModel} +
+
+ Description: {defaultModelDescription} +
+
+
+ + ) +} diff --git a/webview-ui/src/components/settings/providers/index.ts b/webview-ui/src/components/settings/providers/index.ts index fe0e6cecf9..238c5636c9 100644 --- a/webview-ui/src/components/settings/providers/index.ts +++ b/webview-ui/src/components/settings/providers/index.ts @@ -30,3 +30,4 @@ export { Fireworks } from "./Fireworks" export { Featherless } from "./Featherless" export { VercelAiGateway } from "./VercelAiGateway" export { DeepInfra } from "./DeepInfra" +export { WatsonxAI } from "./WatsonxAI" diff --git a/webview-ui/src/components/ui/hooks/useSelectedModel.ts b/webview-ui/src/components/ui/hooks/useSelectedModel.ts index f8a005e86a..cacdb99538 100644 --- a/webview-ui/src/components/ui/hooks/useSelectedModel.ts +++ b/webview-ui/src/components/ui/hooks/useSelectedModel.ts @@ -57,6 +57,8 @@ import { vercelAiGatewayDefaultModelId, BEDROCK_CLAUDE_SONNET_4_MODEL_ID, deepInfraDefaultModelId, + watsonxAiModels, + watsonxAiDefaultModelId, } from "@roo-code/types" import type { ModelRecord, RouterModels } from "@roo/api" @@ -348,6 +350,14 @@ function getSelectedModel({ const info = routerModels["vercel-ai-gateway"]?.[id] return { id, info } } + case "watsonx": { + const id = apiConfiguration.apiModelId ?? watsonxAiDefaultModelId + const info = watsonxAiModels[id as keyof typeof watsonxAiModels] + return { + id, + info: info || undefined, + } + } // case "anthropic": // case "human-relay": // case "fake-ai": diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json index 1be824b37e..0463df26d6 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -62,6 +62,11 @@ "vercelAiGatewayProvider": "Vercel AI Gateway", "vercelAiGatewayApiKeyLabel": "API Key", "vercelAiGatewayApiKeyPlaceholder": "Enter your Vercel AI Gateway API key", + "watsonxProvider": "IBM watsonx", + "watsonxApiKeyLabel": "API Key", + "watsonxApiKeyPlaceholder": "Enter your IBM watsonx API key", + "watsonxProjectIdLabel": "Project ID", + "watsonxProjectIdPlaceholder": "Enter your IBM watsonx project ID", "openaiCompatibleProvider": "OpenAI Compatible", "openAiKeyLabel": "OpenAI API Key", "openAiKeyPlaceholder": "Enter your OpenAI API key", @@ -130,7 +135,8 @@ "vercelAiGatewayApiKeyRequired": "Vercel AI Gateway API key is required", "ollamaBaseUrlRequired": "Ollama base URL is required", "baseUrlRequired": "Base URL is required", - "modelDimensionMinValue": "Model dimension must be greater than 0" + "modelDimensionMinValue": "Model dimension must be greater than 0", + "watsonxApiKeyRequired": "IBM watsonx API key is required" } }, "autoApprove": { diff --git a/webview-ui/src/utils/validate.ts b/webview-ui/src/utils/validate.ts index 58cc8d38e8..9a2b1f862e 100644 --- a/webview-ui/src/utils/validate.ts +++ b/webview-ui/src/utils/validate.ts @@ -145,6 +145,14 @@ function validateModelsAndKeysProvided(apiConfiguration: ProviderSettings): stri return i18next.t("settings:validation.apiKey") } break + case "watsonx": + if (!apiConfiguration.watsonxApiKey) { + return i18next.t("settings:validation.apiKey") + } + if (!apiConfiguration.watsonxProjectId) { + return i18next.t("settings:validation.projectId") + } + break } return undefined