diff --git a/agent-claudecode.py b/agent-claudecode.py index eeb7fbf..48842b4 100644 --- a/agent-claudecode.py +++ b/agent-claudecode.py @@ -99,7 +99,10 @@ def plan(task): response_format={"type": "json_object"} ) try: - plan_data = json.loads(response.choices[0].message.content) + plan_content = response.choices[0].message.content + if not isinstance(plan_content, str): + return "Error: Failed to create plan" + plan_data = json.loads(plan_content) steps = plan_data.get("steps", [task]) current_plan = steps print(f"[Plan] Created {len(steps)} steps") @@ -204,19 +207,21 @@ def run_agent_step(messages, tools, max_iterations=5): if "_argument_error" in function_args: function_response = f"Error: {function_args['_argument_error']}" elif function_name == "plan" and function_impl is not None: - plan_mode = True function_response = function_impl(**function_args) messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": function_response}) if current_plan: results = [] - for i, step in enumerate(current_plan, 1): - print(f"\n[Step {i}/{len(current_plan)}] {step}") - messages.append({"role": "user", "content": step}) - result, messages = run_agent_step(messages, [t for t in tools if t["function"]["name"] != "plan"]) - results.append(result) - print(f"\n{result}") - plan_mode = False - current_plan = [] + plan_mode = True + try: + for i, step in enumerate(current_plan, 1): + print(f"\n[Step {i}/{len(current_plan)}] {step}") + messages.append({"role": "user", "content": step}) + result, messages = run_agent_step(messages, [t for t in tools if t["function"]["name"] != "plan"]) + results.append(result) + print(f"\n{result}") + finally: + plan_mode = False + current_plan = [] return "\n".join(results), messages elif function_impl is not None: function_response = function_impl(**function_args) @@ -246,18 +251,24 @@ def run_agent_claudecode(task, use_plan=False): context_parts.append(f"\n# Previous Context\n{memory}") messages = [{"role": "system", "content": "\n".join(context_parts)}] if use_plan: - plan_mode = True - plan(task) - results = [] - for i, step in enumerate(current_plan, 1): - print(f"\n[Step {i}/{len(current_plan)}] {step}") - messages.append({"role": "user", "content": step}) - result, messages = run_agent_step(messages, [t for t in all_tools if t["function"]["name"] != "plan"]) - results.append(result) - print(f"\n{result}") - plan_mode = False - current_plan = [] - final_result = "\n".join(results) + plan_result = plan(task) + if current_plan: + results = [] + plan_mode = True + try: + for i, step in enumerate(current_plan, 1): + print(f"\n[Step {i}/{len(current_plan)}] {step}") + messages.append({"role": "user", "content": step}) + result, messages = run_agent_step(messages, [t for t in all_tools if t["function"]["name"] != "plan"]) + results.append(result) + print(f"\n{result}") + finally: + plan_mode = False + current_plan = [] + final_result = "\n".join(results) + else: + final_result = plan_result + print(f"\n{final_result}") else: messages.append({"role": "user", "content": task}) final_result, messages = run_agent_step(messages, all_tools) diff --git a/agent.py b/agent.py index 8a6341f..19740eb 100644 --- a/agent.py +++ b/agent.py @@ -1,6 +1,7 @@ import os import json import subprocess +from typing import Any from openai import OpenAI client = OpenAI( @@ -70,6 +71,16 @@ def write_file(path, content): functions = {"execute_bash": execute_bash, "read_file": read_file, "write_file": write_file} +def parse_tool_arguments(raw_arguments: str) -> dict[str, Any]: + if not raw_arguments: + return {} + try: + parsed = json.loads(raw_arguments) + return parsed if isinstance(parsed, dict) else {} + except json.JSONDecodeError as error: + return {"_argument_error": f"Invalid JSON arguments: {error}"} + + def run_agent(user_message, max_iterations=5): messages = [ {"role": "system", "content": "You are a helpful assistant. Be concise."}, @@ -86,11 +97,16 @@ def run_agent(user_message, max_iterations=5): if not message.tool_calls: return message.content for tool_call in message.tool_calls: - name = tool_call.function.name - args = json.loads(tool_call.function.arguments) + function_payload = getattr(tool_call, "function", None) + if function_payload is None: + continue + name = str(getattr(function_payload, "name", "")) + args = parse_tool_arguments(str(getattr(function_payload, "arguments", ""))) print(f"[Tool] {name}({args})") if name not in functions: result = f"Error: Unknown tool '{name}'" + elif "_argument_error" in args: + result = f"Error: {args['_argument_error']}" else: result = functions[name](**args) messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": result}) diff --git a/tests/test_agent.py b/tests/test_agent.py index 2d8a188..f8e1265 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -326,6 +326,33 @@ def fake_create(*, model, messages, tools): self.assertEqual(len(tool_messages), 1) self.assertIn("Invalid JSON arguments", tool_messages[0]["content"]) + def test_run_agent_claudecode_executes_plan_steps(self): + responses = [ + make_response( + SimpleNamespace(content='{"steps": ["step one", "step two"]}') + ), + make_response(SimpleNamespace(content="done 1", tool_calls=[])), + make_response(SimpleNamespace(content="done 2", tool_calls=[])), + ] + + def fake_create(*, model, messages, tools=None, response_format=None): + return responses.pop(0) + + setattr( + self.agent, + "client", + SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=fake_create)) + ), + ) + setattr(self.agent, "save_memory", lambda task, result: None) + + result = self.agent.run_agent_claudecode("test plan", use_plan=True) + + self.assertEqual(result, "done 1\ndone 2") + self.assertEqual(self.agent.current_plan, []) + self.assertFalse(self.agent.plan_mode) + if __name__ == "__main__": unittest.main()