add refresh models button to get the models dynamically

This commit is contained in:
Prasang Prajapati 2025-08-26 15:16:15 -04:00
parent c821b5e13f
commit 24697d18d4
8 changed files with 220 additions and 44 deletions

View file

@ -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 = [

View file

@ -24,7 +24,6 @@ export async function getWatsonxModels(
}),
})
await service.getAuthenticator().authenticate()
let knownModels: Record<string, ModelInfo> = {}
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)

View file

@ -766,6 +766,7 @@ export const webviewMessageHandler = async (
ollama: {},
lmstudio: {},
deepinfra: {},
watsonx: {},
}
const safeGetModels = async (options: GetModelsOptions): Promise<ModelRecord> => {
@ -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)

View file

@ -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" && (
<WatsonxAI apiConfiguration={apiConfiguration} setApiConfigurationField={setApiConfigurationField} />
<WatsonxAI
apiConfiguration={apiConfiguration}
setApiConfigurationField={setApiConfigurationField}
organizationAllowList={organizationAllowList}
modelValidationError={modelValidationError}
/>
)}
{selectedProvider === "human-relay" && (
@ -687,7 +697,7 @@ const ApiOptions = ({
<Featherless apiConfiguration={apiConfiguration} setApiConfigurationField={setApiConfigurationField} />
)}
{selectedProviderModels.length > 0 && (
{selectedProviderModels.length > 0 && selectedProvider !== "watsonx" && (
<>
<div>
<label className="block font-medium mb-1">{t("settings:providers.model")}</label>

View file

@ -37,6 +37,7 @@ type ModelIdKey = keyof Pick<
| "deepInfraModelId"
| "ioIntelligenceModelId"
| "vercelAiGatewayModelId"
| "watsonxModelId"
>
interface ModelPickerProps {

View file

@ -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: <K extends keyof ProviderSettings>(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<string | undefined>()
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<ExtensionMessage>) => {
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(
<E,>(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
</div>
<VSCodeTextField
value={apiConfiguration?.watsonxBaseUrl || ""}
onInput={handleInputChange("watsonxBaseUrl")}
placeholder="https://us-south.ml.cloud.ibm.com"
className="w-full mt-4">
<label className="block font-medium mb-1">IBM watsonx API Base URL (Optional)</label>
</VSCodeTextField>
<div className="text-sm text-vscode-descriptionForeground mt-1 mb-3">
Default: https://us-south.ml.cloud.ibm.com
<h3 className="font-large">Default Model Information</h3>
<div className="text-sm">
<div>
<strong>Model ID:</strong> {defaultModel}
</div>
<div>
<strong>Description:</strong> {defaultModelDescription}
</div>
<div className="w-full mt-4">
<label className="block font-medium mb-1">IBM watsonx Region</label>
<Select value={selectedRegion} onValueChange={handleRegionSelect}>
<SelectTrigger className="w-full">
<SelectValue placeholder="Select a region" />
</SelectTrigger>
<SelectContent>
{Object.entries(WATSONX_REGIONS).map(([regionCode, regionName]) => (
<SelectItem key={regionCode} value={regionCode}>
{regionName}
</SelectItem>
))}
</SelectContent>
</Select>
<div className="text-sm text-vscode-descriptionForeground mt-1">
Selected endpoint: {REGION_TO_URL[selectedRegion as keyof typeof REGION_TO_URL]}
</div>
</div>
<Button
variant="outline"
onClick={handleRefreshModels}
disabled={refreshStatus === "loading" || !apiConfiguration.watsonxApiKey}
className="w-full mt-4">
<div className="flex items-center gap-2">
{refreshStatus === "loading" ? (
<span className="codicon codicon-loading codicon-modifier-spin" />
) : (
<span className="codicon codicon-refresh" />
)}
{t("settings:providers.refreshModels.label") || "Refresh Models"}
</div>
</Button>
{refreshStatus === "loading" && (
<div className="text-sm text-vscode-descriptionForeground">
{t("settings:providers.refreshModels.loading") || "Loading models..."}
</div>
)}
{refreshStatus === "success" && (
<div className="text-sm text-vscode-foreground">
{t("settings:providers.refreshModels.success") || "Models refreshed successfully"}
</div>
)}
{refreshStatus === "error" && (
<div className="text-sm text-vscode-errorForeground">
{refreshError || t("settings:providers.refreshModels.error") || "Failed to refresh models"}
</div>
)}
<ModelPicker
apiConfiguration={apiConfiguration}
defaultModelId={watsonxAiDefaultModelId}
models={routerModels?.watsonx ?? {}}
modelIdKey="watsonxModelId"
serviceName="IBM watsonx"
serviceUrl="https://cloud.ibm.com/apidocs/watsonx-ai#list-foundation-model-specs"
setApiConfigurationField={setApiConfigurationField}
organizationAllowList={organizationAllowList}
errorMessage={modelValidationError}
/>
</>
)
}

View file

@ -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)

View file

@ -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