feat(gemini): Improve tool call processing and error handling

- Enhance the tool call processing in GeminiLLM to better handle function results and planning.
- Update the prompt construction to include plan information when available.
- Implement retry mechanism with user prompt for failed function calls with fallback to Groq.
- Add timeout mechanism for generateContentWithTools and fallback to Groq on timeout.
- Change the default CodeBuddy mode to Ask.

Author: olasunkanmiraymond

src/llms/gemini/gemini.ts

@@ -60,8 +60,8 @@ export class GeminiLLM
   private intializeDisposable(): void {
     this.disposables.push(
       vscode.workspace.onDidChangeConfiguration(() =>
-        this.handleConfigurationChange(),
-      ),
+        this.handleConfigurationChange()
+      )
     );
   }
 
@@ -91,7 +91,7 @@ export class GeminiLLM
 
   public async generateText(
     prompt: string,
-    instruction?: string,
+    instruction?: string
   ): Promise<string> {
     try {
       const model = this.getModel();
@@ -141,15 +141,15 @@ export class GeminiLLM
   }
 
   async generateContentWithTools(
-    userInput: string,
+    userInput: string
   ): Promise<GenerateContentResult> {
     try {
       await this.buildChatHistory(
         userInput,
         undefined,
         undefined,
         undefined,
-        true,
+        true
       );
       const prompt = createPrompt(userInput);
       const contents = Memory.get(COMMON.GEMINI_CHAT_HISTORY) as Content[];
@@ -184,83 +184,80 @@ export class GeminiLLM
    */
   private async processToolCalls(
     toolCalls: FunctionCall[],
-    userInput: string,
+    userInput: string
   ): Promise<any> {
     let finalResult: string | undefined = undefined;
-    let userQuery = userInput;
-    let callCount = 0;
-
-    for (const functionCall of toolCalls) {
-      try {
-        const functionResult =
-          await this.handleSingleFunctionCall(functionCall);
-
-        if (functionResult === undefined) {
-          finalResult = await this.groqLLM.generateText(userInput);
-          if (finalResult) {
-            finalResult = await this.fallBackToGroq(userInput);
-            return finalResult;
-          }
-        }
-
-        userQuery = `Tool result: ${JSON.stringify(functionResult)}. What is your next step?`;
-
-        if (functionCall.name === "think") {
-          const thought = functionResult?.response.content;
-          this.initialThought = thought;
-          if (thought) {
-            this.orchestrator.publish("onStrategizing", thought);
-            this.planSteps = this.parseThought(thought);
-            if (this.planSteps?.length > 0) {
-              userQuery = `Based on the plan, please provide an answer`;
-            } else {
-              userQuery = userInput;
-              this.planSteps = [];
+    try {
+      let userQuery = userInput;
+      let callCount = 0;
+
+      for (const functionCall of toolCalls) {
+        try {
+          const functionResult =
+            await this.handleSingleFunctionCall(functionCall);
+
+          if (functionCall.name === "think") {
+            const thought = functionResult?.response.content;
+            this.initialThought = this.initialThought ?? thought;
+            if (thought) {
+              this.orchestrator.publish("onStrategizing", thought);
+              this.planSteps = this.parseThought(thought);
+              if (this.planSteps?.length > 0) {
+                userQuery = `Tool result: ${JSON.stringify(functionResult)} \n. Based on these plans, Plans: ${this.planSteps} from the tool result \n What is your next step?`;
+              } else {
+                userQuery = `Tool result: ${JSON.stringify(functionResult)}. What is your next step?`;
+                this.planSteps = [];
+              }
             }
           } else {
             userQuery = `Tool result: ${JSON.stringify(functionResult)}. What is your next step?`;
           }
-        }
-
-        await this.buildChatHistory(
-          userQuery,
-          functionCall.name,
-          functionResult,
-          undefined,
-          false,
-        );
 
-        const snapShot = this.createSnapShot({
-          lastQuery: userQuery,
-          lastCall: functionCall.name,
-          lastResult: functionResult,
-          currentStepIndex: this.currentStepIndex,
-          planSteps: this.planSteps,
-        });
-        Memory.set(COMMON.GEMINI_SNAPSHOT, snapShot);
-        callCount++;
-      } catch (error: any) {
-        console.error("Error processing function call", error);
-
-        const retry = await vscode.window.showErrorMessage(
-          `Function call failed: ${error.message}. Retry or abort?`,
-          "Retry",
-          "Abort",
-        );
-
-        if (retry === "Retry") {
-          continue; // Retry the current function call
-        } else {
-          finalResult = `Function call error: ${error.message}. Falling back to last response.`;
-          break; // Exit the loop and return the error result
+          await this.buildChatHistory(
+            userQuery,
+            functionCall.name,
+            functionResult,
+            undefined,
+            false
+          );
+
+          const snapShot = this.createSnapShot({
+            lastQuery: userQuery,
+            lastCall: functionCall.name,
+            lastResult: functionResult,
+            currentStepIndex: this.currentStepIndex,
+            planSteps: this.planSteps,
+          });
+          Memory.set(COMMON.GEMINI_SNAPSHOT, snapShot);
+          callCount++;
+        } catch (error: any) {
+          console.error("Error processing function call", error);
+          // Send this to the webview instead and let the user decide
+          const retry = await vscode.window.showErrorMessage(
+            `Function call failed: ${error.message}. Retry or abort?`,
+            "Retry",
+            "Abort"
+          );
+
+          if (retry === "Retry") {
+            continue; // Retry the current function call
+          } else {
+            finalResult = `Function call error: ${error.message}. Falling back to last response.`;
+            break; // Exit the loop and return the error result
+          }
         }
       }
+      return finalResult;
+    } catch (error) {
+      console.error("Error processing tool calls", error);
+      finalResult = await this.fallBackToGroq(
+        `User Input: ${userInput} \n Plans: ${this.initialThought ?? "Write production ready code to demonstrate your solution"}`
+      );
     }
-    return finalResult;
   }
 
   async processUserQuery(
-    userInput: string,
+    userInput: string
   ): Promise<string | GenerateContentResult | undefined> {
     let finalResult: string | GenerateContentResult | undefined;
     let userQuery = userInput;
@@ -285,8 +282,8 @@ export class GeminiLLM
         const timeoutPromise = new Promise((_, reject) =>
           setTimeout(
             () => reject(new Error("TImeout Exceeded")),
-            this.timeOutMs,
-          ),
+            this.timeOutMs
+          )
         );
         const responsePromise = await this.generateContentWithTools(userQuery);
         const result = (await Promise.race([
@@ -320,7 +317,7 @@ export class GeminiLLM
           this.lastFunctionCalls.add(currentCallSignatures);
           if (this.lastFunctionCalls.size > 10) {
             this.lastFunctionCalls = new Set(
-              [...this.lastFunctionCalls].slice(-10),
+              [...this.lastFunctionCalls].slice(-10)
             );
           }
           if (toolCalls && toolCalls.length > 0) {
@@ -357,7 +354,7 @@ export class GeminiLLM
       if (snapshot?.length > 0) {
         Memory.removeItems(
           COMMON.GEMINI_SNAPSHOT,
-          Memory.get(COMMON.GEMINI_SNAPSHOT).length,
+          Memory.get(COMMON.GEMINI_SNAPSHOT).length
         );
       }
 
@@ -369,15 +366,15 @@ export class GeminiLLM
       // );
       console.log("Error processing user query", error);
       finalResult = await this.fallBackToGroq(
-        `${userInput} \n ${this.initialThought ?? "Write production ready code to demonstrate your solution"}`,
+        `${userInput} \n ${this.initialThought ?? "Write production ready code to demonstrate your solution"}`
       );
       console.log("Model not responding at this time, please try again", error);
     }
   }
 
   private async handleSingleFunctionCall(
     functionCall: FunctionCall,
-    attempt: number = 0,
+    attempt: number = 0
   ): Promise<any> {
     const MAX_RETRIES = 3;
     const args = functionCall.args as Record<string, any>;
@@ -402,7 +399,7 @@ export class GeminiLLM
       if (attempt < MAX_RETRIES) {
         console.warn(
           `Retry attempt ${attempt + 1} for function ${name}`,
-          JSON.stringify({ error, args }),
+          JSON.stringify({ error, args })
         );
         return this.handleSingleFunctionCall(functionCall, attempt + 1);
       }
@@ -424,7 +421,7 @@ export class GeminiLLM
     functionCall?: any,
     functionResponse?: any,
     chat?: ChatSession,
-    isInitialQuery: boolean = false,
+    isInitialQuery: boolean = false
   ): Promise<Content[]> {
     // Check if it makes sense to kind of seperate agent and Edit Mode memory, when switching.
     let chatHistory: any = Memory.get(COMMON.GEMINI_CHAT_HISTORY) || [];
@@ -450,17 +447,17 @@ export class GeminiLLM
         Message.of({
           role: "model",
           parts: [{ functionCall }],
-        }),
+        })
       );
 
       const observationResult = await chat.sendMessage(
-        `Tool result: ${JSON.stringify(functionResponse)}`,
+        `Tool result: ${JSON.stringify(functionResponse)}`
       );
       chatHistory.push(
         Message.of({
           role: "user",
           parts: [{ text: observationResult.response.text() }],
-        }),
+        })
       );
     }
     if (chatHistory.length > 50) chatHistory = chatHistory.slice(-50);

src/llms/groq/groq.ts

@@ -26,6 +26,7 @@ export class GroqLLM extends BaseLLM<any> implements vscode.Disposable {
     return [1, 2];
   }
 
+  // TODO Implement function call, especially think for this model.
   async generateText(message: string): Promise<string> {
     try {
       const { temperature, top_p, stop } = GROQ_CONFIG;

webviewUi/src/components/webview.tsx

@@ -49,7 +49,7 @@ export const WebviewUI = () => {
   const css = getChatCss("tokyo night");
   updateStyles(css);
   const [selectedModel, setSelectedModel] = useState("Gemini");
-  const [selectedCodeBuddyMode, setSelectedCodeBuddyMode] = useState("Edit");
+  const [selectedCodeBuddyMode, setSelectedCodeBuddyMode] = useState("Ask");
   const [userInput, setUserInput] = useState("");
   const [messages, setMessages] = useState<Message[]>([]);
   const [isBotLoading, setIsBotLoading] = useState(false);

webviewUi/src/constants/constant.ts

@@ -7,5 +7,5 @@ export const modelOptions = [
 
 export const codeBuddyMode = [
   { value: "Agent", label: "Agent" },
-  { value: "Edit", label: "Edit" },
+  { value: "Ask", label: "Ask" },
 ];