Sfoglia il codice sorgente

post-Roo merge Ollama fixes

Christiaan Arnoldus 4 mesi fa
parent
commit
4ff00d588d

+ 1 - 1
src/api/providers/fetchers/modelCache.ts

@@ -101,7 +101,7 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
 				break
 			// kilocode_change end
 			case "ollama":
-				models = await getOllamaModels(options.baseUrl, options.apiKey)
+				models = await getOllamaModels(options.baseUrl, options.apiKey, options.numCtx /*kilocode_change*/)
 				break
 			case "lmstudio":
 				models = await getLMStudioModels(options.baseUrl)

+ 10 - 2
src/api/providers/fetchers/ollama.ts

@@ -39,7 +39,10 @@ type OllamaModelInfoResponse = z.infer<typeof OllamaModelInfoResponseSchema>
 
 export const parseOllamaModel = (
 	rawModel: OllamaModelInfoResponse,
-	baseUrl?: string, // kilocode_change
+	// kilocode_change start
+	baseUrl?: string,
+	numCtx?: number,
+	// kilocode_change end
 ): ModelInfo => {
 	// kilocode_change start
 	const contextKey = Object.keys(rawModel.model_info).find((k) => k.includes("context_length"))
@@ -54,6 +57,7 @@ export const parseOllamaModel = (
 	const contextLengthFromEnvironment = parseInt(process.env.OLLAMA_CONTEXT_LENGTH ?? "", 10) || undefined
 
 	const contextWindow =
+		numCtx ??
 		(baseUrl?.toLowerCase().startsWith("https://ollama.com") ? contextLengthFromModelInfo : undefined) ??
 		contextLengthFromEnvironment ??
 		(contextLengthFromModelParameters !== 40960 ? contextLengthFromModelParameters : undefined) ?? // Alledgedly Ollama sometimes returns an undefind context as 40960
@@ -75,6 +79,7 @@ export const parseOllamaModel = (
 export async function getOllamaModels(
 	baseUrl = "http://localhost:11434",
 	apiKey?: string,
+	numCtx?: number, // kilocode_change
 ): Promise<Record<string, ModelInfo>> {
 	const models: Record<string, ModelInfo> = {}
 
@@ -110,7 +115,10 @@ export async function getOllamaModels(
 						.then((ollamaModelInfo) => {
 							models[ollamaModel.name] = parseOllamaModel(
 								ollamaModelInfo.data,
-								baseUrl, // kilocode_change
+								// kilocode_change start
+								baseUrl,
+								numCtx,
+								// kilocode_change end
 							)
 						}),
 				)

+ 11 - 4
src/api/providers/native-ollama.ts

@@ -211,9 +211,10 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
 
 		// kilocode_change start
 		const estimatedTokenCount = estimateOllamaTokenCount(ollamaMessages)
-		if (modelInfo.maxTokens && estimatedTokenCount > modelInfo.maxTokens) {
+		const maxTokens = this.options.ollamaNumCtx ?? modelInfo.contextWindow
+		if (maxTokens && estimatedTokenCount > maxTokens) {
 			throw new Error(
-				`Input message is too long for the selected model. Estimated tokens: ${estimatedTokenCount}, Max tokens: ${modelInfo.maxTokens}. To increase the context window size, see: https://kilocode.ai/docs/providers/ollama#configure-the-context-size`,
+				`Input message is too long for the selected model. Estimated tokens: ${estimatedTokenCount}, Max tokens: ${maxTokens}. To increase the context window size, see: https://kilocode.ai/docs/providers/ollama#configure-the-context-size`,
 			)
 		}
 		// kilocode_change end
@@ -307,8 +308,14 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
 	}
 
 	async fetchModel() {
-		this.models = await getOllamaModels(this.options.ollamaBaseUrl, this.options.ollamaApiKey)
-		return this.models // kilocode_change
+		// kilocode_change start
+		this.models = await getOllamaModels(
+			this.options.ollamaBaseUrl,
+			this.options.ollamaApiKey,
+			this.options.ollamaNumCtx,
+		)
+		return this.models
+		// kilocode_change end
 	}
 
 	override getModel(): { id: string; info: ModelInfo } {

+ 1 - 0
src/core/webview/webviewMessageHandler.ts

@@ -931,6 +931,7 @@ export const webviewMessageHandler = async (
 					provider: "ollama",
 					baseUrl: ollamaApiConfig.ollamaBaseUrl,
 					apiKey: ollamaApiConfig.ollamaApiKey,
+					numCtx: ollamaApiConfig.ollamaNumCtx, // kilocode_change
 				})
 
 				if (Object.keys(ollamaModels).length > 0) {

+ 1 - 1
src/shared/api.ts

@@ -163,7 +163,7 @@ const dynamicProviderExtras = {
 	requesty: {} as { apiKey?: string; baseUrl?: string },
 	unbound: {} as { apiKey?: string },
 	glama: {} as {}, // eslint-disable-line @typescript-eslint/no-empty-object-type
-	ollama: {} as {}, // eslint-disable-line @typescript-eslint/no-empty-object-type
+	ollama: {} as { numCtx?: number }, // kilocode_change
 	lmstudio: {} as {}, // eslint-disable-line @typescript-eslint/no-empty-object-type
 } as const satisfies Record<RouterName, object>