Refactor app.py and update import paths in test_agent.py to improve code organization. Introduce new files for agent configuration, graph definition, and tools, enhancing the overall structure and functionality of the agent system.

agent.py → api/runner.py

@@ -4,7 +4,7 @@ import uuid
 
 from langgraph.types import Command
 
-from graph import agent_graph
+from services.graph import agent_graph
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)  # Default to INFO level
@@ -86,32 +86,20 @@ class AgentRunner:
                 }
                 logger.info(f"Initial state: {initial_state}")
 
-                # Use stream to get interrupt information
+                # Use stream to get results
                 logger.info("Starting graph stream for initial question")
                 for chunk in self.graph.stream(initial_state, config):
                     logger.debug(f"Received chunk: {chunk}")
-
                     if isinstance(chunk, dict):
                         if "__interrupt__" in chunk:
                             logger.info("Detected interrupt in stream")
                             logger.info(f"Interrupt details: {chunk['__interrupt__']}")
-
-                            # If we hit an interrupt, resume with 'c'
-                            logger.info("Resuming with 'c' command")
-                            for result in self.graph.stream(
-                                Command(resume="c"), config
-                            ):
-                                logger.debug(f"Received resume result: {result}")
-                                if isinstance(result, dict):
-                                    answer = self._extract_answer(result)
-                                    if answer:
-                                        self.last_state = result
-                                        return answer
-                        else:
-                            answer = self._extract_answer(chunk)
-                            if answer:
-                                self.last_state = chunk
-                                return answer
+                            # Let the graph handle the interrupt naturally
+                            continue
+                        answer = self._extract_answer(chunk)
+                        if answer:
+                            self.last_state = chunk
+                            return answer
                     else:
                         logger.debug(f"Skipping chunk without answer: {chunk}")
             else:

app.py

@@ -3,7 +3,6 @@ import os
 import gradio as gr
 import pandas as pd
 import requests
-
 from agent import AgentRunner
 
 # (Keep Constants as is)

graph.py → services/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 configuration import Configuration
-from tools import tools
+from services.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.abspath(__file__))
+current_dir = os.path.dirname(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")
 
@@ -182,9 +182,7 @@ class StepCallbackNode:
                 logger.info(f"Current answer: {state['answer']}")
                 return state
             elif user_input.lower() == "c":
-                # If we have an answer, mark as complete
-                if state["answer"]:
-                    state["is_complete"] = True
+                # Continue without marking as complete
                 return state
             else:
                 logger.warning("Invalid input. Please use 'c', 'q', or 'i'.")
@@ -192,9 +190,7 @@ class StepCallbackNode:
 
         except Exception as e:
             logger.warning(f"Error during interrupt: {str(e)}")
-            # If we have an answer, mark as complete
-            if state["answer"]:
-                state["is_complete"] = True
+            # Continue without marking as complete
             return state
 
 
@@ -213,19 +209,19 @@ 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 an answer and it's complete, we're done
-        if state["answer"] and state["is_complete"]:
-            logger.info(f"Found complete answer: {state['answer']}")
-            return END
+        # If we have no answer, continue
+        if not state["answer"]:
+            logger.info("No answer found, continuing")
+            return "agent"
 
         # If we have an answer but it's not complete, continue
-        if state["answer"]:
+        if not state["is_complete"]:
             logger.info(f"Found answer but not complete: {state['answer']}")
             return "agent"
 
-        # If we have no answer, continue
-        logger.info("No answer found, continuing")
-        return "agent"
+        # If we have an answer and it's complete, we're done
+        logger.info(f"Found complete answer: {state['answer']}")
+        return END
 
     workflow.add_conditional_edges(
         "callback", should_continue, {END: END, "agent": "agent"}

test_agent.py

@@ -2,7 +2,7 @@ import logging
 
 import pytest
 
-from agent import AgentRunner
+from api.runner import AgentRunner
 
 # Configure test logger
 test_logger = logging.getLogger("test_agent")
@@ -194,9 +194,29 @@ def test_simple_math_calculation_with_steps():
 
     # Verify final answer
     expected_result = 1302.678
+
+    # Extract all numbers from the response
+    import re
+
+    # First check for LaTeX formatting
+    latex_match = re.search(r"\\boxed{([^}]+)}", response)
+    if latex_match:
+        # Extract number from LaTeX box
+        latex_content = latex_match.group(1)
+        numbers = re.findall(r"\d+\.?\d*", latex_content)
+    else:
+        # Extract all numbers from the response
+        numbers = re.findall(r"\d+\.?\d*", response)
+
+    assert numbers, "Response should contain at least one number"
+
+    # Check if any number matches the expected result
+    has_correct_result = any(abs(float(n) - expected_result) < 0.001 for n in numbers)
     assert (
-        str(expected_result) in response
-    ), f"Response should contain the result {expected_result}"
+        has_correct_result
+    ), f"Response should contain the result {expected_result}, got {response}"
+
+    # Verify the response indicates it's a final answer
     assert (
         "final_answer" in response.lower()
     ), "Response should indicate it's using final_answer"