Update agent.py

agent.py

@@ -108,14 +108,30 @@ builder.add_edge("environment", "llm_call")  # after running the tools go back t
 
 gaia_agent = builder.compile()  # converts my builder into a runnable agent by using gaia_agent.invoke()
 
-# Call the Agent
-def agent_run(question: str):
-    input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
-    result = gaia_agent.invoke(input_state, {"recursion_limit": 30})  # prevents infinite looping when the LLM keeps calling tools over and over
-    final_response = result["messages"][-1].content
-    
-    try:
-        return final_response.split("FINAL ANSWER:")[-1].strip()  #  parse out only what's after "FINAL ANSWER:"
-    except Exception:
-        print("Could not split on 'FINAL ANSWER:'")
-        return final_response
+                        # Call the Agent
+                        #def agent_run(question: str):
+                            input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
+                            result = gaia_agent.invoke(input_state, {"recursion_limit": 30})  # prevents infinite looping when the LLM keeps calling tools over and over
+                            final_response = result["messages"][-1].content
+                            
+                            try:
+                                return final_response.split("FINAL ANSWER:")[-1].strip()  #  parse out only what's after "FINAL ANSWER:"
+                            except Exception:
+                                print("Could not split on 'FINAL ANSWER:'")
+                                return final_response
+
+class LangGraphAgent:
+    def __init__(self):
+        print("LangGraphAgent initialized.")
+
+    def __call__(self, question: str) -> str:
+        input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
+        print(f"Running LangGraphAgent with input: {question[:50]}...")
+        result = gaia_agent.invoke(input_state, {"recursion_limit": 30})  # prevents infinite looping when the LLM keeps calling tools over and over
+        final_response = result["messages"][-1].content
+        
+        try:
+            return final_response.split("FINAL ANSWER:")[-1].strip()  #  parse out only what's after "FINAL ANSWER:"
+        except Exception:
+            print("Could not split on 'FINAL ANSWER:'")
+            return final_response