Enhance AgentRunner and graph functionality by adding answer extraction logic and improving logging throughout the processing flow. Update the handling of interrupts and state management to ensure clarity in debug output. Refactor the should_continue function in graph.py to better manage completion states and improve user interaction.

agent.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import uuid
+
 from langgraph.types import Command
 
 from graph import agent_graph
@@ -28,7 +29,29 @@ 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
+        self.thread_id = str(
+            uuid.uuid4()
+        )  # Generate a unique thread_id for this runner
+        logger.info(f"Created AgentRunner with thread_id: {self.thread_id}")
+
+    def _extract_answer(self, state: dict) -> str:
+        """Extract the answer from the state."""
+        if not state:
+            return None
+
+        # First try to get answer from direct answer field
+        if "answer" in state and state["answer"]:
+            logger.info(f"Found answer in direct field: {state['answer']}")
+            return state["answer"]
+
+        # Then try to get answer from messages
+        if "messages" in state and state["messages"]:
+            for msg in reversed(state["messages"]):
+                if hasattr(msg, "content") and msg.content:
+                    logger.info(f"Found answer in message: {msg.content}")
+                    return msg.content
+
+        return None
 
     def __call__(self, input_data) -> str:
         """Process a question through the agent graph and return the answer.
@@ -41,10 +64,11 @@ class AgentRunner:
         """
         try:
             config = {"configurable": {"thread_id": self.thread_id}}
-            
+            logger.info(f"Using config: {config}")
+
             if isinstance(input_data, str):
                 # Initial question
-                logger.info(f"Processing question: {input_data}")
+                logger.info(f"Processing initial question: {input_data}")
                 initial_state = {
                     "question": input_data,
                     "messages": [],
@@ -60,22 +84,52 @@ class AgentRunner:
                     "error_count": 0,
                     "success_count": 0,
                 }
-                
+                logger.info(f"Initial state: {initial_state}")
+
                 # Use stream to get interrupt information
+                logger.info("Starting graph stream for initial question")
                 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")
+                    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
+                    else:
+                        logger.debug(f"Skipping chunk without answer: {chunk}")
             else:
                 # Resuming from interrupt
-                logger.info("Resuming from interrupt")
+                logger.info(f"Resuming from interrupt with input: {input_data}")
                 for result in self.graph.stream(input_data, config):
-                    self.last_state = result
-                    return result.get("answer", "No answer generated")
+                    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:
+                        logger.debug(f"Skipping result without answer: {result}")
+
+            # If we get here, we didn't find an answer
+            logger.warning("No answer generated from stream")
+            return "No answer generated"
 
         except Exception as e:
             logger.error(f"Error processing input: {str(e)}")

graph.py

@@ -182,6 +182,9 @@ 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
                 return state
             else:
                 logger.warning("Invalid input. Please use 'c', 'q', or 'i'.")
@@ -189,6 +192,9 @@ 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
             return state
 
 
@@ -207,12 +213,18 @@ 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 state["is_complete"]:
-            return END
-        # If we have an answer and no errors, we're done
-        if state["answer"] and state["error_count"] == 0:
+        # 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
-        # Otherwise continue to agent
+
+        # If we have an answer but it's not complete, continue
+        if state["answer"]:
+            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"
 
     workflow.add_conditional_edges(

test_agent.py

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