diff --git a/package.json b/package.json index f3c355a..2502f47 100644 --- a/package.json +++ b/package.json @@ -76,6 +76,11 @@ "description": "Ollama Server Endpoint. Empty for local instance. Example: http://192.168.0.100:11434", "order": 1 }, + "inference.bearerToken": { + "type": "string", + "default": "", + "description": "Auth Bearer token that should be used for secure requests. Leave empty if not desired." + }, "inference.model": { "type": "string", "enum": [ diff --git a/src/config.ts b/src/config.ts index 7d0a714..a6294b1 100644 --- a/src/config.ts +++ b/src/config.ts @@ -15,6 +15,7 @@ class Config { if (endpoint === '') { endpoint = 'http://127.0.0.1:11434'; } + let bearerToken = config.get('bearerToken') as string; // Load general paremeters let maxLines = config.get('maxLines') as number; @@ -39,6 +40,7 @@ class Config { return { endpoint, + bearerToken, maxLines, maxTokens, temperature, diff --git a/src/modules/lineGenerator.ts b/src/modules/lineGenerator.ts index e61b9a6..f20081f 100644 --- a/src/modules/lineGenerator.ts +++ b/src/modules/lineGenerator.ts @@ -1,14 +1,16 @@ -export async function* lineGenerator(url: string, data: any): AsyncGenerator { - +export async function* lineGenerator(url: string, data: any, bearerToken: string): AsyncGenerator { // Request const controller = new AbortController(); let res = await fetch(url, { - method: 'POST', - body: JSON.stringify(data), - headers: { - "Content-Type": "application/json", - }, - signal: controller.signal + method: 'POST', + body: JSON.stringify(data), + headers: bearerToken ? { + 'Content-Type': 'application/json', + Authorization: `Bearer ${bearerToken}`, + } : { + 'Content-Type': 'application/json', + }, + signal: controller.signal, }); if (!res.ok || !res.body) { throw Error('Unable to connect to backend'); diff --git a/src/modules/ollamaCheckModel.ts b/src/modules/ollamaCheckModel.ts index 86ae281..5aedb8f 100644 --- a/src/modules/ollamaCheckModel.ts +++ b/src/modules/ollamaCheckModel.ts @@ -1,9 +1,12 @@ import { info } from "./log"; -export async function ollamaCheckModel(endpoint: string, model: string) { - +export async function ollamaCheckModel(endpoint: string, model: string, bearerToken: string) { // Check if exists - let res = await fetch(endpoint + '/api/tags'); + let res = await fetch(endpoint + '/api/tags', { + headers: bearerToken ? { + Authorization: `Bearer ${bearerToken}`, + } : {}, + }); if (!res.ok) { info(await res.text()); info(endpoint + '/api/tags'); diff --git a/src/modules/ollamaDownloadModel.ts b/src/modules/ollamaDownloadModel.ts index 31e2d92..7e6eccb 100644 --- a/src/modules/ollamaDownloadModel.ts +++ b/src/modules/ollamaDownloadModel.ts @@ -1,9 +1,9 @@ import { lineGenerator } from "./lineGenerator"; import { info } from "./log"; -export async function ollamaDownloadModel(endpoint: string, model: string) { +export async function ollamaDownloadModel(endpoint: string, model: string, bearerToken: string) { info('Downloading model from ollama: ' + model); - for await (let line of lineGenerator(endpoint + '/api/pull', { name: model })) { + for await (let line of lineGenerator(endpoint + '/api/pull', { name: model }, bearerToken)) { info('[DOWNLOAD] ' + line); } } \ No newline at end of file diff --git a/src/modules/ollamaTokenGenerator.ts b/src/modules/ollamaTokenGenerator.ts index 04e4b52..f57fef9 100644 --- a/src/modules/ollamaTokenGenerator.ts +++ b/src/modules/ollamaTokenGenerator.ts @@ -7,8 +7,8 @@ export type OllamaToken = { done: boolean }; -export async function* ollamaTokenGenerator(url: string, data: any): AsyncGenerator { - for await (let line of lineGenerator(url, data)) { +export async function* ollamaTokenGenerator(url: string, data: any, bearerToken: string): AsyncGenerator { + for await (let line of lineGenerator(url, data, bearerToken)) { info('Receive line: ' + line); let parsed: OllamaToken; try { diff --git a/src/prompts/autocomplete.ts b/src/prompts/autocomplete.ts index b5b9d8f..b6a8089 100644 --- a/src/prompts/autocomplete.ts +++ b/src/prompts/autocomplete.ts @@ -5,6 +5,7 @@ import { ModelFormat, adaptPrompt } from './processors/models'; export async function autocomplete(args: { endpoint: string, + bearerToken: string, model: string, format: ModelFormat, prefix: string, @@ -33,7 +34,7 @@ export async function autocomplete(args: { let res = ''; let totalLines = 1; let blockStack: ('[' | '(' | '{')[] = []; - outer: for await (let tokens of ollamaTokenGenerator(args.endpoint + '/api/generate', data)) { + outer: for await (let tokens of ollamaTokenGenerator(args.endpoint + '/api/generate', data, args.bearerToken)) { if (args.canceled && args.canceled()) { break; } diff --git a/src/prompts/provider.ts b/src/prompts/provider.ts index 7820b04..ed4be76 100644 --- a/src/prompts/provider.ts +++ b/src/prompts/provider.ts @@ -122,7 +122,7 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider { try { // Check model exists - let modelExists = await ollamaCheckModel(inferenceConfig.endpoint, inferenceConfig.modelName); + let modelExists = await ollamaCheckModel(inferenceConfig.endpoint, inferenceConfig.modelName, inferenceConfig.bearerToken); if (token.isCancellationRequested) { info(`Canceled after AI completion.`); return; @@ -147,7 +147,7 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider { // Perform download this.update('sync~spin', 'Downloading'); - await ollamaDownloadModel(inferenceConfig.endpoint, inferenceConfig.modelName); + await ollamaDownloadModel(inferenceConfig.endpoint, inferenceConfig.modelName, inferenceConfig.bearerToken); this.update('sync~spin', 'Llama Coder') } if (token.isCancellationRequested) { @@ -161,6 +161,7 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider { prefix: prepared.prefix, suffix: prepared.suffix, endpoint: inferenceConfig.endpoint, + bearerToken: inferenceConfig.bearerToken, model: inferenceConfig.modelName, format: inferenceConfig.modelFormat, maxLines: inferenceConfig.maxLines,