Update agent.py

agent.py

@@ -12,7 +12,6 @@ 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()
 
@@ -105,54 +104,16 @@ class LangGraphAgent:
         self.graph = build_graph()
         print("LangGraphAgent initialized with tools.")
         
-    def __call__(self, question: str, task_id: str = None) -> dict:
-        """Run the agent on a question and return structured answer with reasoning trace"""
+    def __call__(self, question: str) -> str:
+        """Run the agent on a question and return the answer"""
         try:
             messages = [HumanMessage(content=question)]
             result = self.graph.invoke({"messages": messages})
-            
-            # Print all messages for debugging
             for m in result["messages"]:
                 m.pretty_print()
-            
-            # 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
-            }
-            
+            return result["messages"][-1].content
         except Exception as e:
-            return {
-                "task_id": task_id or "unknown",
-                "model_answer": f"Error: {str(e)}",
-                "reasoning_trace": f"Error occurred during processing: {str(e)}"
-            }
+            return f"Error: {str(e)}"
 
 if __name__ == "__main__":
     agent = LangGraphAgent()