feat: ask to download model

zhutony · Dec 22, 2023 · 665b7f4 · 665b7f4
1 parent 1d22c2a
commit 665b7f4
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -42,10 +42,15 @@ Currently Llama Coder supports only Codellama. Model is quantized in different w
 * m - slow on MacOS
 * g - slow on older NVidia cards (pre 30xx)
 
+## Troubleshooting
+
+Most of the problems could be seen in output of a plugin in VS Code extension output.
+
 ## Changelog
 
 ## [0.0.10]
 - Adding ability to pick a custom model
+- Asking user if he wants to download model if it is not available
 
 ## [0.0.9]
 - Adding deepseek 1b model and making it default

diff --git a/src/extension.ts b/src/extension.ts
@@ -21,7 +21,7 @@ export function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(statusBarItem);
 
 	// Create provider
-	const provider = new PromptProvider(statusBarItem);
+	const provider = new PromptProvider(statusBarItem, context);
 	let disposable = vscode.languages.registerInlineCompletionItemProvider({ pattern: '**', }, provider);
 	context.subscriptions.push(disposable);
 }

diff --git a/src/prompts/provider.ts b/src/prompts/provider.ts
@@ -13,9 +13,11 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
 
     lock = new AsyncLock();
     statusbar: vscode.StatusBarItem;
+    context: vscode.ExtensionContext;
 
-    constructor(statusbar: vscode.StatusBarItem) {
+    constructor(statusbar: vscode.StatusBarItem, context: vscode.ExtensionContext) {
         this.statusbar = statusbar;
+        this.context = context;
     }
 
     async provideInlineCompletionItems(document: vscode.TextDocument, position: vscode.Position, context: vscode.InlineCompletionContext, token: vscode.CancellationToken): Promise<vscode.InlineCompletionItem[] | vscode.InlineCompletionList | undefined | null> {
@@ -87,6 +89,22 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
 
                         // Download model if not exists
                         if (!modelExists) {
+
+                            // Check if user asked to ignore download
+                            if (this.context.globalState.get('llama-coder-download-ignored')) {
+                                info(`Ingoring since user asked to ignore download.`);
+                                return;
+                            }
+
+                            // Ask for download
+                            let download = await vscode.window.showInformationMessage(`Model ${inferenceConfig.modelName} is not downloaded. Do you want to download it? Answering "No" would require you to manually download model.`, 'Yes', 'No');
+                            if (download === 'No') {
+                                info(`Ingoring since user asked to ignore download.`);
+                                this.context.globalState.update('llama-coder-download-ignored', true);
+                                return;
+                            }
+
+                            // Perform download
                             this.statusbar.text = `$(sync~spin) Downloading`;
                             await ollamaDownloadModel(inferenceConfig.endpoint, inferenceConfig.modelName);
                             this.statusbar.text = `$(sync~spin) Llama Coder`;