Update agent.py

agent.py

@@ -19,15 +19,8 @@ from langchain_core.runnables import RunnableConfig  # for LangSmith tracking
 langsmith_api_key = os.getenv("LANGSMITH_API_KEY")
 langsmith_tracing = os.getenv("LANGSMITH_TRACING")
 
-# Add a global counter to track the number of questions answered
-question_counter = 0
-
-# Modify get_llm to accept a key parameter
-def get_llm(api_key=None):
-    if api_key is None:
-        api_keys = [os.getenv("OPENROUTER_API_KEY"), os.getenv("OPENROUTER_API_KEY_1")]
-    else:
-        api_keys = [api_key]
+def get_llm():
+    api_keys = [os.getenv("OPENROUTER_API_KEY"), os.getenv("OPENROUTER_API_KEY_1")]
     last_exception = None
     for key in api_keys:
         if not key:
@@ -39,52 +32,15 @@ def get_llm(api_key=None):
                 model="qwen/qwen3-coder:free",
                 temperature=1
             )
+            # Optionally, test the key with a trivial call to ensure it's valid
+            # llm.invoke([SystemMessage(content="ping")])
             return llm
         except Exception as e:
             last_exception = e
             continue
     raise RuntimeError(f"All OpenRouter API keys failed: {last_exception}")
 
-# Remove the global llm instance
-# llm = get_llm()
-
-# In the LangGraphAgent class, select the key based on the counter
-class LangGraphAgent:
-    def __init__(self):
-        print("LangGraphAgent initialized.")
-        self.counter = 0
-        self.total = 0  # Set this to the total number of GAIA questions
-
-    def __call__(self, question: str) -> str:
-        # Decide which key to use
-        if self.total == 0:
-            self.total = 100  # Replace with actual total if known
-        halfway = self.total // 2
-        if self.counter < halfway:
-            api_key = os.getenv("OPENROUTER_API_KEY")
-        else:
-            api_key = os.getenv("OPENROUTER_API_KEY_1")
-        llm = get_llm(api_key)
-        llm_with_tools = llm.bind_tools(tools)
-        input_state = {"messages": [HumanMessage(content=question)]}
-        print(f"Running LangGraphAgent with input: {question[:150]}...")
-        config = RunnableConfig(
-            config={
-                "run_name": "GAIA Agent",
-                "tags": ["gaia", "langgraph", "agent"],
-                "metadata": {"user_input": question},
-                "recursion_limit": 30,
-                "tracing": True
-            }
-        )
-        result = gaia_agent.invoke(input_state, config)
-        final_response = result["messages"][-1].content
-        self.counter += 1
-        try:
-            return final_response.split("FINAL ANSWER:")[-1].strip()
-        except Exception:
-            print("Could not split on 'FINAL ANSWER:'")
-            return final_response
+llm = get_llm()
 python_tool = PythonAstREPLTool()
 search_tool = BraveSearch.from_api_key(
                     api_key=os.getenv("BRAVE_SEARCH_API"),
@@ -163,12 +119,54 @@ 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:
-        input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
-        print(f"Running LangGraphAgent with input: {question[:150]}...")
+        # 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="qwen/qwen3-coder: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
         
-        # tracing configuration for LangSmith
+        print(f"Running LangGraphAgent with input: {question[:150]}... (Using API key {self.question_count % 2 + 1})")
+        
+        # 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)]}
         config = RunnableConfig(
             config={
                 "run_name": "GAIA Agent",
@@ -178,11 +176,13 @@ class LangGraphAgent:
                 "tracing": True
             }
         )
-        result = gaia_agent.invoke(input_state, config)  # prevents infinite looping when the LLM keeps calling tools over and over
+        
+        # Run the agent
+        result = custom_agent.invoke(input_state, config)
         final_response = result["messages"][-1].content
         
         try:
-            return final_response.split("FINAL ANSWER:")[-1].strip()  #  parse out only what's after "FINAL ANSWER:"
+            return final_response.split("FINAL ANSWER:")[-1].strip()
         except Exception:
             print("Could not split on 'FINAL ANSWER:'")
             return final_response