From aec01bffdb5b652a3127969aeb691dbbc9a70990 Mon Sep 17 00:00:00 2001 From: Ismail Pelaseyed Date: Mon, 21 Jul 2025 10:20:39 +0200 Subject: [PATCH] increase max tool execution --- src/agent/grok-agent.ts | 278 +++++++++++++++++++++++----------------- 1 file changed, 157 insertions(+), 121 deletions(-) diff --git a/src/agent/grok-agent.ts b/src/agent/grok-agent.ts index 5009066..81676a6 100644 --- a/src/agent/grok-agent.ts +++ b/src/agent/grok-agent.ts @@ -1,13 +1,13 @@ -import { GrokClient, GrokMessage, GrokToolCall } from '../grok/client'; -import { GROK_TOOLS } from '../grok/tools'; -import { TextEditorTool, BashTool, TodoTool, ConfirmationTool } from '../tools'; -import { ToolResult } from '../types'; -import { EventEmitter } from 'events'; -import { createTokenCounter, TokenCounter } from '../utils/token-counter'; -import { loadCustomInstructions } from '../utils/custom-instructions'; +import { GrokClient, GrokMessage, GrokToolCall } from "../grok/client"; +import { GROK_TOOLS } from "../grok/tools"; +import { TextEditorTool, BashTool, TodoTool, ConfirmationTool } from "../tools"; +import { ToolResult } from "../types"; +import { EventEmitter } from "events"; +import { createTokenCounter, TokenCounter } from "../utils/token-counter"; +import { loadCustomInstructions } from "../utils/custom-instructions"; export interface ChatEntry { - type: 'user' | 'assistant' | 'tool_result'; + type: "user" | "assistant" | "tool_result"; content: string; timestamp: Date; toolCalls?: GrokToolCall[]; @@ -17,7 +17,7 @@ export interface ChatEntry { } export interface StreamingChunk { - type: 'content' | 'tool_calls' | 'tool_result' | 'done' | 'token_count'; + type: "content" | "tool_calls" | "tool_result" | "done" | "token_count"; content?: string; toolCalls?: GrokToolCall[]; toolCall?: GrokToolCall; @@ -42,17 +42,17 @@ export class GrokAgent extends EventEmitter { this.bash = new BashTool(); this.todoTool = new TodoTool(); this.confirmationTool = new ConfirmationTool(); - this.tokenCounter = createTokenCounter('grok-4-latest'); - + this.tokenCounter = createTokenCounter("grok-4-latest"); + // Load custom instructions const customInstructions = loadCustomInstructions(); - const customInstructionsSection = customInstructions + const customInstructionsSection = customInstructions ? `\n\nCUSTOM INSTRUCTIONS:\n${customInstructions}\n\nThe above custom instructions should be followed alongside the standard instructions below.` - : ''; - + : ""; + // Initialize with system message this.messages.push({ - role: 'system', + role: "system", content: `You are Grok CLI, an AI assistant that helps with file editing, coding tasks, and system operations.${customInstructionsSection} You have access to these tools: @@ -105,89 +105,107 @@ IMPORTANT RESPONSE GUIDELINES: - Keep responses concise and focused on the actual work being done - If a tool execution completes the user's request, you can remain silent or give a brief confirmation -Current working directory: ${process.cwd()}` +Current working directory: ${process.cwd()}`, }); } async processUserMessage(message: string): Promise { // Add user message to conversation const userEntry: ChatEntry = { - type: 'user', + type: "user", content: message, - timestamp: new Date() + timestamp: new Date(), }; this.chatHistory.push(userEntry); - this.messages.push({ role: 'user', content: message }); + this.messages.push({ role: "user", content: message }); const newEntries: ChatEntry[] = [userEntry]; const maxToolRounds = 10; // Prevent infinite loops let toolRounds = 0; try { - let currentResponse = await this.grokClient.chat(this.messages, GROK_TOOLS); - + let currentResponse = await this.grokClient.chat( + this.messages, + GROK_TOOLS + ); + // Agent loop - continue until no more tool calls or max rounds reached while (toolRounds < maxToolRounds) { const assistantMessage = currentResponse.choices[0]?.message; if (!assistantMessage) { - throw new Error('No response from Grok'); + throw new Error("No response from Grok"); } // Handle tool calls - if (assistantMessage.tool_calls && assistantMessage.tool_calls.length > 0) { + if ( + assistantMessage.tool_calls && + assistantMessage.tool_calls.length > 0 + ) { toolRounds++; - + // Add assistant message with tool calls const assistantEntry: ChatEntry = { - type: 'assistant', - content: assistantMessage.content || 'Using tools to help you...', + type: "assistant", + content: assistantMessage.content || "Using tools to help you...", timestamp: new Date(), - toolCalls: assistantMessage.tool_calls + toolCalls: assistantMessage.tool_calls, }; this.chatHistory.push(assistantEntry); newEntries.push(assistantEntry); // Add assistant message to conversation this.messages.push({ - role: 'assistant', - content: assistantMessage.content || '', - tool_calls: assistantMessage.tool_calls + role: "assistant", + content: assistantMessage.content || "", + tool_calls: assistantMessage.tool_calls, } as any); // Execute tool calls for (const toolCall of assistantMessage.tool_calls) { const result = await this.executeTool(toolCall); - + const toolResultEntry: ChatEntry = { - type: 'tool_result', - content: result.success ? result.output || 'Success' : result.error || 'Error occurred', + type: "tool_result", + content: result.success + ? result.output || "Success" + : result.error || "Error occurred", timestamp: new Date(), toolCall: toolCall, - toolResult: result + toolResult: result, }; this.chatHistory.push(toolResultEntry); newEntries.push(toolResultEntry); // Add tool result to messages with proper format (needed for AI context) this.messages.push({ - role: 'tool', - content: result.success ? (result.output || 'Success') : (result.error || 'Error'), - tool_call_id: toolCall.id + role: "tool", + content: result.success + ? result.output || "Success" + : result.error || "Error", + tool_call_id: toolCall.id, }); } // Get next response - this might contain more tool calls - currentResponse = await this.grokClient.chat(this.messages, GROK_TOOLS); + currentResponse = await this.grokClient.chat( + this.messages, + GROK_TOOLS + ); } else { // No more tool calls, add final response const finalEntry: ChatEntry = { - type: 'assistant', - content: assistantMessage.content || 'I understand, but I don\'t have a specific response.', - timestamp: new Date() + type: "assistant", + content: + assistantMessage.content || + "I understand, but I don't have a specific response.", + timestamp: new Date(), }; this.chatHistory.push(finalEntry); - this.messages.push({ role: 'assistant', content: assistantMessage.content || '' }); + this.messages.push({ + role: "assistant", + content: assistantMessage.content || "", + }); newEntries.push(finalEntry); break; // Exit the loop } @@ -195,9 +213,10 @@ Current working directory: ${process.cwd()}` if (toolRounds >= maxToolRounds) { const warningEntry: ChatEntry = { - type: 'assistant', - content: 'Maximum tool execution rounds reached. Stopping to prevent infinite loops.', - timestamp: new Date() + type: "assistant", + content: + "Maximum tool execution rounds reached. Stopping to prevent infinite loops.", + timestamp: new Date(), }; this.chatHistory.push(warningEntry); newEntries.push(warningEntry); @@ -206,9 +225,9 @@ Current working directory: ${process.cwd()}` return newEntries; } catch (error: any) { const errorEntry: ChatEntry = { - type: 'assistant', + type: "assistant", content: `Sorry, I encountered an error: ${error.message}`, - timestamp: new Date() + timestamp: new Date(), }; this.chatHistory.push(errorEntry); return [userEntry, errorEntry]; @@ -227,7 +246,7 @@ Current working directory: ${process.cwd()}` delete arr.index; } } - } else if (typeof acc[key] === 'string' && typeof value === 'string') { + } else if (typeof acc[key] === "string" && typeof value === "string") { (acc[key] as string) += value; } else if (Array.isArray(acc[key]) && Array.isArray(value)) { const accArray = acc[key] as any[]; @@ -235,34 +254,38 @@ Current working directory: ${process.cwd()}` if (!accArray[i]) accArray[i] = {}; accArray[i] = reduce(accArray[i], value[i]); } - } else if (typeof acc[key] === 'object' && typeof value === 'object') { + } else if (typeof acc[key] === "object" && typeof value === "object") { acc[key] = reduce(acc[key], value); } } return acc; }; - + return reduce(previous, item.choices[0]?.delta || {}); } - async *processUserMessageStream(message: string): AsyncGenerator { + async *processUserMessageStream( + message: string + ): AsyncGenerator { // Add user message to conversation const userEntry: ChatEntry = { - type: 'user', + type: "user", content: message, - timestamp: new Date() + timestamp: new Date(), }; this.chatHistory.push(userEntry); - this.messages.push({ role: 'user', content: message }); + this.messages.push({ role: "user", content: message }); // Calculate input tokens - const inputTokens = this.tokenCounter.countMessageTokens(this.messages as any); + const inputTokens = this.tokenCounter.countMessageTokens( + this.messages as any + ); yield { - type: 'token_count', - tokenCount: inputTokens + type: "token_count", + tokenCount: inputTokens, }; - const maxToolRounds = 10; // Prevent infinite loops + const maxToolRounds = 30; // Prevent infinite loops let toolRounds = 0; let totalOutputTokens = 0; @@ -272,101 +295,108 @@ Current working directory: ${process.cwd()}` // Stream response and accumulate const stream = this.grokClient.chatStream(this.messages, GROK_TOOLS); let accumulatedMessage: any = {}; - let accumulatedContent = ''; + let accumulatedContent = ""; let toolCallsYielded = false; - + for await (const chunk of stream) { if (!chunk.choices?.[0]) continue; - + // Accumulate the message using reducer accumulatedMessage = this.messageReducer(accumulatedMessage, chunk); - + // Check for tool calls - yield when we have complete tool calls with function names if (!toolCallsYielded && accumulatedMessage.tool_calls?.length > 0) { // Check if we have at least one complete tool call with a function name - const hasCompleteTool = accumulatedMessage.tool_calls.some((tc: any) => tc.function?.name); + const hasCompleteTool = accumulatedMessage.tool_calls.some( + (tc: any) => tc.function?.name + ); if (hasCompleteTool) { yield { - type: 'tool_calls', - toolCalls: accumulatedMessage.tool_calls + type: "tool_calls", + toolCalls: accumulatedMessage.tool_calls, }; toolCallsYielded = true; } } - + // Stream content as it comes if (chunk.choices[0].delta?.content) { accumulatedContent += chunk.choices[0].delta.content; - + // Update token count in real-time - const currentOutputTokens = this.tokenCounter.estimateStreamingTokens(accumulatedContent); + const currentOutputTokens = + this.tokenCounter.estimateStreamingTokens(accumulatedContent); totalOutputTokens = currentOutputTokens; - + yield { - type: 'content', - content: chunk.choices[0].delta.content + type: "content", + content: chunk.choices[0].delta.content, }; - + // Emit token count update yield { - type: 'token_count', - tokenCount: inputTokens + totalOutputTokens + type: "token_count", + tokenCount: inputTokens + totalOutputTokens, }; } } // Add assistant entry to history const assistantEntry: ChatEntry = { - type: 'assistant', - content: accumulatedMessage.content || 'Using tools to help you...', + type: "assistant", + content: accumulatedMessage.content || "Using tools to help you...", timestamp: new Date(), - toolCalls: accumulatedMessage.tool_calls || undefined + toolCalls: accumulatedMessage.tool_calls || undefined, }; this.chatHistory.push(assistantEntry); // Add accumulated message to conversation this.messages.push({ - role: 'assistant', - content: accumulatedMessage.content || '', - tool_calls: accumulatedMessage.tool_calls + role: "assistant", + content: accumulatedMessage.content || "", + tool_calls: accumulatedMessage.tool_calls, } as any); // Handle tool calls if present if (accumulatedMessage.tool_calls?.length > 0) { toolRounds++; - + // Only yield tool_calls if we haven't already yielded them during streaming if (!toolCallsYielded) { yield { - type: 'tool_calls', - toolCalls: accumulatedMessage.tool_calls + type: "tool_calls", + toolCalls: accumulatedMessage.tool_calls, }; } // Execute tools for (const toolCall of accumulatedMessage.tool_calls) { const result = await this.executeTool(toolCall); - + const toolResultEntry: ChatEntry = { - type: 'tool_result', - content: result.success ? result.output || 'Success' : result.error || 'Error occurred', + type: "tool_result", + content: result.success + ? result.output || "Success" + : result.error || "Error occurred", timestamp: new Date(), toolCall: toolCall, - toolResult: result + toolResult: result, }; this.chatHistory.push(toolResultEntry); - + yield { - type: 'tool_result', + type: "tool_result", toolCall, - toolResult: result + toolResult: result, }; // Add tool result with proper format (needed for AI context) this.messages.push({ - role: 'tool', - content: result.success ? (result.output || 'Success') : (result.error || 'Error'), - tool_call_id: toolCall.id + role: "tool", + content: result.success + ? result.output || "Success" + : result.error || "Error", + tool_call_id: toolCall.id, }); } @@ -379,63 +409,69 @@ Current working directory: ${process.cwd()}` if (toolRounds >= maxToolRounds) { yield { - type: 'content', - content: '\n\nMaximum tool execution rounds reached. Stopping to prevent infinite loops.' + type: "content", + content: + "\n\nMaximum tool execution rounds reached. Stopping to prevent infinite loops.", }; } - yield { type: 'done' }; + yield { type: "done" }; } catch (error: any) { const errorEntry: ChatEntry = { - type: 'assistant', + type: "assistant", content: `Sorry, I encountered an error: ${error.message}`, - timestamp: new Date() + timestamp: new Date(), }; this.chatHistory.push(errorEntry); yield { - type: 'content', - content: errorEntry.content + type: "content", + content: errorEntry.content, }; - yield { type: 'done' }; + yield { type: "done" }; } } - private async executeTool(toolCall: GrokToolCall): Promise { try { const args = JSON.parse(toolCall.function.arguments); - + switch (toolCall.function.name) { - case 'view_file': - const range: [number, number] | undefined = args.start_line && args.end_line ? [args.start_line, args.end_line] : undefined; + case "view_file": + const range: [number, number] | undefined = + args.start_line && args.end_line + ? [args.start_line, args.end_line] + : undefined; return await this.textEditor.view(args.path, range); - - case 'create_file': + + case "create_file": return await this.textEditor.create(args.path, args.content); - - case 'str_replace_editor': - return await this.textEditor.strReplace(args.path, args.old_str, args.new_str); - - case 'bash': + + case "str_replace_editor": + return await this.textEditor.strReplace( + args.path, + args.old_str, + args.new_str + ); + + case "bash": return await this.bash.execute(args.command); - - case 'create_todo_list': + + case "create_todo_list": return await this.todoTool.createTodoList(args.todos); - - case 'update_todo_list': + + case "update_todo_list": return await this.todoTool.updateTodoList(args.updates); - - + default: return { success: false, - error: `Unknown tool: ${toolCall.function.name}` + error: `Unknown tool: ${toolCall.function.name}`, }; } } catch (error: any) { return { success: false, - error: `Tool execution error: ${error.message}` + error: `Tool execution error: ${error.message}`, }; } } @@ -462,4 +498,4 @@ Current working directory: ${process.cwd()}` this.tokenCounter.dispose(); this.tokenCounter = createTokenCounter(model); } -} \ No newline at end of file +}