Enhance AgentRunner and graph functionality by introducing memory management and improved state handling. Update __call__ method to support both question input and resuming from interrupts, while adding new memory-related fields to track context, actions, and success/error counts. Refactor step callback logic for better user interaction and state management.

agent.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import uuid
+from langgraph.types import Command
 
 from graph import agent_graph
 
@@ -27,34 +28,55 @@ class AgentRunner:
         logger.info("Initializing AgentRunner")
         self.graph = agent_graph
         self.last_state = None  # Store the last state for testing/debugging
+        self.thread_id = str(uuid.uuid4())  # Generate a unique thread_id for this runner
 
-    def __call__(self, question: str) -> str:
+    def __call__(self, input_data) -> str:
         """Process a question through the agent graph and return the answer.
 
         Args:
-            question: The question to process
+            input_data: Either a question string or a Command object for resuming
 
         Returns:
             str: The agent's response
         """
         try:
-            logger.info(f"Processing question: {question}")
-            initial_state = {
-                "question": question,
-                "messages": [],
-                "answer": None,
-                "step_logs": [],
-                "is_complete": False,  # Initialize is_complete
-                "step_count": 0,  # Initialize step_count
-            }
-
-            # Generate a unique thread_id for this interaction
-            thread_id = str(uuid.uuid4())
-            config = {"configurable": {"thread_id": thread_id}}
-
-            final_state = self.graph.invoke(initial_state, config)
-            self.last_state = final_state  # Store the final state
-            return final_state.get("answer", "No answer generated")
+            config = {"configurable": {"thread_id": self.thread_id}}
+            
+            if isinstance(input_data, str):
+                # Initial question
+                logger.info(f"Processing question: {input_data}")
+                initial_state = {
+                    "question": input_data,
+                    "messages": [],
+                    "answer": None,
+                    "step_logs": [],
+                    "is_complete": False,
+                    "step_count": 0,
+                    # Initialize new memory fields
+                    "context": {},
+                    "memory_buffer": [],
+                    "last_action": None,
+                    "action_history": [],
+                    "error_count": 0,
+                    "success_count": 0,
+                }
+                
+                # Use stream to get interrupt information
+                for chunk in self.graph.stream(initial_state, config):
+                    if isinstance(chunk, tuple) and len(chunk) > 0 and hasattr(chunk[0], '__interrupt__'):
+                        # If we hit an interrupt, resume with 'c'
+                        for result in self.graph.stream(Command(resume="c"), config):
+                            self.last_state = result
+                            return result.get("answer", "No answer generated")
+                    self.last_state = chunk
+                    return chunk.get("answer", "No answer generated")
+            else:
+                # Resuming from interrupt
+                logger.info("Resuming from interrupt")
+                for result in self.graph.stream(input_data, config):
+                    self.last_state = result
+                    return result.get("answer", "No answer generated")
+
         except Exception as e:
-            logger.error(f"Error processing question: {str(e)}")
+            logger.error(f"Error processing input: {str(e)}")
             raise

graph.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+from datetime import datetime
 from typing import Dict, List, Optional, TypedDict, Union
 
 import yaml
@@ -66,6 +67,13 @@ class AgentState(TypedDict):
     step_logs: List[Dict]
     is_complete: bool
     step_count: int
+    # Add memory-related fields
+    context: Dict[str, any]  # For storing contextual information
+    memory_buffer: List[Dict]  # For storing important information across steps
+    last_action: Optional[str]  # Track the last action taken
+    action_history: List[Dict]  # History of actions taken
+    error_count: int  # Track error frequency
+    success_count: int  # Track successful operations
 
 
 class AgentNode:
@@ -92,18 +100,43 @@ class AgentNode:
         # Log execution start
         logger.info("Starting agent execution")
 
-        # Run the agent
-        result = self.agent.run(state["question"])
-
-        # Log result
-        logger.info(f"Agent execution result type: {type(result)}")
-        logger.info(f"Agent execution result value: {result}")
+        try:
+            # Run the agent
+            result = self.agent.run(state["question"])
+
+            # Update memory-related fields
+            new_state = state.copy()
+            new_state["messages"].append(AIMessage(content=result))
+            new_state["answer"] = result
+            new_state["step_count"] += 1
+            new_state["last_action"] = "agent_response"
+            new_state["action_history"].append(
+                {
+                    "step": state["step_count"],
+                    "action": "agent_response",
+                    "result": result,
+                }
+            )
+            new_state["success_count"] += 1
+
+            # Store important information in memory buffer
+            if result:
+                new_state["memory_buffer"].append(
+                    {
+                        "step": state["step_count"],
+                        "content": result,
+                        "timestamp": datetime.now().isoformat(),
+                    }
+                )
 
-        # Update state
-        new_state = state.copy()
-        new_state["messages"].append(AIMessage(content=result))
-        new_state["answer"] = result
-        new_state["step_count"] += 1
+        except Exception as e:
+            logger.error(f"Error during agent execution: {str(e)}")
+            new_state = state.copy()
+            new_state["error_count"] += 1
+            new_state["action_history"].append(
+                {"step": state["step_count"], "action": "error", "error": str(e)}
+            )
+            raise
 
         # Log updated state
         logger.info("Updated state after processing:")
@@ -134,10 +167,11 @@ class StepCallbackNode:
         state["step_logs"].append(step_log)
 
         try:
-            # Use interrupt for user input
-            user_input = interrupt(
+            # 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
 
             if user_input.lower() == "q":
                 state["is_complete"] = True
@@ -146,12 +180,12 @@ class StepCallbackNode:
                 logger.info(f"Current step: {state['step_count']}")
                 logger.info(f"Question: {state['question']}")
                 logger.info(f"Current answer: {state['answer']}")
-                return self(state, config)  # Recursively call for new input
+                return state
             elif user_input.lower() == "c":
                 return state
             else:
                 logger.warning("Invalid input. Please use 'c', 'q', or 'i'.")
-                return self(state, config)  # Recursively call for new input
+                return state
 
         except Exception as e:
             logger.warning(f"Error during interrupt: {str(e)}")
@@ -169,10 +203,20 @@ def build_agent_graph(agent: AgentNode) -> StateGraph:
 
     # Add edges
     workflow.add_edge("agent", "callback")
+
+    # Add conditional edges for callback
+    def should_continue(state: AgentState) -> str:
+        """Determine the next node based on state."""
+        if state["is_complete"]:
+            return END
+        # If we have an answer and no errors, we're done
+        if state["answer"] and state["error_count"] == 0:
+            return END
+        # Otherwise continue to agent
+        return "agent"
+
     workflow.add_conditional_edges(
-        "callback",
-        lambda x: END if x["is_complete"] else "agent",
-        {True: END, False: "agent"},
+        "callback", should_continue, {END: END, "agent": "agent"}
     )
 
     # Set entry point

test_agent.py

@@ -1,6 +1,7 @@
 import logging
 
 import pytest
+from langgraph.types import Command
 
 from agent import AgentRunner