Update agent.py

agent.py

@@ -12,6 +12,7 @@ from langchain_core.tools import tool
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_tavily import TavilySearch
+import json
 
 load_dotenv()
 
@@ -104,16 +105,54 @@ class LangGraphAgent:
         self.graph = build_graph()
         print("LangGraphAgent initialized with tools.")
         
-    def __call__(self, question: str) -> str:
-        """Run the agent on a question and return the answer"""
+    def __call__(self, question: str, task_id: str = None) -> dict:
+        """Run the agent on a question and return structured answer with reasoning trace"""
         try:
             messages = [HumanMessage(content=question)]
             result = self.graph.invoke({"messages": messages})
+            
+            # Print all messages for debugging
             for m in result["messages"]:
                 m.pretty_print()
-            return result["messages"][-1].content
+            
+            # Extract the final answer and build reasoning trace
+            final_answer = result["messages"][-1].content
+            
+            # Build reasoning trace from all messages
+            reasoning_steps = []
+            for i, msg in enumerate(result["messages"]):
+                if isinstance(msg, SystemMessage):
+                    reasoning_steps.append(f"Step {i+1}: System prompt loaded for ReAct methodology")
+                elif isinstance(msg, HumanMessage):
+                    reasoning_steps.append(f"Step {i+1}: Received question: {msg.content}")
+                elif isinstance(msg, AIMessage):
+                    if msg.tool_calls:
+                        for tool_call in msg.tool_calls:
+                            reasoning_steps.append(f"Step {i+1}: Called tool '{tool_call['name']}' with args: {tool_call['args']}")
+                    else:
+                        reasoning_steps.append(f"Step {i+1}: AI reasoning: {msg.content}")
+                else:
+                    reasoning_steps.append(f"Step {i+1}: Tool response: {str(msg.content)[:200]}...")
+            
+            reasoning_trace = " | ".join(reasoning_steps)
+            
+            # Extract the final answer from the AI response
+            model_answer = final_answer
+            if "FINAL ANSWER:" in final_answer:
+                model_answer = final_answer.split("FINAL ANSWER:")[-1].strip()
+            
+            return {
+                "task_id": task_id or "unknown",
+                "model_answer": model_answer,
+                "reasoning_trace": reasoning_trace
+            }
+            
         except Exception as e:
-            return f"Error: {str(e)}"
+            return {
+                "task_id": task_id or "unknown",
+                "model_answer": f"Error: {str(e)}",
+                "reasoning_trace": f"Error occurred during processing: {str(e)}"
+            }
 
 if __name__ == "__main__":
     agent = LangGraphAgent()