Error handling and other improvements (#56)

* safety error handling and other improvements * lift vscode api from provider * return success
CodingTrain · Apr 2, 2024 · 966d0ec · 966d0ec
1 parent 2818d2e
commit 966d0ec
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/src/lib/agent/Agent.js b/src/lib/agent/Agent.js
@@ -2,7 +2,6 @@ const { typeRealistically } = require('../../util/realisticTyping');
 const vscode = require('vscode');
 const { getProvider } = require('./providers/providerInstance');
 const { speak } = require('../../util/speak');
-const config = require('../../../config');
 
 class Agent {
   constructor() {
@@ -16,25 +15,37 @@ class Agent {
     this.isNewPrompt = false;
     this.promptingTemplate =
       'Dan says: {prompt}\nCurrent code in the editor:\n```\n{currentCode}\n```';
+    this.isStreaming = false;
   }
 
   /**
    * Prompt something to the AI
    * @param {string} input The prompt to be processed
    */
   prompt(input) {
+    if (this.isStreaming || this.actionsQueue.length > 0) {
+      return vscode.window.showErrorMessage(
+        'Please wait for the current prompt to finish processing before sending another one.'
+      );
+    }
     const editor = vscode.window.activeTextEditor;
 
     this.isNewPrompt = true;
     const prompt = this.promptingTemplate
       .replace('{prompt}', input)
       .replace('{currentCode}', editor.document.getText());
-    this.provider.queryStream(prompt, (response) =>
-      this.consumeStream(response)
-    );
+    this.isStreaming = true;
+    this.provider
+      .queryStream(prompt, (response) => this.consumeStream(response))
+      .then((out) => {
+        this.isStreaming = false;
+        if (out.blocked) {
+          vscode.window.showErrorMessage(`Prompt blocked: ${out.blockReason}`);
+        }
+      });
   }
 
-  async consumeStream(response) {
+  consumeStream(response) {
     const text = response.response;
     const event = response.event;
 

diff --git a/src/lib/agent/providers/geminiProvider.js b/src/lib/agent/providers/geminiProvider.js
@@ -47,12 +47,23 @@ class GeminiProvider extends ModelProvider {
     const result = await this.chat.sendMessageStream(prompt);
     let fullResponse = '';
     for await (const chunk of result.stream) {
+      if (chunk.promptFeedback?.blockReason) {
+        // prompt blocked
+        this.messageHistory.pop();
+        return {
+          blocked: true,
+          blockReason: chunk.promptFeedback.blockReason,
+        };
+      }
       const text = chunk.text();
       fullResponse += text;
       await process({ response: text, event: 'output' });
     }
     await process({ response: '', event: 'done' });
     this.messageHistory.push({ role: 'assistant', content: fullResponse });
+    return {
+      success: true,
+    };
   }
 }