Add configuration, graph, runner, and tools modules to enhance agent functionality. Introduce a Configuration class for managing parameters, implement an AgentRunner for executing the agent graph, and create tools for general search and mathematical calculations. Update test_agent.py to reflect new import paths and improve overall code organization.

services/graph.py → graph.py

@@ -6,14 +6,14 @@ from datetime import datetime
 from typing import Dict, List, Optional, TypedDict, Union
 
 import yaml
-from services.configuration import Configuration
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
 from langgraph.graph import END, StateGraph
 from langgraph.types import interrupt
 from smolagents import CodeAgent, LiteLLMModel
 
-from services.tools import tools
+from configuration import Configuration
+from tools import tools
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -33,7 +33,7 @@ else:
 litellm.drop_params = True
 
 # Load default prompt templates from local file
-current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__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, "code_agent.yaml")
 
@@ -150,48 +150,50 @@ class AgentNode:
 class StepCallbackNode:
     """Node that handles step callbacks and user interaction."""
 
-    def __call__(
-        self, state: AgentState, config: Optional[RunnableConfig] = None
-    ) -> AgentState:
-        """Handle step callback and user interaction."""
-        # Get configuration
-        cfg = Configuration.from_runnable_config(config)
+    def __init__(self, name: str):
+        self.name = name
 
-        # Log the step
-        step_log = {
-            "step": state["step_count"],
-            "messages": [msg.content for msg in state["messages"]],
-            "question": state["question"],
-            "answer": state["answer"],
-        }
-        state["step_logs"].append(step_log)
+    def __call__(self, state: dict) -> dict:
+        """Process the state and handle user interaction."""
+        print(f"\nCurrent step: {state.get('step_count', 0)}")
+        print(f"Question: {state.get('question', 'No question')}")
+        print(f"Current answer: {state.get('answer', 'No answer yet')}\n")
 
-        try:
-            # Use interrupt for user input and unpack the tuple
-            interrupt_result = interrupt(
-                "Press 'c' to continue, 'q' to quit, or 'i' for more info: "
-            )
-            user_input = interrupt_result[0]  # Get the actual user input
+        while True:
+            choice = input(
+                "Enter 'c' to continue, 'q' to quit, 'i' for more info, or 'r' to reject answer: "
+            ).lower()
 
-            if user_input.lower() == "q":
+            if choice == "c":
+                # Mark as complete to continue
                 state["is_complete"] = True
                 return state
-            elif user_input.lower() == "i":
-                logger.info(f"Current step: {state['step_count']}")
-                logger.info(f"Question: {state['question']}")
-                logger.info(f"Current answer: {state['answer']}")
+            elif choice == "q":
+                # Mark as complete and set answer to None to quit
+                state["is_complete"] = True
+                state["answer"] = None
                 return state
-            elif user_input.lower() == "c":
-                # Continue without marking as complete
+            elif choice == "i":
+                # Show more information but don't mark as complete
+                print("\nAdditional Information:")
+                print(f"Messages: {state.get('messages', [])}")
+                print(f"Step Logs: {state.get('step_logs', [])}")
+                print(f"Context: {state.get('context', {})}")
+                print(f"Memory Buffer: {state.get('memory_buffer', [])}")
+                print(f"Last Action: {state.get('last_action', None)}")
+                print(f"Action History: {state.get('action_history', [])}")
+                print(f"Error Count: {state.get('error_count', 0)}")
+                print(f"Success Count: {state.get('success_count', 0)}\n")
+            elif choice == "r":
+                # Reject the current answer and continue execution
+                print("\nRejecting current answer and continuing execution...")
+                # Clear the message history to prevent confusion
+                state["messages"] = []
+                state["answer"] = None
+                state["is_complete"] = False
                 return state
             else:
-                logger.warning("Invalid input. Please use 'c', 'q', or 'i'.")
-                return state
-
-        except Exception as e:
-            logger.warning(f"Error during interrupt: {str(e)}")
-            # Continue without marking as complete
-            return state
+                print("Invalid choice. Please enter 'c', 'q', 'i', or 'r'.")
 
 
 def build_agent_graph(agent: AgentNode) -> StateGraph:
@@ -201,7 +203,7 @@ def build_agent_graph(agent: AgentNode) -> StateGraph:
 
     # Add nodes
     workflow.add_node("agent", agent)
-    workflow.add_node("callback", StepCallbackNode())
+    workflow.add_node("callback", StepCallbackNode("callback"))
 
     # Add edges
     workflow.add_edge("agent", "callback")
@@ -209,22 +211,24 @@ def build_agent_graph(agent: AgentNode) -> StateGraph:
     # Add conditional edges for callback
     def should_continue(state: AgentState) -> str:
         """Determine the next node based on state."""
-        # If we have no answer, continue
+        # If we have no answer, continue to agent
         if not state["answer"]:
-            logger.info("No answer found, continuing")
-            return "agent"
-
-        # If we have an answer but it's not complete, continue
-        if not state["is_complete"]:
-            logger.info(f"Found answer but not complete: {state['answer']}")
+            logger.info("No answer found, continuing to agent")
             return "agent"
 
         # If we have an answer and it's complete, we're done
-        logger.info(f"Found complete answer: {state['answer']}")
-        return END
+        if state["is_complete"]:
+            logger.info(f"Found complete answer: {state['answer']}")
+            return END
+
+        # Otherwise, go to callback for user input
+        logger.info(f"Waiting for user input for answer: {state['answer']}")
+        return "callback"
 
     workflow.add_conditional_edges(
-        "callback", should_continue, {END: END, "agent": "agent"}
+        "callback",
+        should_continue,
+        {END: END, "agent": "agent", "callback": "callback"},
     )
 
     # Set entry point

api/runner.py → runner.py

@@ -1,10 +1,11 @@
 import logging
 import os
+import re
 import uuid
 
 from langgraph.types import Command
 
-from services.graph import agent_graph
+from graph import agent_graph
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)  # Default to INFO level
@@ -48,6 +49,26 @@ class AgentRunner:
         if "messages" in state and state["messages"]:
             for msg in reversed(state["messages"]):
                 if hasattr(msg, "content") and msg.content:
+                    # Look for code blocks that might contain the answer
+                    if "```" in msg.content:
+                        # Extract code between ```py and ``` or ```python and ```
+                        code_match = re.search(
+                            r"```(?:py|python)?\s*\n(.*?)\n```", msg.content, re.DOTALL
+                        )
+                        if code_match:
+                            code = code_match.group(1)
+                            # Look for final_answer call
+                            final_answer_match = re.search(
+                                r"final_answer\((.*?)\)", code
+                            )
+                            if final_answer_match:
+                                answer = final_answer_match.group(1)
+                                logger.info(
+                                    f"Found answer in final_answer call: {answer}"
+                                )
+                                return answer
+
+                    # If no code block with final_answer, use the content
                     logger.info(f"Found answer in message: {msg.content}")
                     return msg.content
 
@@ -99,7 +120,9 @@ class AgentRunner:
                         answer = self._extract_answer(chunk)
                         if answer:
                             self.last_state = chunk
-                            return answer
+                            # If the state is complete, return the answer
+                            if chunk.get("is_complete", False):
+                                return answer
                     else:
                         logger.debug(f"Skipping chunk without answer: {chunk}")
             else:
@@ -111,7 +134,9 @@ class AgentRunner:
                         answer = self._extract_answer(result)
                         if answer:
                             self.last_state = result
-                            return answer
+                            # If the state is complete, return the answer
+                            if result.get("is_complete", False):
+                                return answer
                     else:
                         logger.debug(f"Skipping result without answer: {result}")
 
@@ -122,3 +147,34 @@ class AgentRunner:
         except Exception as e:
             logger.error(f"Error processing input: {str(e)}")
             raise
+
+
+if __name__ == "__main__":
+    import argparse
+
+    from langgraph.types import Command
+
+    # Set up argument parser
+    parser = argparse.ArgumentParser(description="Run the agent with a question")
+    parser.add_argument("question", type=str, help="The question to ask the agent")
+    parser.add_argument(
+        "--resume",
+        type=str,
+        help="Value to resume with after an interrupt",
+        default=None,
+    )
+    args = parser.parse_args()
+
+    # Create agent runner
+    runner = AgentRunner()
+
+    if args.resume:
+        # Resume from interrupt with provided value
+        print(f"\nResuming with value: {args.resume}")
+        response = runner(Command(resume=args.resume))
+    else:
+        # Initial run with question
+        print(f"\nAsking question: {args.question}")
+        response = runner(args.question)
+
+    print(f"\nFinal response: {response}")

test_agent.py

@@ -2,7 +2,7 @@ import logging
 
 import pytest
 
-from api.runner import AgentRunner
+from runner import AgentRunner
 
 # Configure test logger
 test_logger = logging.getLogger("test_agent")

services/tools.py → tools.py

@@ -47,9 +47,39 @@ class GeneralSearchTool(Tool):
         return "\n\n---\n\n".join(output)
 
 
+class MathTool(Tool):
+    name = "math"
+    description = """Performs mathematical calculations and returns the result."""
+    inputs = {
+        "expression": {
+            "type": "string",
+            "description": "The mathematical expression to evaluate.",
+        }
+    }
+    output_type = "string"
+
+    def forward(self, expression: str) -> str:
+        try:
+            # Use eval with a restricted set of builtins for safety
+            safe_dict = {
+                "__builtins__": {
+                    "abs": abs,
+                    "round": round,
+                    "min": min,
+                    "max": max,
+                    "sum": sum,
+                }
+            }
+            result = eval(expression, safe_dict)
+            return str(result)
+        except Exception as e:
+            raise Exception(f"Error evaluating expression: {str(e)}")
+
+
 # Export all tools
 tools = [
     # DuckDuckGoSearchTool(),
     GeneralSearchTool(),
+    MathTool(),
     # WikipediaSearchTool(),
 ]