From 9eaf7530002a9918ad3d3652d0d23df0e85c3f02 Mon Sep 17 00:00:00 2001 From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com> Date: Fri, 6 Sep 2024 04:51:17 -0400 Subject: [PATCH] Add ollama models polling --- package.json | 2 +- src/api/ollama.ts | 2 +- src/providers/ClaudeDevProvider.ts | 29 ++++++++ src/shared/ExtensionMessage.ts | 3 +- src/shared/WebviewMessage.ts | 1 + src/shared/api.ts | 1 + webview-ui/src/components/ApiOptions.tsx | 85 +++++++++++++++++++++--- 7 files changed, 111 insertions(+), 12 deletions(-) diff --git a/package.json b/package.json index 5928c78e29..401312a887 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "claude-dev", "displayName": "Claude Dev", "description": "Autonomous coding agent right in your IDE, capable of creating/editing files, executing commands, and more with your permission every step of the way.", - "version": "1.5.31", + "version": "1.5.32", "icon": "icon.png", "engines": { "vscode": "^1.84.0" diff --git a/src/api/ollama.ts b/src/api/ollama.ts index 323f99e3ea..d2bd0fef12 100644 --- a/src/api/ollama.ts +++ b/src/api/ollama.ts @@ -11,7 +11,7 @@ export class OllamaHandler implements ApiHandler { constructor(options: ApiHandlerOptions) { this.options = options this.client = new OpenAI({ - baseURL: "http://localhost:11434/v1", + baseURL: (this.options.ollamaBaseUrl || "http://localhost:11434") + "/v1", apiKey: "ollama", }) } diff --git a/src/providers/ClaudeDevProvider.ts b/src/providers/ClaudeDevProvider.ts index e0daf3ade1..f44e97f543 100644 --- a/src/providers/ClaudeDevProvider.ts +++ b/src/providers/ClaudeDevProvider.ts @@ -30,6 +30,7 @@ type GlobalStateKey = | "openAiBaseUrl" | "openAiModelId" | "ollamaModelId" + | "ollamaBaseUrl" export class ClaudeDevProvider implements vscode.WebviewViewProvider { public static readonly sideBarId = "claude-dev.SidebarProvider" // used in package.json as the view's id. This value cannot be changed due to how vscode caches views based on their id, and updating the id would break existing instances of the extension. @@ -321,6 +322,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { openAiApiKey, openAiModelId, ollamaModelId, + ollamaBaseUrl, } = message.apiConfiguration await this.updateGlobalState("apiProvider", apiProvider) await this.updateGlobalState("apiModelId", apiModelId) @@ -336,6 +338,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { await this.storeSecret("openAiApiKey", openAiApiKey) await this.updateGlobalState("openAiModelId", openAiModelId) await this.updateGlobalState("ollamaModelId", ollamaModelId) + await this.updateGlobalState("ollamaBaseUrl", ollamaBaseUrl) this.claudeDev?.updateApi(message.apiConfiguration) } await this.postStateToWebview() @@ -385,6 +388,10 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { case "resetState": await this.resetState() break + case "requestOllamaModels": + const models = await this.getOllamaModels(message.text) + this.postMessageToWebview({ type: "ollamaModels", models }) + break // Add more switch case statements here as more webview message commands // are created within the webview context (i.e. inside media/main.js) } @@ -394,6 +401,25 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { ) } + // Ollama + + async getOllamaModels(baseUrl?: string) { + try { + if (!baseUrl) { + baseUrl = "http://localhost:11434" + } + if (!URL.canParse(baseUrl)) { + return [] + } + const response = await axios.get(`${baseUrl}/api/tags`) + const modelsArray = response.data?.models?.map((model: any) => model.name) || [] + const models = [...new Set(modelsArray)] + return models + } catch (error) { + return [] + } + } + // OpenRouter async handleOpenRouterCallback(code: string) { @@ -627,6 +653,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { openAiApiKey, openAiModelId, ollamaModelId, + ollamaBaseUrl, lastShownAnnouncementId, customInstructions, alwaysAllowReadOnly, @@ -646,6 +673,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { this.getSecret("openAiApiKey") as Promise, this.getGlobalState("openAiModelId") as Promise, this.getGlobalState("ollamaModelId") as Promise, + this.getGlobalState("ollamaBaseUrl") as Promise, this.getGlobalState("lastShownAnnouncementId") as Promise, this.getGlobalState("customInstructions") as Promise, this.getGlobalState("alwaysAllowReadOnly") as Promise, @@ -682,6 +710,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { openAiApiKey, openAiModelId, ollamaModelId, + ollamaBaseUrl, }, lastShownAnnouncementId, customInstructions, diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index 6c562b02bf..100100f8a9 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -5,11 +5,12 @@ import { HistoryItem } from "./HistoryItem" // webview will hold state export interface ExtensionMessage { - type: "action" | "state" | "selectedImages" + type: "action" | "state" | "selectedImages" | "ollamaModels" text?: string action?: "chatButtonTapped" | "settingsButtonTapped" | "historyButtonTapped" | "didBecomeVisible" state?: ExtensionState images?: string[] + models?: string[] } export interface ExtensionState { diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index d8a47626de..804e0aaa45 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -16,6 +16,7 @@ export interface WebviewMessage { | "deleteTaskWithId" | "exportTaskWithId" | "resetState" + | "requestOllamaModels" text?: string askResponse?: ClaudeAskResponse apiConfiguration?: ApiConfiguration diff --git a/src/shared/api.ts b/src/shared/api.ts index cac46aaba3..fa71b1827e 100644 --- a/src/shared/api.ts +++ b/src/shared/api.ts @@ -14,6 +14,7 @@ export interface ApiHandlerOptions { openAiApiKey?: string openAiModelId?: string ollamaModelId?: string + ollamaBaseUrl?: string } export type ApiConfiguration = ApiHandlerOptions & { diff --git a/webview-ui/src/components/ApiOptions.tsx b/webview-ui/src/components/ApiOptions.tsx index fce7ac64d1..0d87edf7f6 100644 --- a/webview-ui/src/components/ApiOptions.tsx +++ b/webview-ui/src/components/ApiOptions.tsx @@ -1,5 +1,12 @@ -import { VSCodeDropdown, VSCodeLink, VSCodeOption, VSCodeTextField } from "@vscode/webview-ui-toolkit/react" -import React, { useMemo } from "react" +import { + VSCodeDropdown, + VSCodeLink, + VSCodeOption, + VSCodeRadio, + VSCodeRadioGroup, + VSCodeTextField, +} from "@vscode/webview-ui-toolkit/react" +import React, { useCallback, useEffect, useMemo, useState } from "react" import { ApiConfiguration, ModelInfo, @@ -15,6 +22,9 @@ import { } from "../../../src/shared/api" import { useExtensionState } from "../context/ExtensionStateContext" import VSCodeButtonLink from "./VSCodeButtonLink" +import { ExtensionMessage } from "../../../src/shared/ExtensionMessage" +import { useEvent, useInterval } from "react-use" +import { vscode } from "../utils/vscode" interface ApiOptionsProps { showModelOptions: boolean @@ -23,6 +33,8 @@ interface ApiOptionsProps { const ApiOptions: React.FC = ({ showModelOptions, apiErrorMessage }) => { const { apiConfiguration, setApiConfiguration, uriScheme } = useExtensionState() + const [ollamaModels, setOllamaModels] = useState([]) + const handleInputChange = (field: keyof ApiConfiguration) => (event: any) => { setApiConfiguration({ ...apiConfiguration, [field]: event.target.value }) } @@ -31,6 +43,27 @@ const ApiOptions: React.FC = ({ showModelOptions, apiErrorMessa return normalizeApiConfiguration(apiConfiguration) }, [apiConfiguration]) + // Poll ollama models + const requestOllamaModels = useCallback(() => { + if (selectedProvider === "ollama") { + vscode.postMessage({ type: "requestOllamaModels", text: apiConfiguration?.ollamaBaseUrl }) + } + }, [selectedProvider, apiConfiguration?.ollamaBaseUrl]) + useEffect(() => { + if (selectedProvider === "ollama") { + requestOllamaModels() + } + }, [selectedProvider, requestOllamaModels]) + useInterval(requestOllamaModels, selectedProvider === "ollama" ? 2000 : null) + + const handleMessage = useCallback((event: MessageEvent) => { + const message: ExtensionMessage = event.data + if (message.type === "ollamaModels" && message.models) { + setOllamaModels(message.models) + } + }, []) + useEvent("message", handleMessage) + /* VSCodeDropdown has an open bug where dynamically rendered options don't auto select the provided value prop. You can see this for yourself by comparing it with normal select/option elements, which work as expected. https://github.com/microsoft/vscode-webview-ui-toolkit/issues/433 @@ -295,8 +328,8 @@ const ApiOptions: React.FC = ({ showModelOptions, apiErrorMessa }}> You can use any OpenAI compatible API with models that support tool use.{" "} - (Note: Claude Dev uses complex prompts, so less - capable models may not work as expected.) + (Note: Claude Dev uses complex prompts and works + best with Claude models. Less capable models may not work as expected.)

@@ -304,6 +337,14 @@ const ApiOptions: React.FC = ({ showModelOptions, apiErrorMessa {selectedProvider === "ollama" && (
+ + Base URL (optional) + = ({ showModelOptions, apiErrorMessa placeholder={"e.g. llama3.1"}> Model ID + {ollamaModels.length > 0 && ( + { + const value = (e.target as HTMLInputElement)?.value + // need to check value first since radio group returns empty string sometimes + if (value) { + handleInputChange("ollamaModelId")({ + target: { value }, + }) + } + }}> + {ollamaModels.map((model) => ( + + {model} + + ))} + + )}

- Ollama allows you to run models locally on your computer. For instructions on how to get started - with Ollama, see their + Ollama allows you to run models locally on your computer. For instructions on how to get + started, see their quickstart guide. {" "} - You can use any models that support{" "} + You can use any model that supports{" "} tool use. - (Note: Claude Dev uses complex prompts, so less - capable models may not work as expected.) + (Note: Claude Dev uses complex prompts and works + best with Claude models. Less capable models may not work as expected.)