diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index 89d5b168d7..e4f700b5b4 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -21,7 +21,9 @@ export const CODEBASE_INDEX_DEFAULTS = { export const codebaseIndexConfigSchema = z.object({ codebaseIndexEnabled: z.boolean().optional(), codebaseIndexQdrantUrl: z.string().optional(), - codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible", "gemini", "mistral"]).optional(), + codebaseIndexEmbedderProvider: z + .enum(["openai", "ollama", "openai-compatible", "gemini", "mistral", "watsonx"]) + .optional(), codebaseIndexEmbedderBaseUrl: z.string().optional(), codebaseIndexEmbedderModelId: z.string().optional(), codebaseIndexEmbedderModelDimension: z.number().optional(), @@ -48,6 +50,7 @@ export const codebaseIndexModelsSchema = z.object({ "openai-compatible": z.record(z.string(), z.object({ dimension: z.number() })).optional(), gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(), mistral: 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 @@ -64,6 +67,8 @@ export const codebaseIndexProviderSchema = z.object({ codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(), codebaseIndexGeminiApiKey: z.string().optional(), codebaseIndexMistralApiKey: 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 dc5a9e6744..39dfb7fd93 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -187,6 +187,9 @@ export const SECRET_STATE_KEYS = [ "codebaseIndexGeminiApiKey", "codebaseIndexMistralApiKey", "huggingFaceApiKey", + "watsonxApiKey", + "codebaseIndexWatsonxApiKey", + "codebaseIndexWatsonxProjectId", ] as const satisfies readonly (keyof ProviderSettings)[] export type SecretState = Pick diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 8cdb5296b2..4a5a431439 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -33,6 +33,7 @@ export const providerNames = [ "chutes", "litellm", "huggingface", + "watsonx", ] as const export const providerNamesSchema = z.enum(providerNames) @@ -241,6 +242,13 @@ const litellmSchema = baseProviderSettingsSchema.extend({ litellmUsePromptCache: z.boolean().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(), }) @@ -271,6 +279,7 @@ export const providerSettingsSchemaDiscriminated = z.discriminatedUnion("apiProv huggingFaceSchema.merge(z.object({ apiProvider: z.literal("huggingface") })), chutesSchema.merge(z.object({ apiProvider: z.literal("chutes") })), litellmSchema.merge(z.object({ apiProvider: z.literal("litellm") })), + watsonxSchema.merge(z.object({ apiProvider: z.literal("watsonx") })), defaultSchema, ]) @@ -302,6 +311,7 @@ export const providerSettingsSchema = z.object({ ...chutesSchema.shape, ...litellmSchema.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 f5061f152c..55827bbb68 100644 --- a/packages/types/src/providers/index.ts +++ b/packages/types/src/providers/index.ts @@ -19,3 +19,4 @@ export * from "./unbound.js" export * from "./vertex.js" export * from "./vscode-llm.js" export * from "./xai.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 bda390848c..a28cde55ea 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -30,6 +30,7 @@ import { ChutesHandler, LiteLLMHandler, ClaudeCodeHandler, + WatsonxAIHandler, } from "./providers" export interface SingleCompletionHandler { @@ -115,6 +116,8 @@ export function buildApiHandler(configuration: ProviderSettings): ApiHandler { return new ChutesHandler(options) case "litellm": return new LiteLLMHandler(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 1cefd0616b..e16a0be7f5 100644 --- a/src/api/providers/index.ts +++ b/src/api/providers/index.ts @@ -23,3 +23,4 @@ export { UnboundHandler } from "./unbound" export { VertexHandler } from "./vertex" export { VsCodeLmHandler } from "./vscode-lm" export { XAIHandler } from "./xai" +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 763e118125..bad54c365a 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -43,7 +43,6 @@ import { getVsCodeLmModels } from "../../api/providers/vscode-lm" import { openMention } from "../mentions" import { TelemetrySetting } from "../../shared/TelemetrySetting" import { getWorkspacePath } from "../../utils/path" -import { ensureSettingsDirectoryExists } from "../../utils/globalContext" import { Mode, defaultModeSlug } from "../../shared/modes" import { getModels, flushModels } from "../../api/providers/fetchers/modelCache" import { GetModelsOptions } from "../../shared/api" @@ -54,6 +53,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, @@ -674,6 +674,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") @@ -2036,6 +2059,18 @@ export const webviewMessageHandler = async ( settings.codebaseIndexMistralApiKey, ) } + 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({ @@ -2157,6 +2192,7 @@ export const webviewMessageHandler = async ( )) const hasGeminiApiKey = !!(await provider.context.secrets.get("codebaseIndexGeminiApiKey")) const hasMistralApiKey = !!(await provider.context.secrets.get("codebaseIndexMistralApiKey")) + const hasWatsonxApiKey = !!(await provider.context.secrets.get("codebaseIndexWatsonxApiKey")) provider.postMessageToWebview({ type: "codeIndexSecretStatus", @@ -2166,6 +2202,7 @@ export const webviewMessageHandler = async ( hasOpenAiCompatibleApiKey, hasGeminiApiKey, hasMistralApiKey, + hasWatsonxApiKey, }, }) break diff --git a/src/i18n/locales/en/embeddings.json b/src/i18n/locales/en/embeddings.json index 66465d8c35..84dd4fe97c 100644 --- a/src/i18n/locales/en/embeddings.json +++ b/src/i18n/locales/en/embeddings.json @@ -47,6 +47,7 @@ "openAiCompatibleConfigMissing": "OpenAI Compatible configuration missing for embedder creation", "geminiConfigMissing": "Gemini configuration missing for embedder creation", "mistralConfigMissing": "Mistral 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 1723f1c2a0..b0ca2d9e64 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -19,6 +19,10 @@ export class CodeIndexConfigManager { private openAiCompatibleOptions?: { baseUrl: string; apiKey: string } private geminiOptions?: { apiKey: string } private mistralOptions?: { apiKey: string } + private watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string private searchMinScore?: number @@ -69,6 +73,8 @@ export class CodeIndexConfigManager { const openAiCompatibleApiKey = this.contextProxy?.getSecret("codebaseIndexOpenAiCompatibleApiKey") ?? "" const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? "" const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? "" + const codebaseIndexWatsonxApiKey = this.contextProxy?.getSecret("codebaseIndexWatsonxApiKey") ?? "" + const codebaseIndexWatsonxProjectId = this.contextProxy?.getSecret("codebaseIndexWatsonxProjectId") ?? "" // Update instance variables with configuration this.codebaseIndexEnabled = codebaseIndexEnabled ?? true @@ -96,14 +102,17 @@ export class CodeIndexConfigManager { this.openAiOptions = { openAiNativeApiKey: openAiKey } // Set embedder provider with support for openai-compatible - if (codebaseIndexEmbedderProvider === "ollama") { + const provider = codebaseIndexEmbedderProvider as string + if (provider === "ollama") { this.embedderProvider = "ollama" - } else if (codebaseIndexEmbedderProvider === "openai-compatible") { + } else if (provider === "openai-compatible") { this.embedderProvider = "openai-compatible" - } else if (codebaseIndexEmbedderProvider === "gemini") { + } else if (provider === "gemini") { this.embedderProvider = "gemini" - } else if (codebaseIndexEmbedderProvider === "mistral") { + } else if (provider === "mistral") { this.embedderProvider = "mistral" + } else if (provider === "watsonx") { + this.embedderProvider = "watsonx" } else { this.embedderProvider = "openai" } @@ -124,6 +133,15 @@ export class CodeIndexConfigManager { this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined this.mistralOptions = mistralApiKey ? { apiKey: mistralApiKey } : undefined + if (codebaseIndexWatsonxApiKey) { + this.watsonxOptions = { + codebaseIndexWatsonxApiKey: codebaseIndexWatsonxApiKey, + codebaseIndexWatsonxProjectId: codebaseIndexWatsonxProjectId, + } + this.contextProxy.storeSecret("codebaseIndexWatsonxProjectId", codebaseIndexWatsonxProjectId) + } else { + this.watsonxOptions = undefined + } } /** @@ -141,6 +159,10 @@ export class CodeIndexConfigManager { openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } + watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -160,6 +182,8 @@ export class CodeIndexConfigManager { openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "", geminiApiKey: this.geminiOptions?.apiKey ?? "", mistralApiKey: this.mistralOptions?.apiKey ?? "", + codebaseIndexWatsonxApiKey: this.watsonxOptions?.codebaseIndexWatsonxApiKey ?? "", + codebaseIndexWatsonxProjectId: this.watsonxOptions?.codebaseIndexWatsonxProjectId ?? "", qdrantUrl: this.qdrantUrl ?? "", qdrantApiKey: this.qdrantApiKey ?? "", } @@ -184,6 +208,7 @@ export class CodeIndexConfigManager { openAiCompatibleOptions: this.openAiCompatibleOptions, geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, + watsonxOptions: this.watsonxOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, @@ -221,6 +246,11 @@ export class CodeIndexConfigManager { const qdrantUrl = this.qdrantUrl const isConfigured = !!(apiKey && qdrantUrl) return isConfigured + } else if (this.embedderProvider === "watsonx") { + const apiKey = this.watsonxOptions?.codebaseIndexWatsonxApiKey + const qdrantUrl = this.qdrantUrl + const isConfigured = !!(apiKey && qdrantUrl) + return isConfigured } return false // Should not happen if embedderProvider is always set correctly } @@ -255,6 +285,8 @@ export class CodeIndexConfigManager { const prevModelDimension = prev?.modelDimension const prevGeminiApiKey = prev?.geminiApiKey ?? "" const prevMistralApiKey = prev?.mistralApiKey ?? "" + const prevWatsonxApiKey = prev?.codebaseIndexWatsonxApiKey ?? "" + const prevWatsonxProjectId = prev?.codebaseIndexWatsonxProjectId ?? "" const prevQdrantUrl = prev?.qdrantUrl ?? "" const prevQdrantApiKey = prev?.qdrantApiKey ?? "" @@ -292,6 +324,8 @@ export class CodeIndexConfigManager { const currentModelDimension = this.modelDimension const currentGeminiApiKey = this.geminiOptions?.apiKey ?? "" const currentMistralApiKey = this.mistralOptions?.apiKey ?? "" + const currentWatsonxApiKey = this.watsonxOptions?.codebaseIndexWatsonxApiKey ?? "" + const currentWatsonxProjectId = this.watsonxOptions?.codebaseIndexWatsonxProjectId ?? "" const currentQdrantUrl = this.qdrantUrl ?? "" const currentQdrantApiKey = this.qdrantApiKey ?? "" @@ -318,6 +352,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 @@ -375,6 +413,7 @@ export class CodeIndexConfigManager { openAiCompatibleOptions: this.openAiCompatibleOptions, geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, + 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 9098a60091..18e31ceedf 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -14,6 +14,10 @@ export interface CodeIndexConfig { openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } + watsonxOptions?: { + codebaseIndexWatsonxApiKey: string + codebaseIndexWatsonxProjectId?: string + } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -35,6 +39,8 @@ export type PreviousConfigSnapshot = { openAiCompatibleApiKey?: string geminiApiKey?: string mistralApiKey?: 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 c5653ea2b7..48fbb6af83 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -28,7 +28,7 @@ export interface EmbeddingResponse { } } -export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" +export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "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 fd3b2bfdda..a442d99791 100644 --- a/src/services/code-index/interfaces/manager.ts +++ b/src/services/code-index/interfaces/manager.ts @@ -70,7 +70,7 @@ export interface ICodeIndexManager { } export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error" -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" +export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "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 68b0f5c0bc..bddc9d28d1 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -4,6 +4,7 @@ import { CodeIndexOllamaEmbedder } from "./embedders/ollama" import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible" import { GeminiEmbedder } from "./embedders/gemini" import { MistralEmbedder } from "./embedders/mistral" +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" @@ -70,6 +71,15 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.mistralConfigMissing")) } return new MistralEmbedder(config.mistralOptions.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 67f8782e19..76f69a2298 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -76,6 +76,7 @@ export interface ExtensionMessage { | "ollamaModels" | "lmStudioModels" | "vsCodeLmModels" + | "watsonxModels" | "huggingFaceModels" | "vsCodeLmApiAvailable" | "updatePrompt" @@ -148,6 +149,7 @@ export interface ExtensionMessage { ollamaModels?: string[] lmStudioModels?: string[] 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 9fc527c15a..79adf7eb78 100644 --- a/src/shared/ProfileValidator.ts +++ b/src/shared/ProfileValidator.ts @@ -84,6 +84,8 @@ export class ProfileValidator { return profile.ollamaModelId case "requesty": return profile.requestyModelId + 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 a91d1af7ba..db95977426 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -67,6 +67,7 @@ export interface WebviewMessage { | "requestOllamaModels" | "requestLmStudioModels" | "requestVsCodeLmModels" + | "requestWatsonxModels" | "requestHuggingFaceModels" | "openImage" | "saveImage" @@ -269,6 +270,8 @@ export interface WebviewMessage { codebaseIndexOpenAiCompatibleApiKey?: string codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string + codebaseIndexWatsonxApiKey?: string + codebaseIndexWatsonxProjectId?: string } } diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index a3cd61e659..b3172814c9 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -2,7 +2,7 @@ * Defines profiles for different embedding models, including their dimensions. */ -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" // Add other providers as needed +export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "watsonx" // Add other providers as needed export interface EmbeddingModelProfile { dimension: number @@ -53,6 +53,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { mistral: { "codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 }, }, + watsonx: { + "ibm/slate-125m-english-rtrvr-v2": { dimension: 768, scoreThreshold: 0.4 }, + }, } /** @@ -143,6 +146,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string { case "mistral": return "codestral-embed-2505" + 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 c85aaf6ea5..be39c0a173 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -70,6 +70,8 @@ interface LocalCodeIndexSettings { codebaseIndexOpenAiCompatibleApiKey?: string codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string + codebaseIndexWatsonxApiKey?: string + codebaseIndexWatsonxProjectId?: string } // Validation schema for codebase index settings @@ -136,6 +138,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 } @@ -180,6 +191,8 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexOpenAiCompatibleApiKey: "", codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", + codebaseIndexWatsonxApiKey: "", + codebaseIndexWatsonxProjectId: "", }) // Initial settings state - stores the settings when popover opens @@ -214,6 +227,8 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexOpenAiCompatibleApiKey: "", codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", + codebaseIndexWatsonxApiKey: "", + codebaseIndexWatsonxProjectId: "", } setInitialSettings(settings) setCurrentSettings(settings) @@ -231,11 +246,32 @@ export const CodeIndexPopover: React.FC = ({ } }, [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") { @@ -268,6 +304,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) } } @@ -308,6 +348,17 @@ export const CodeIndexPopover: React.FC = ({ if (!prev.codebaseIndexMistralApiKey || prev.codebaseIndexMistralApiKey === SECRET_PLACEHOLDER) { updated.codebaseIndexMistralApiKey = secretStatus.hasMistralApiKey ? SECRET_PLACEHOLDER : "" } + 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 + : "" + } return updated } @@ -380,7 +431,8 @@ export const CodeIndexPopover: React.FC = ({ key === "codeIndexOpenAiKey" || key === "codebaseIndexOpenAiCompatibleApiKey" || key === "codebaseIndexGeminiApiKey" || - key === "codebaseIndexMistralApiKey" + key === "codebaseIndexMistralApiKey" || + key === "codebaseIndexWatsonxApiKey" ) { dataToValidate[key] = "placeholder-valid" } @@ -490,6 +542,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 [] @@ -628,6 +686,9 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.mistralProvider")} + + {t("settings:codeIndex.watsonxProvider")} + @@ -673,10 +734,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}{" "} @@ -930,10 +991,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}{" "} @@ -995,10 +1056,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 977822cac3..e89d5072e7 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -25,6 +25,7 @@ import { chutesDefaultModelId, bedrockDefaultModelId, vertexDefaultModelId, + watsonxAiDefaultModelId, } from "@roo-code/types" import { vscode } from "@src/utils/vscode" @@ -72,6 +73,7 @@ import { Unbound, Vertex, VSCodeLM, + WatsonxAI, XAI, } from "./providers" @@ -300,6 +302,7 @@ const ApiOptions = ({ openai: { field: "openAiModelId" }, ollama: { field: "ollamaModelId" }, lmstudio: { field: "lmStudioModelId" }, + watsonx: { field: "apiModelId", default: watsonxAiDefaultModelId }, } const config = PROVIDER_MODEL_CONFIG[value] @@ -509,6 +512,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 995f591034..f64ade1b52 100644 --- a/webview-ui/src/components/settings/constants.ts +++ b/webview-ui/src/components/settings/constants.ts @@ -13,6 +13,7 @@ import { xaiModels, groqModels, chutesModels, + watsonxAiModels, } from "@roo-code/types" export const MODELS_BY_PROVIDER: Partial>> = { @@ -28,6 +29,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 6c6fdddaee..fcd1fa6c0c 100644 --- a/webview-ui/src/components/settings/providers/index.ts +++ b/webview-ui/src/components/settings/providers/index.ts @@ -20,3 +20,4 @@ export { Vertex } from "./Vertex" export { VSCodeLM } from "./VSCodeLM" export { XAI } from "./XAI" export { LiteLLM } from "./LiteLLM" +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 8dceb6e117..fe5f842831 100644 --- a/webview-ui/src/components/ui/hooks/useSelectedModel.ts +++ b/webview-ui/src/components/ui/hooks/useSelectedModel.ts @@ -34,6 +34,8 @@ import { litellmDefaultModelId, claudeCodeDefaultModelId, claudeCodeModels, + watsonxAiModels, + watsonxAiDefaultModelId, } from "@roo-code/types" import type { RouterModels } from "@roo/api" @@ -224,6 +226,14 @@ function getSelectedModel({ const info = claudeCodeModels[id as keyof typeof claudeCodeModels] return { id, info: { ...openAiModelInfoSaneDefaults, ...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 7c58e679c6..6f31b5e3a8 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -55,6 +55,11 @@ "mistralProvider": "Mistral", "mistralApiKeyLabel": "API Key:", "mistralApiKeyPlaceholder": "Enter your Mistral 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", @@ -122,7 +127,8 @@ "mistralApiKeyRequired": "Mistral 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 ed546cccc7..785a149adb 100644 --- a/webview-ui/src/utils/validate.ts +++ b/webview-ui/src/utils/validate.ts @@ -110,6 +110,14 @@ function validateModelsAndKeysProvided(apiConfiguration: ProviderSettings): stri return i18next.t("settings:validation.modelId") } break + case "watsonx": + if (!apiConfiguration.watsonxApiKey) { + return i18next.t("settings:validation.apiKey") + } + if (!apiConfiguration.watsonxProjectId) { + return i18next.t("settings:validation.projectId") + } + break } return undefined