feat: Enhance tool call handling in AICodeAgentCog with TaskComplete support and improved response parsing

This commit is contained in:
Slipstream 2025-05-31 16:00:32 -06:00
parent b11a974b64
commit ea93b85fc5
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD

View File

@ -42,6 +42,7 @@ AGENT_SYSTEM_PROMPT = """You are an expert AI Coding Agent. Your primary functio
**Inline Tool Call Syntax:**
When you need to use a tool, your response should *only* contain the tool call block, formatted exactly as specified below. The system will parse this, execute the tool, and then feed the output back to you in a subsequent message prefixed with "ToolResponse:".
IMPORTANT: Do NOT wrap your tool calls in markdown code blocks (e.g., ```tool ... ``` or ```json ... ```). Output the raw tool syntax directly, starting with the tool name (e.g., `ReadFile:`).
1. **ReadFile:** Reads the content of a specified file.
```tool
@ -90,6 +91,13 @@ When you need to use a tool, your response should *only* contain the tool call b
```
(System will provide search results or error in ToolResponse)
7. **TaskComplete:** Signals that the current multi-step task is considered complete by the AI.
```tool
TaskComplete:
message: <string: A brief summary of what was accomplished or the final status.>
```
(System will acknowledge and stop the current interaction loop.)
**Workflow and Rules:**
- **Tool Preference:** For modifying existing files, ALWAYS prefer `ApplyDiff` if the changes are targeted. Use `WriteFile` for new files or if `ApplyDiff` is unsuitable or fails repeatedly.
- **Direct Operation:** You operate directly. No explicit user confirmation is needed for individual tool actions after the initial user prompt.
@ -355,19 +363,28 @@ class AICodeAgentCog(commands.Cog):
async def _parse_and_execute_tool_call(self, ctx: commands.Context, ai_response_text: str) -> Optional[str]:
"""
Parses AI response for an inline tool call, executes it, and returns the tool's output string.
Returns None if no tool call is found.
Returns a tuple: (status: str, data: Optional[str]).
Status can be "TOOL_OUTPUT", "TASK_COMPLETE", "NO_TOOL".
Data is the tool output string, completion message, or original AI text.
"""
tool_output_str = None
tool_executed = False # Flag to indicate if any tool was matched and attempted
# Order of checks might matter if syntax is ambiguous, but designed to be distinct.
# --- ReadFile ---
read_file_match = re.search(r"^\s*ReadFile:\s*path:\s*(.+?)\s*$", ai_response_text, re.IGNORECASE | re.MULTILINE)
if read_file_match:
# --- TaskComplete ---
# Check for TaskComplete first as it's a terminal operation for the loop.
task_complete_match = re.search(r"^\s*TaskComplete:\s*message:\s*(.*)\s*$", ai_response_text, re.IGNORECASE | re.DOTALL)
if task_complete_match:
tool_executed = True
file_path = read_file_match.group(1).strip()
tool_output = await self._execute_tool_read_file(file_path)
tool_output_str = f"ToolResponse: ReadFile\nPath: {file_path}\n---\n{tool_output}"
completion_message = task_complete_match.group(1).strip()
return "TASK_COMPLETE", completion_message
# --- ReadFile ---
if not tool_executed:
read_file_match = re.search(r"^\s*ReadFile:\s*path:\s*(.+?)\s*$", ai_response_text, re.IGNORECASE | re.MULTILINE)
if read_file_match:
tool_executed = True
file_path = read_file_match.group(1).strip()
tool_output = await self._execute_tool_read_file(file_path)
return "TOOL_OUTPUT", f"ToolResponse: ReadFile\nPath: {file_path}\n---\n{tool_output}"
# --- WriteFile ---
if not tool_executed:
@ -375,15 +392,15 @@ class AICodeAgentCog(commands.Cog):
if write_file_match:
tool_executed = True
file_path = write_file_match.group(1).strip()
content = write_file_match.group(2).strip() # Ensure .strip() is appropriate for all content
content = write_file_match.group(2).strip()
snapshot_branch = await self._create_programmatic_snapshot()
if not snapshot_branch:
tool_output_str = "ToolResponse: SystemError\n---\nFailed to create project snapshot. WriteFile operation aborted."
return "TOOL_OUTPUT", "ToolResponse: SystemError\n---\nFailed to create project snapshot. WriteFile operation aborted."
else:
await ctx.send(f"AICodeAgent: Created snapshot: {snapshot_branch} before writing to {file_path}")
tool_output = await self._execute_tool_write_file(file_path, content)
tool_output_str = f"ToolResponse: WriteFile\nPath: {file_path}\n---\n{tool_output}"
return "TOOL_OUTPUT", f"ToolResponse: WriteFile\nPath: {file_path}\n---\n{tool_output}"
# --- ApplyDiff ---
if not tool_executed:
@ -395,11 +412,11 @@ class AICodeAgentCog(commands.Cog):
snapshot_branch = await self._create_programmatic_snapshot()
if not snapshot_branch:
tool_output_str = "ToolResponse: SystemError\n---\nFailed to create project snapshot. ApplyDiff operation aborted."
return "TOOL_OUTPUT", "ToolResponse: SystemError\n---\nFailed to create project snapshot. ApplyDiff operation aborted."
else:
await ctx.send(f"AICodeAgent: Created snapshot: {snapshot_branch} before applying diff to {file_path}")
tool_output = await self._execute_tool_apply_diff(file_path, diff_block)
tool_output_str = f"ToolResponse: ApplyDiff\nPath: {file_path}\n---\n{tool_output}"
return "TOOL_OUTPUT", f"ToolResponse: ApplyDiff\nPath: {file_path}\n---\n{tool_output}"
# --- ExecuteCommand ---
if not tool_executed:
@ -408,7 +425,7 @@ class AICodeAgentCog(commands.Cog):
tool_executed = True
command_str = exec_command_match.group(1).strip()
tool_output = await self._execute_tool_execute_command(command_str)
tool_output_str = f"ToolResponse: ExecuteCommand\nCommand: {command_str}\n---\n{tool_output}"
return "TOOL_OUTPUT", f"ToolResponse: ExecuteCommand\nCommand: {command_str}\n---\n{tool_output}"
# --- ListFiles ---
if not tool_executed:
@ -419,7 +436,7 @@ class AICodeAgentCog(commands.Cog):
recursive_str = list_files_match.group(2)
recursive = recursive_str.lower() == 'true' if recursive_str else False
tool_output = await self._execute_tool_list_files(file_path, recursive)
tool_output_str = f"ToolResponse: ListFiles\nPath: {file_path}\nRecursive: {recursive}\n---\n{tool_output}"
return "TOOL_OUTPUT", f"ToolResponse: ListFiles\nPath: {file_path}\nRecursive: {recursive}\n---\n{tool_output}"
# --- WebSearch ---
if not tool_executed:
@ -428,12 +445,12 @@ class AICodeAgentCog(commands.Cog):
tool_executed = True
query_str = web_search_match.group(1).strip()
tool_output = await self._execute_tool_web_search(query_str)
tool_output_str = f"ToolResponse: WebSearch\nQuery: {query_str}\n---\n{tool_output}"
return "TOOL_OUTPUT", f"ToolResponse: WebSearch\nQuery: {query_str}\n---\n{tool_output}"
return tool_output_str
return "NO_TOOL", ai_response_text # No tool call found, return original AI text
# --- Placeholder Tool Execution Methods ---
# These will be properly implemented later. For now, they return placeholder strings.
# --- Tool Execution Methods ---
# (Implementations for _execute_tool_... methods remain the same)
async def _execute_tool_read_file(self, path: str) -> str:
print(f"AICodeAgentCog: Placeholder _execute_tool_read_file for path: {path}")
@ -679,22 +696,44 @@ class AICodeAgentCog(commands.Cog):
print(f"AICodeAgentCog: AI Raw Response:\n{ai_response_text}")
# Parse for inline tool call
tool_output_str = await self._parse_and_execute_tool_call(ctx, ai_response_text)
# _parse_and_execute_tool_call now returns -> Tuple[str, Optional[str]]
# status can be "TOOL_OUTPUT", "TASK_COMPLETE", "NO_TOOL"
# data is the tool output string, completion message, or original AI text
parse_status, parsed_data = await self._parse_and_execute_tool_call(ctx, ai_response_text)
if tool_output_str:
if parse_status == "TASK_COMPLETE":
completion_message = parsed_data if parsed_data is not None else "Task marked as complete by AI."
await ctx.send(f"AICodeAgent: Task Complete!\n{completion_message}")
# Log AI's completion signal to history (optional, but good for context)
# self._add_to_conversation_history(user_id, role="model", text_content=f"TaskComplete: message: {completion_message}")
return # End of interaction
elif parse_status == "TOOL_OUTPUT":
tool_output_str = parsed_data
if tool_output_str is None: # Should not happen if status is TOOL_OUTPUT but defensive
tool_output_str = "Error: Tool executed but returned no output string."
print(f"AICodeAgentCog: Tool Output:\n{tool_output_str}")
self._add_to_conversation_history(user_id, role="user", text_content=tool_output_str) # Feed tool output back as 'user'
iteration_count += 1
# Potentially send tool output to Discord for transparency if desired
# await ctx.send(f"```\n{tool_output_str}\n```")
# Optionally send tool output to Discord for transparency if desired
# if len(tool_output_str) < 1900 : await ctx.send(f"```{tool_output_str}```")
continue # Loop back to AI with tool output in history
else:
elif parse_status == "NO_TOOL":
# No tool call found, this is the final AI response for this turn
if len(ai_response_text) > 1950:
await ctx.send(ai_response_text[:1950] + "\n...(message truncated)")
final_ai_text = parsed_data # This is the original ai_response_text
if final_ai_text is None: # Should not happen
final_ai_text = "AI provided no textual response."
if len(final_ai_text) > 1950:
await ctx.send(final_ai_text[:1950] + "\n...(message truncated)")
else:
await ctx.send(ai_response_text)
await ctx.send(final_ai_text)
return # End of interaction
else: # Should not happen
await ctx.send("AICodeAgent: Internal error - unknown parse status from tool parser.")
return
except google_exceptions.GoogleAPICallError as e:
await ctx.send(f"AICodeAgent: Vertex AI API call failed: {e}")