fix context length for lmstudio and ollama (#2462)

Author: thecolorblue

packages/types/src/providers/index.ts

@@ -8,6 +8,7 @@ export * from "./groq.js"
 export * from "./lite-llm.js"
 export * from "./lm-studio.js"
 export * from "./mistral.js"
+export * from "./ollama.js"
 export * from "./openai.js"
 export * from "./openrouter.js"
 export * from "./requesty.js"

packages/types/src/providers/lm-studio.ts

@@ -1 +1,19 @@
+import type { ModelInfo } from "../model.js"
+
 export const LMSTUDIO_DEFAULT_TEMPERATURE = 0
+
+// LM Studio
+// https://lmstudio.ai/docs/cli/ls
+export const lMStudioDefaultModelId = "mistralai/devstral-small-2505"
+export const lMStudioDefaultModelInfo: ModelInfo = {
+	maxTokens: 8192,
+	contextWindow: 200_000,
+	supportsImages: true,
+	supportsComputerUse: true,
+	supportsPromptCache: true,
+	inputPrice: 0,
+	outputPrice: 0,
+	cacheWritesPrice: 0,
+	cacheReadsPrice: 0,
+	description: "LM Studio hosted models",
+}

packages/types/src/providers/ollama.ts

@@ -0,0 +1,17 @@
+import type { ModelInfo } from "../model.js"
+
+// Ollama
+// https://ollama.com/models
+export const ollamaDefaultModelId = "devstral:24b"
+export const ollamaDefaultModelInfo: ModelInfo = {
+	maxTokens: 4096,
+	contextWindow: 200_000,
+	supportsImages: true,
+	supportsComputerUse: true,
+	supportsPromptCache: true,
+	inputPrice: 0,
+	outputPrice: 0,
+	cacheWritesPrice: 0,
+	cacheReadsPrice: 0,
+	description: "Ollama hosted models",
+}

pnpm-lock.yaml

@@ -0,0 +1,52 @@
+import { ModelInfo, lMStudioDefaultModelInfo } from "@roo-code/types"
+import { LLMInfo, LMStudioClient } from "@lmstudio/sdk"
+import axios from "axios"
+
+export const parseLMStudioModel = (rawModel: LLMInfo): ModelInfo => {
+	const modelInfo: ModelInfo = Object.assign({}, lMStudioDefaultModelInfo, {
+		description: `${rawModel.displayName} - ${rawModel.paramsString} - ${rawModel.path}`,
+		contextWindow: rawModel.maxContextLength,
+		supportsPromptCache: true,
+		supportsImages: rawModel.vision,
+		supportsComputerUse: false,
+		maxTokens: rawModel.maxContextLength,
+	})
+
+	return modelInfo
+}
+
+export async function getLMStudioModels(baseUrl = "http://localhost:1234"): Promise<Record<string, ModelInfo>> {
+	// clearing the input can leave an empty string; use the default in that case
+	baseUrl = baseUrl === "" ? "http://localhost:1234" : baseUrl
+
+	const models: Record<string, ModelInfo> = {}
+	// ws is required to connect using the LMStudio library
+	const lmsUrl = baseUrl.replace(/^http:\/\//, "ws://").replace(/^https:\/\//, "wss://")
+
+	try {
+		if (!URL.canParse(lmsUrl)) {
+			return models
+		}
+
+		// test the connection to LM Studio first
+		// errors will be caught further down
+		await axios.get(`${baseUrl}/v1/models`)
+
+		const client = new LMStudioClient({ baseUrl: lmsUrl })
+		const response = (await client.system.listDownloadedModels()) as Array<LLMInfo>
+
+		for (const lmstudioModel of response) {
+			models[lmstudioModel.modelKey] = parseLMStudioModel(lmstudioModel)
+		}
+	} catch (error) {
+		if (error.code === "ECONNREFUSED") {
+			console.error(`Error connecting to LMStudio at ${baseUrl}`)
+		} else {
+			console.error(
+				`Error fetching LMStudio models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`,
+			)
+		}
+	}
+
+	return models
+}

src/api/providers/fetchers/lmstudio.ts

@@ -14,6 +14,9 @@ import { getGlamaModels } from "./glama"
 import { getUnboundModels } from "./unbound"
 import { getLiteLLMModels } from "./litellm"
 import { GetModelsOptions } from "../../../shared/api"
+import { getOllamaModels } from "./ollama"
+import { getLMStudioModels } from "./lmstudio"
+
 const memoryCache = new NodeCache({ stdTTL: 5 * 60, checkperiod: 5 * 60 })
 
 async function writeModels(router: RouterName, data: ModelRecord) {
@@ -68,6 +71,12 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
 				// Type safety ensures apiKey and baseUrl are always provided for litellm
 				models = await getLiteLLMModels(options.apiKey, options.baseUrl)
 				break
+			case "ollama":
+				models = await getOllamaModels(options.baseUrl)
+				break
+			case "lmstudio":
+				models = await getLMStudioModels(options.baseUrl)
+				break
 			default: {
 				// Ensures router is exhaustively checked if RouterName is a strict union
 				const exhaustiveCheck: never = provider

src/api/providers/fetchers/modelCache.ts

@@ -0,0 +1,97 @@
+import axios from "axios"
+import { ModelInfo, ollamaDefaultModelInfo } from "@roo-code/types"
+import { z } from "zod"
+
+const OllamaModelDetailsSchema = z.object({
+	family: z.string(),
+	families: z.array(z.string()),
+	format: z.string(),
+	parameter_size: z.string(),
+	parent_model: z.string(),
+	quantization_level: z.string(),
+})
+
+const OllamaModelSchema = z.object({
+	details: OllamaModelDetailsSchema,
+	digest: z.string(),
+	model: z.string(),
+	modified_at: z.string(),
+	name: z.string(),
+	size: z.number(),
+})
+
+const OllamaModelInfoResponseSchema = z.object({
+	modelfile: z.string(),
+	parameters: z.string(),
+	template: z.string(),
+	details: OllamaModelDetailsSchema,
+	model_info: z.record(z.string(), z.any()),
+	capabilities: z.array(z.string()).optional(),
+})
+
+const OllamaModelsResponseSchema = z.object({
+	models: z.array(OllamaModelSchema),
+})
+
+type OllamaModelsResponse = z.infer<typeof OllamaModelsResponseSchema>
+
+type OllamaModelInfoResponse = z.infer<typeof OllamaModelInfoResponseSchema>
+
+export const parseOllamaModel = (rawModel: OllamaModelInfoResponse): ModelInfo => {
+	const contextKey = Object.keys(rawModel.model_info).find((k) => k.includes("context_length"))
+	const contextWindow = contextKey ? rawModel.model_info[contextKey] : undefined
+
+	const modelInfo: ModelInfo = Object.assign({}, ollamaDefaultModelInfo, {
+		description: `Family: ${rawModel.details.family}, Context: ${contextWindow}, Size: ${rawModel.details.parameter_size}`,
+		contextWindow: contextWindow || ollamaDefaultModelInfo.contextWindow,
+		supportsPromptCache: true,
+		supportsImages: rawModel.capabilities?.includes("vision"),
+		supportsComputerUse: false,
+		maxTokens: contextWindow || ollamaDefaultModelInfo.contextWindow,
+	})
+
+	return modelInfo
+}
+
+export async function getOllamaModels(baseUrl = "http://localhost:11434"): Promise<Record<string, ModelInfo>> {
+	const models: Record<string, ModelInfo> = {}
+
+	// clearing the input can leave an empty string; use the default in that case
+	baseUrl = baseUrl === "" ? "http://localhost:11434" : baseUrl
+
+	try {
+		if (!URL.canParse(baseUrl)) {
+			return models
+		}
+
+		const response = await axios.get<OllamaModelsResponse>(`${baseUrl}/api/tags`)
+		const parsedResponse = OllamaModelsResponseSchema.safeParse(response.data)
+		let modelInfoPromises = []
+
+		if (parsedResponse.success) {
+			for (const ollamaModel of parsedResponse.data.models) {
+				modelInfoPromises.push(
+					axios
+						.post<OllamaModelInfoResponse>(`${baseUrl}/api/show`, {
+							model: ollamaModel.model,
+						})
+						.then((ollamaModelInfo) => {
+							models[ollamaModel.name] = parseOllamaModel(ollamaModelInfo.data)
+						}),
+				)
+			}
+
+			await Promise.all(modelInfoPromises)
+		} else {
+			console.error(`Error parsing Ollama models response: ${JSON.stringify(parsedResponse.error, null, 2)}`)
+		}
+	} catch (error) {
+		if (error.code === "ECONNREFUSED") {
+			console.info(`Failed connecting to Ollama at ${baseUrl}`)
+		} else {
+			console.warn(`Error fetching Ollama models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`)
+		}
+	}
+
+	return models
+}

src/api/providers/fetchers/ollama.ts

@@ -1,6 +1,5 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import OpenAI from "openai"
-import axios from "axios"
 
 import { type ModelInfo, openAiModelInfoSaneDefaults, DEEP_SEEK_DEFAULT_TEMPERATURE } from "@roo-code/types"
 
@@ -111,17 +110,3 @@ export class OllamaHandler extends BaseProvider implements SingleCompletionHandl
 		}
 	}
 }
-
-export async function getOllamaModels(baseUrl = "http://localhost:11434") {
-	try {
-		if (!URL.canParse(baseUrl)) {
-			return []
-		}
-
-		const response = await axios.get(`${baseUrl}/api/tags`)
-		const modelsArray = response.data?.models?.map((model: any) => model.name) || []
-		return [...new Set<string>(modelsArray)]
-	} catch (error) {
-		return []
-	}
-}

src/api/providers/ollama.ts

@@ -29,9 +29,7 @@ import { singleCompletionHandler } from "../../utils/single-completion-handler"
 import { searchCommits } from "../../utils/git"
 import { exportSettings, importSettings } from "../config/importExport"
 import { getOpenAiModels } from "../../api/providers/openai"
-import { getOllamaModels } from "../../api/providers/ollama"
 import { getVsCodeLmModels } from "../../api/providers/vscode-lm"
-import { getLmStudioModels } from "../../api/providers/lm-studio"
 import { openMention } from "../mentions"
 import { TelemetrySetting } from "../../shared/TelemetrySetting"
 import { getWorkspacePath } from "../../utils/path"
@@ -379,6 +377,19 @@ export const webviewMessageHandler = async (
 
 				if (result.status === "fulfilled") {
 					fetchedRouterModels[routerName] = result.value.models
+
+					// Ollama and LM Studio settings pages still need these events
+					if (routerName === "ollama" && Object.keys(result.value.models).length > 0) {
+						provider.postMessageToWebview({
+							type: "ollamaModels",
+							ollamaModels: Object.keys(result.value.models),
+						})
+					} else if (routerName === "lmstudio" && Object.keys(result.value.models).length > 0) {
+						provider.postMessageToWebview({
+							type: "lmStudioModels",
+							lmStudioModels: Object.keys(result.value.models),
+						})
+					}
 				} else {
 					// Handle rejection: Post a specific error message for this provider
 					const errorMessage = result.reason instanceof Error ? result.reason.message : String(result.reason)
@@ -399,6 +410,7 @@ export const webviewMessageHandler = async (
 				type: "routerModels",
 				routerModels: fetchedRouterModels as Record<RouterName, ModelRecord>,
 			})
+
 			break
 		case "requestOpenAiModels":
 			if (message?.values?.baseUrl && message?.values?.apiKey) {
@@ -411,16 +423,6 @@ export const webviewMessageHandler = async (
 				provider.postMessageToWebview({ type: "openAiModels", openAiModels })
 			}
 
-			break
-		case "requestOllamaModels":
-			const ollamaModels = await getOllamaModels(message.text)
-			// TODO: Cache like we do for OpenRouter, etc?
-			provider.postMessageToWebview({ type: "ollamaModels", ollamaModels })
-			break
-		case "requestLmStudioModels":
-			const lmStudioModels = await getLmStudioModels(message.text)
-			// TODO: Cache like we do for OpenRouter, etc?
-			provider.postMessageToWebview({ type: "lmStudioModels", lmStudioModels })
 			break
 		case "requestVsCodeLmModels":
 			const vsCodeLmModels = await getVsCodeLmModels()

src/core/webview/webviewMessageHandler.ts

@@ -369,6 +369,7 @@
 		"@aws-sdk/client-bedrock-runtime": "^3.779.0",
 		"@aws-sdk/credential-providers": "^3.806.0",
 		"@google/genai": "^1.0.0",
+		"@lmstudio/sdk": "^1.1.1",
 		"@mistralai/mistralai": "^1.3.6",
 		"@modelcontextprotocol/sdk": "^1.9.0",
 		"@qdrant/js-client-rest": "^1.14.0",