Update agent.py

agent.py

@@ -25,6 +25,7 @@ llm = ChatOpenAI(
     model="moonshotai/kimi-k2:free", # Model must support function calling in OpenRouter
     temperature=1
 )
+
 python_tool = PythonAstREPLTool()
 search_tool = BraveSearch.from_api_key(
                     api_key=os.getenv("BRAVE_SEARCH_API"),
@@ -103,70 +104,25 @@ gaia_agent = builder.compile()  # converts my builder into a runnable agent by u
 class LangGraphAgent:
     def __init__(self):
         print("LangGraphAgent initialized.")
-        self.question_count = 0  # Track the number of questions processed
 
     def __call__(self, question: str) -> str:
-        # Determine which API key to use based on question count
-        # First 50% of questions use OPENROUTER_API_KEY, rest use OPENROUTER_API_KEY_1
-        api_key = os.getenv("OPENROUTER_API_KEY") if self.question_count % 2 == 0 else os.getenv("OPENROUTER_API_KEY_1")
-        
-        # Create a new LLM instance with the selected API key
-        current_llm = ChatOpenAI(
-            base_url="https://openrouter.ai/api/v1",
-            api_key=api_key,
-            model="moonshotai/kimi-k2:free",
-            temperature=1
-        )
-        
-        # Bind tools to the current LLM
-        current_llm_with_tools = current_llm.bind_tools(tools)
-        
-        # Increment question counter for next call
-        self.question_count += 1
-        
-        print(f"Running LangGraphAgent with input: {question[:150]}... (Using API key {self.question_count % 2 + 1})")
+        input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
+        print(f"Running LangGraphAgent with input: {question[:150]}...")
         
-        # Create a custom LLM node for this specific question
-        def custom_llm_call(state: MessagesState):
-            return {
-                "messages": [
-                    current_llm_with_tools.invoke(
-                        [SystemMessage(content=system_prompt)] + state["messages"]
-                    )
-                ]
-            }
-        
-        # Build a new workflow with the custom LLM
-        custom_builder = StateGraph(MessagesState)
-        custom_builder.add_node("llm_call", custom_llm_call)
-        custom_builder.add_node("environment", tool_node)
-        custom_builder.add_edge(START, "llm_call")
-        custom_builder.add_conditional_edges(
-            "llm_call",
-            should_continue,
-            {"Action": "environment", END: END}
-        )
-        custom_builder.add_edge("environment", "llm_call")
-        custom_agent = custom_builder.compile()
-        
-        # Prepare the initial state and config
-        input_state = {"messages": [HumanMessage(content=question)]}
+        # tracing configuration for LangSmith
         config = RunnableConfig(
             config={
                 "run_name": "GAIA Agent",
                 "tags": ["gaia", "langgraph", "agent"],
                 "metadata": {"user_input": question},
-                "recursion_limit": 30,
-                "tracing": True
+                "recursion_limit": 30
             }
         )
-        
-        # Run the agent
-        result = custom_agent.invoke(input_state, config)
+        result = gaia_agent.invoke(input_state, config)  # 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()
+            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