Refactor AgentNode in graph.py to utilize CodeAgent instead of ToolCallingAgent, update prompt template loading to use code_agent.yaml, and enhance step callback logging for improved debugging.

graph.py

@@ -1,7 +1,7 @@
 import logging
-from typing import TypedDict
+from typing import Callable, List, Optional, TypedDict
 from langgraph.graph import StateGraph, END
-from smolagents import ToolCallingAgent, LiteLLMModel
+from smolagents import CodeAgent, ToolCallingAgent, LiteLLMModel
 from tools import tools
 import yaml
 import os
@@ -25,7 +25,8 @@ class AgentNode:
         # Load default prompt templates from local file
         current_dir = os.path.dirname(os.path.abspath(__file__))
         prompts_dir = os.path.join(current_dir, "prompts")
-        yaml_path = os.path.join(prompts_dir, "toolcalling_agent.yaml")
+        # yaml_path = os.path.join(prompts_dir, "toolcalling_agent.yaml")
+        yaml_path = os.path.join(prompts_dir, "code_agent.yaml")
         
         with open(yaml_path, 'r') as f:
             prompt_templates = yaml.safe_load(f)
@@ -55,11 +56,25 @@ class AgentNode:
             model_id="ollama/codellama",
         )
         
-        self.agent = ToolCallingAgent(
+        # self.agent = ToolCallingAgent(
+        #     max_steps=1,
+        #     model=self.model,
+        #     prompt_templates=prompt_templates,
+        #     tools=tools
+        # )
+
+        step_callbacks: Optional[List[Callable]] = [
+            lambda step: logger.info(f"Step {step.step_number} completed: {step.action_step.action}")
+        ]
+
+        self.agent = CodeAgent(
+            add_base_tools=True,
             max_steps=1,
             model=self.model,
             prompt_templates=prompt_templates,
-            tools=tools
+            step_callbacks=step_callbacks,
+            tools=tools,
+            verbosity_level=logging.DEBUG
         )
 
     def __call__(self, state: AgentState) -> AgentState: