From f4ceb53738250b194c3231579eb8f826b57d366a Mon Sep 17 00:00:00 2001 From: w1gs Date: Wed, 18 Sep 2024 02:02:43 -0400 Subject: [PATCH] Added Ollama integration --- extensions/void/package.json | 10 +- extensions/void/src/SidebarWebviewProvider.ts | 5 +- extensions/void/src/common/sendLLMMessage.ts | 124 ++++++++++-------- extensions/void/src/extension.ts | 3 +- 4 files changed, 83 insertions(+), 59 deletions(-) diff --git a/extensions/void/package.json b/extensions/void/package.json index ebccfe72..21be74ed 100644 --- a/extensions/void/package.json +++ b/extensions/void/package.json @@ -40,10 +40,15 @@ "default": "", "description": "Greptile - Github PAT (gives Greptile access to your repo)" }, - "void.ollamaSettings": { + "void.ollamaSettings.endpoint": { "type": "string", "default": "", - "description": "Ollama settings (coming soon...)" + "description": "Ollama Endpoint - Local API server can be started with `OLLAMA_ORIGINS=* ollama serve`" + }, + "void.ollamaSettings.model": { + "type": "string", + "default": "", + "description": "Ollama model to use" } } }, @@ -126,7 +131,6 @@ "eslint-plugin-react-hooks": "^4.6.2", "globals": "^15.9.0", "marked": "^14.1.0", - "ollama": "^0.5.8", "postcss": "^8.4.41", "react": "^18.3.1", "react-dom": "^18.3.1", diff --git a/extensions/void/src/SidebarWebviewProvider.ts b/extensions/void/src/SidebarWebviewProvider.ts index fe01c800..46c45085 100644 --- a/extensions/void/src/SidebarWebviewProvider.ts +++ b/extensions/void/src/SidebarWebviewProvider.ts @@ -54,8 +54,9 @@ export class SidebarWebviewProvider implements vscode.WebviewViewProvider { const nonce = getNonce(); // only scripts with the nonce are allowed to run, this is a recommended security measure - - const allowed_urls = ['https://api.anthropic.com', 'https://api.openai.com', 'https://api.greptile.com'] + // Allow Ollama endpoint + const ollamaEndpoint = vscode.workspace.getConfiguration('void').get('ollamaSettings.endpoint') || 'http://localhost:11434' + const allowed_urls = ['https://api.anthropic.com', 'https://api.openai.com', 'https://api.greptile.com', ollamaEndpoint ] webview.html = ` diff --git a/extensions/void/src/common/sendLLMMessage.ts b/extensions/void/src/common/sendLLMMessage.ts index ca3c34bd..bf87f40b 100644 --- a/extensions/void/src/common/sendLLMMessage.ts +++ b/extensions/void/src/common/sendLLMMessage.ts @@ -1,8 +1,6 @@ import Anthropic from '@anthropic-ai/sdk'; import OpenAI from 'openai'; -// import ollama from 'ollama' - export type ApiConfig = { anthropic: { apikey: string, @@ -20,7 +18,8 @@ export type ApiConfig = { } }, ollama: { - // TODO + endpoint: string, + model: string }, whichApi: string } @@ -220,66 +219,85 @@ const sendGreptileMsg: SendLLMMessageFnTypeInternal = ({ messages, onText, onFin export const sendLLMMessage: SendLLMMessageFnTypeExternal = ({ messages, onText, onFinalMessage, apiConfig }) => { if (!apiConfig) return { abort: () => { } } - const whichApi = apiConfig.whichApi + const whichApi = apiConfig.whichApi; - if (whichApi === 'anthropic') { - return sendClaudeMsg({ messages, onText, onFinalMessage, apiConfig }) + switch (whichApi) { + case 'anthropic': + return sendClaudeMsg({ messages, onText, onFinalMessage, apiConfig }); + case 'openai': + return sendOpenAIMsg({ messages, onText, onFinalMessage, apiConfig }); + case 'greptile': + return sendGreptileMsg({ messages, onText, onFinalMessage, apiConfig }); + case 'ollama': + return sendOllamaMsg({ messages, onText, onFinalMessage, apiConfig }); + default: + console.error(`Error: whichApi was ${whichApi}, which is not recognized!`); + return sendClaudeMsg({ messages, onText, onFinalMessage, apiConfig }); // TODO } - else if (whichApi === 'openai') { - return sendOpenAIMsg({ messages, onText, onFinalMessage, apiConfig }) - } - else if (whichApi === 'greptile') { - return sendGreptileMsg({ messages, onText, onFinalMessage, apiConfig }) - } - else if (whichApi === 'ollama') { - return sendClaudeMsg({ messages, onText, onFinalMessage, apiConfig }) // TODO - } - else { - console.error(`Error: whichApi was ${whichApi}, which is not recognized!`) - return sendClaudeMsg({ messages, onText, onFinalMessage, apiConfig }) // TODO - } - } // Ollama -// const sendOllamaMsg: sendMsgFnType = ({ messages, onText, onFinalMessage }) => { +export const sendOllamaMsg: SendLLMMessageFnTypeInternal = ({ messages, onText, onFinalMessage, apiConfig }) => { + let didAbort = false; + let fullText = ""; -// let did_abort = false -// let fullText = '' + // if abort is called, onFinalMessage is NOT called, and no later onTexts are called either + const abort = () => { + didAbort = true; + }; -// // if abort is called, onFinalMessage is NOT called, and no later onTexts are called either -// let abort: () => void = () => { -// did_abort = true -// } + const handleError = (error: any) => { + console.error('Error:', error); + onFinalMessage(fullText); + }; -// ollama.chat({ model: 'llama3.1', messages: messages, stream: true }) -// .then(async response => { + if (apiConfig.ollama.endpoint.endsWith('/')) { + apiConfig.ollama.endpoint = apiConfig.ollama.endpoint.slice(0, -1); + } -// abort = () => { -// // response.abort() // this isn't needed now, to keep consistency with claude will leave it commented for now -// did_abort = true; -// } + fetch(`${apiConfig.ollama.endpoint}/api/chat`, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model: apiConfig.ollama.model, + messages: messages, + stream: true, + }), + }) + .then(response => { + if (didAbort) return; + const reader = response.body?.getReader(); + if (!reader) { + onFinalMessage(fullText); + return; + } + return reader; + }) + .then(reader => { + if (!reader) return; -// // when receive text -// try { -// for await (const part of response) { -// if (did_abort) return -// let newText = part.message.content -// fullText += newText -// onText(newText, fullText) -// } -// } -// // when error/fail -// catch (e) { -// onFinalMessage(fullText) -// return -// } + const readStream = async () => { + try { + let done, value; + while ({ done, value } = await reader.read(), !done) { + if (didAbort) return; + const stringedResponse = new TextDecoder().decode(value); + const newText = JSON.parse(stringedResponse).message.content; + fullText += newText; + onText(newText, fullText); + } + onFinalMessage(fullText); + } catch (error) { + handleError(error); + } + }; -// // when we get the final message on this stream -// onFinalMessage(fullText) -// }) - -// return { abort }; -// }; + readStream(); + }) + .catch(handleError); + return { abort }; +}; diff --git a/extensions/void/src/extension.ts b/extensions/void/src/extension.ts index 4cd8ac26..d95c2873 100644 --- a/extensions/void/src/extension.ts +++ b/extensions/void/src/extension.ts @@ -25,7 +25,8 @@ const getApiConfig = () => { } }, ollama: { - // apikey: vscode.workspace.getConfiguration('void').get('ollamaSettings') ?? '', + endpoint: vscode.workspace.getConfiguration('void').get('ollamaSettings.endpoint') ?? '', + model: vscode.workspace.getConfiguration('void').get('ollamaSettings.model') ?? '', }, whichApi: vscode.workspace.getConfiguration('void').get('whichApi') ?? '' }