diff --git a/src/api/providers/fetchers/litellm.ts b/src/api/providers/fetchers/litellm.ts index abe6ee1f59..a97d94fb7e 100644 --- a/src/api/providers/fetchers/litellm.ts +++ b/src/api/providers/fetchers/litellm.ts @@ -1,5 +1,5 @@ import axios from "axios" -import { ModelRecord } from "../../../shared/api" +import { COMPUTER_USE_MODELS, ModelRecord } from "../../../shared/api" /** * Fetches available models from a LiteLLM server @@ -23,6 +23,8 @@ export async function getLiteLLMModels(apiKey: string, baseUrl: string): Promise const response = await axios.get(`${baseUrl}/v1/model/info`, { headers, timeout: 15000 }) const models: ModelRecord = {} + const computerModels = Array.from(COMPUTER_USE_MODELS) + // Process the model info from the response if (response.data && response.data.data && Array.isArray(response.data.data)) { for (const model of response.data.data) { @@ -43,7 +45,9 @@ export async function getLiteLLMModels(apiKey: string, baseUrl: string): Promise maxTokens: determinedMaxTokens, contextWindow: modelInfo.max_input_tokens || 200000, supportsImages: Boolean(modelInfo.supports_vision), - supportsComputerUse: Boolean(modelInfo.supports_computer_use), + supportsComputerUse: computerModels.some((computer_model) => + litellmModelName.endsWith(computer_model), + ), supportsPromptCache: Boolean(modelInfo.supports_prompt_caching), inputPrice: modelInfo.input_cost_per_token ? modelInfo.input_cost_per_token * 1000000 : undefined, outputPrice: modelInfo.output_cost_per_token