From 24697d18d4ab28f38a696241cf042320475eeb67 Mon Sep 17 00:00:00 2001 From: Prasang Prajapati Date: Tue, 26 Aug 2025 15:16:15 -0400 Subject: [PATCH] add refresh models button to get the models dynamically --- packages/types/src/provider-settings.ts | 7 +- src/api/providers/fetchers/watsonx.ts | 8 +- src/core/webview/webviewMessageHandler.ts | 19 ++ .../src/components/settings/ApiOptions.tsx | 16 +- .../src/components/settings/ModelPicker.tsx | 1 + .../settings/providers/WatsonxAI.tsx | 205 +++++++++++++++--- .../components/ui/hooks/useSelectedModel.ts | 3 + webview-ui/src/utils/validate.ts | 5 + 8 files changed, 220 insertions(+), 44 deletions(-) diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 8dbdd6f010..f0e5372c7b 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -25,7 +25,6 @@ import { vscodeLlmModels, xaiModels, internationalZAiModels, - watsonxAiModels, } from "./providers/index.js" /** @@ -578,11 +577,6 @@ export const MODELS_BY_PROVIDER: Record< label: "VS Code LM API", models: Object.keys(vscodeLlmModels), }, - watsonx: { - id: "watsonx", - label: "IBM watsonx", - models: Object.keys(watsonxAiModels), - }, xai: { id: "xai", label: "xAI (Grok)", models: Object.keys(xaiModels) }, zai: { id: "zai", label: "Zai", models: Object.keys(internationalZAiModels) }, @@ -595,6 +589,7 @@ export const MODELS_BY_PROVIDER: Record< unbound: { id: "unbound", label: "Unbound", models: [] }, deepinfra: { id: "deepinfra", label: "DeepInfra", models: [] }, "vercel-ai-gateway": { id: "vercel-ai-gateway", label: "Vercel AI Gateway", models: [] }, + watsonx: { id: "watsonx", label: "IBM watsonx", models: [] }, } export const dynamicProviders = [ diff --git a/src/api/providers/fetchers/watsonx.ts b/src/api/providers/fetchers/watsonx.ts index dcf049eb29..98da1003a7 100644 --- a/src/api/providers/fetchers/watsonx.ts +++ b/src/api/providers/fetchers/watsonx.ts @@ -24,7 +24,6 @@ export async function getWatsonxModels( }), }) - await service.getAuthenticator().authenticate() let knownModels: Record = {} try { @@ -37,12 +36,7 @@ export async function getWatsonxModels( 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") - ) { + if (modelId && !modelInfo.includes("embed") && !modelInfo.includes("rtrvr")) { const contextWindow = model.context_length || model.max_input_tokens || 8192 const maxTokens = model.max_output_tokens || Math.floor(contextWindow / 2) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 30c9b9a697..8aa90ab4b6 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -766,6 +766,7 @@ export const webviewMessageHandler = async ( ollama: {}, lmstudio: {}, deepinfra: {}, + watsonx: {}, } const safeGetModels = async (options: GetModelsOptions): Promise => { @@ -791,6 +792,14 @@ export const webviewMessageHandler = async ( }, }, { key: "glama", options: { provider: "glama" } }, + { + key: "watsonx", + options: { + provider: "watsonx", + apiKey: apiConfiguration.watsonxApiKey!, + baseUrl: apiConfiguration.watsonxBaseUrl!, + }, + }, { key: "unbound", options: { provider: "unbound", apiKey: apiConfiguration.unboundApiKey } }, { key: "vercel-ai-gateway", options: { provider: "vercel-ai-gateway" } }, { @@ -824,6 +833,16 @@ export const webviewMessageHandler = async ( }) } + const watsonxApiKey = apiConfiguration.watsonxApiKey + const watsonxBaseUrl = apiConfiguration.watsonxBaseUrl + + if (watsonxApiKey && watsonxBaseUrl) { + modelFetchPromises.push({ + key: "watsonx", + options: { provider: "watsonx", apiKey: watsonxApiKey, baseUrl: watsonxBaseUrl }, + }) + } + const results = await Promise.allSettled( modelFetchPromises.map(async ({ key, options }) => { const models = await safeGetModels(options) diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index 9dcb6cb4e7..578c3377e3 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -232,6 +232,8 @@ const ApiOptions = ({ vscode.postMessage({ type: "requestRouterModels" }) } else if (selectedProvider === "deepinfra") { vscode.postMessage({ type: "requestRouterModels" }) + } else if (selectedProvider === "watsonx") { + vscode.postMessage({ type: "requestWatsonxModels" }) } }, 250, @@ -246,6 +248,9 @@ const ApiOptions = ({ apiConfiguration?.litellmApiKey, apiConfiguration?.deepInfraApiKey, apiConfiguration?.deepInfraBaseUrl, + apiConfiguration.watsonxApiKey, + apiConfiguration.watsonxProjectId, + apiConfiguration.watsonxBaseUrl, customHeaders, ], ) @@ -350,7 +355,7 @@ const ApiOptions = ({ openai: { field: "openAiModelId" }, ollama: { field: "ollamaModelId" }, lmstudio: { field: "lmStudioModelId" }, - watsonx: { field: "apiModelId", default: watsonxAiDefaultModelId }, + watsonx: { field: "watsonxModelId", default: watsonxAiDefaultModelId }, } const config = PROVIDER_MODEL_CONFIG[value] @@ -646,7 +651,12 @@ const ApiOptions = ({ )} {selectedProvider === "watsonx" && ( - + )} {selectedProvider === "human-relay" && ( @@ -687,7 +697,7 @@ const ApiOptions = ({ )} - {selectedProviderModels.length > 0 && ( + {selectedProviderModels.length > 0 && selectedProvider !== "watsonx" && ( <>
diff --git a/webview-ui/src/components/settings/ModelPicker.tsx b/webview-ui/src/components/settings/ModelPicker.tsx index 74e3d31f00..5f49ccd9fa 100644 --- a/webview-ui/src/components/settings/ModelPicker.tsx +++ b/webview-ui/src/components/settings/ModelPicker.tsx @@ -37,6 +37,7 @@ type ModelIdKey = keyof Pick< | "deepInfraModelId" | "ioIntelligenceModelId" | "vercelAiGatewayModelId" + | "watsonxModelId" > interface ModelPickerProps { diff --git a/webview-ui/src/components/settings/providers/WatsonxAI.tsx b/webview-ui/src/components/settings/providers/WatsonxAI.tsx index 5c44223c77..2258e698cc 100644 --- a/webview-ui/src/components/settings/providers/WatsonxAI.tsx +++ b/webview-ui/src/components/settings/providers/WatsonxAI.tsx @@ -1,21 +1,112 @@ -import { useCallback } from "react" +import { useCallback, useState, useEffect, useRef } 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 { watsonxAiDefaultModelId, type ProviderSettings } from "@roo-code/types" import { useAppTranslation } from "@src/i18n/TranslationContext" import { VSCodeButtonLink } from "@src/components/common/VSCodeButtonLink" +import { vscode } from "@src/utils/vscode" +import { Button, Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@src/components/ui" +import { ExtensionMessage } from "@roo/ExtensionMessage" import { inputEventTransform } from "../transforms" +import { OrganizationAllowList } from "@roo/cloud" +import { useExtensionState } from "@src/context/ExtensionStateContext" +import { RouterName } from "@roo/api" +import { ModelPicker } from "../ModelPicker" + +// Define the available regions +const WATSONX_REGIONS = { + "us-south": "Dallas (us-south.ml.cloud.ibm.com)", + "eu-de": "Frankfurt (eu-de.ml.cloud.ibm.com)", + "eu-gb": "London (eu-gb.ml.cloud.ibm.com)", + "jp-tok": "Tokyo (jp-tok.ml.cloud.ibm.com)", + "au-syd": "Sydney (au-syd.ml.cloud.ibm.com)", + "ca-tor": "Toronto (ca-tor.ml.cloud.ibm.com)", + "ap-south-1": "Mumbai (ap-south-1.aws.wxai.ibm.com)", +} + +// Map region codes to full URLs +const REGION_TO_URL = { + "us-south": "https://us-south.ml.cloud.ibm.com", + "eu-de": "https://eu-de.ml.cloud.ibm.com", + "eu-gb": "https://eu-gb.ml.cloud.ibm.com", + "jp-tok": "https://jp-tok.ml.cloud.ibm.com", + "au-syd": "https://au-syd.ml.cloud.ibm.com", + "ca-tor": "https://ca-tor.ml.cloud.ibm.com", + "ap-south-1": "https://ap-south-1.aws.wxai.ibm.com", + custom: "", // For custom URL input +} type WatsonxAIProps = { apiConfiguration: ProviderSettings - setApiConfigurationField: (field: K, value: ProviderSettings[K]) => void + setApiConfigurationField: (field: keyof ProviderSettings, value: ProviderSettings[keyof ProviderSettings]) => void + organizationAllowList: OrganizationAllowList + modelValidationError?: string } -export const WatsonxAI = ({ apiConfiguration, setApiConfigurationField }: WatsonxAIProps) => { +export const WatsonxAI = ({ + apiConfiguration, + setApiConfigurationField, + organizationAllowList, + modelValidationError, +}: WatsonxAIProps) => { const { t } = useAppTranslation() + const { routerModels } = useExtensionState() + const [refreshStatus, setRefreshStatus] = useState<"idle" | "loading" | "success" | "error">("idle") + const [refreshError, setRefreshError] = useState() + const watsonxErrorJustReceived = useRef(false) + + // Determine the current region based on the base URL + const getCurrentRegion = () => { + const baseUrl = apiConfiguration?.watsonxBaseUrl || "" + + // Find the region that matches the current base URL + const regionEntry = Object.entries(REGION_TO_URL).find(([_, url]) => url === baseUrl) + + // Return the region code or 'us-south' as default if not found + return regionEntry ? regionEntry[0] : "us-south" + } + + const [selectedRegion, setSelectedRegion] = useState(getCurrentRegion()) + + // Handle region selection + const handleRegionSelect = useCallback( + (region: string) => { + setSelectedRegion(region) + + // Update the base URL in the API configuration + const baseUrl = REGION_TO_URL[region as keyof typeof REGION_TO_URL] || "" + setApiConfigurationField("watsonxBaseUrl", baseUrl) + }, + [setApiConfigurationField], + ) + + useEffect(() => { + const handleMessage = (event: MessageEvent) => { + const message = event.data + if (message.type === "singleRouterModelFetchResponse" && !message.success) { + const providerName = message.values?.provider as RouterName + if (providerName === "watsonx") { + watsonxErrorJustReceived.current = true + setRefreshStatus("error") + setRefreshError(message.error) + } + } else if (message.type === "routerModels") { + // When router models are updated, update the refresh status + if (refreshStatus === "loading") { + if (!watsonxErrorJustReceived.current) { + setRefreshStatus("success") + } + } + } + } + + window.addEventListener("message", handleMessage) + return () => { + window.removeEventListener("message", handleMessage) + } + }, [refreshStatus, refreshError, t]) const handleInputChange = useCallback( (field: keyof ProviderSettings, transform: (event: E) => any = inputEventTransform) => @@ -25,12 +116,29 @@ export const WatsonxAI = ({ apiConfiguration, setApiConfigurationField }: Watson [setApiConfigurationField], ) - const defaultModel = watsonxAiDefaultModelId - const modelInfo = watsonxAiModels[defaultModel] || {} - const defaultModelDescription = - typeof modelInfo === "object" && "contextWindow" in modelInfo - ? `Context window: ${modelInfo.contextWindow} tokens` - : "IBM watsonx model" + const handleRefreshModels = useCallback(() => { + setRefreshStatus("loading") + setRefreshError(undefined) + + const apiKey = apiConfiguration.watsonxApiKey + const projectId = apiConfiguration.watsonxProjectId + const baseUrl = REGION_TO_URL[selectedRegion as keyof typeof REGION_TO_URL] + + if (!apiKey) { + setRefreshStatus("error") + setRefreshError(t("settings:providers.refreshModels.missingConfig")) + return + } + + vscode.postMessage({ + type: "requestRouterModels", + values: { + watsonxApiKey: apiKey, + watsonxProjectId: projectId, + watsonxBaseUrl: baseUrl, + }, + }) + }, [apiConfiguration, setRefreshStatus, setRefreshError, t, selectedRegion]) return ( <> @@ -62,25 +170,66 @@ export const WatsonxAI = ({ apiConfiguration, setApiConfigurationField }: Watson Project ID is required for IBM watsonx integration
- - - -
- Default: https://us-south.ml.cloud.ibm.com -

Default Model Information

-
-
- Model ID: {defaultModel} -
-
- Description: {defaultModelDescription} -
+
+ + +
+ Selected endpoint: {REGION_TO_URL[selectedRegion as keyof typeof REGION_TO_URL]}
+ + + {refreshStatus === "loading" && ( +
+ {t("settings:providers.refreshModels.loading") || "Loading models..."} +
+ )} + {refreshStatus === "success" && ( +
+ {t("settings:providers.refreshModels.success") || "Models refreshed successfully"} +
+ )} + {refreshStatus === "error" && ( +
+ {refreshError || t("settings:providers.refreshModels.error") || "Failed to refresh models"} +
+ )} + + ) } diff --git a/webview-ui/src/components/ui/hooks/useSelectedModel.ts b/webview-ui/src/components/ui/hooks/useSelectedModel.ts index cacdb99538..7b1709852a 100644 --- a/webview-ui/src/components/ui/hooks/useSelectedModel.ts +++ b/webview-ui/src/components/ui/hooks/useSelectedModel.ts @@ -59,6 +59,8 @@ import { deepInfraDefaultModelId, watsonxAiModels, watsonxAiDefaultModelId, + watsonxAiDefaultModelId, + watsonxAiModels, } from "@roo-code/types" import type { ModelRecord, RouterModels } from "@roo/api" @@ -75,6 +77,7 @@ export const useSelectedModel = (apiConfiguration?: ProviderSettings) => { const ollamaModelId = provider === "ollama" ? apiConfiguration?.ollamaModelId : undefined const routerModels = useRouterModels() + const openRouterModelProviders = useOpenRouterModelProviders(openRouterModelId) const lmStudioModels = useLmStudioModels(lmStudioModelId) const ollamaModels = useOllamaModels(ollamaModelId) diff --git a/webview-ui/src/utils/validate.ts b/webview-ui/src/utils/validate.ts index 9a2b1f862e..1ce37c6ff1 100644 --- a/webview-ui/src/utils/validate.ts +++ b/webview-ui/src/utils/validate.ts @@ -225,6 +225,8 @@ function getModelIdForProvider(apiConfiguration: ProviderSettings, provider: str return apiConfiguration.ioIntelligenceModelId case "vercel-ai-gateway": return apiConfiguration.vercelAiGatewayModelId + case "watsonx": + return apiConfiguration.watsonxModelId default: return apiConfiguration.apiModelId } @@ -298,6 +300,9 @@ export function validateModelId(apiConfiguration: ProviderSettings, routerModels case "litellm": modelId = apiConfiguration.litellmModelId break + case "watsonx": + modelId = apiConfiguration.watsonxModelId + break case "io-intelligence": modelId = apiConfiguration.ioIntelligenceModelId break