Update agent.py

agent.py

@@ -1,92 +1,110 @@
 import os
-import itertools
-from typing import TypedDict, Annotated, Literal
 
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage
 from langgraph.graph.message import add_messages
-from langgraph.graph import MessagesState, StateGraph, START, END
-from langchain_core.runnables import RunnableConfig  # for LangSmith tracking
+from langgraph.graph import MessagesState
+from langgraph.graph import StateGraph, START, END
+from typing import TypedDict, Annotated, Literal
 
 from langchain_community.tools import BraveSearch  # web search
 from langchain_experimental.tools.python.tool import PythonAstREPLTool  # for logic/math problems
 
-from tools import (
-    calculator_basic, datetime_tools, transcribe_audio,
-    transcribe_youtube, query_image, webpage_content, read_excel
-)
+from tools import (calculator_basic, datetime_tools, transcribe_audio, transcribe_youtube, query_image, webpage_content, read_excel)
 from prompt import system_prompt
 
-# --------------------------------------------------------------------
-# 1. API Key Rotation Setup
-# --------------------------------------------------------------------
-api_keys = [k for k in [
-    os.getenv("OPENROUTER_API_KEY"),
-    os.getenv("OPENROUTER_API_KEY_1")
-] if k]
-
-if not api_keys:
-    raise EnvironmentError("No OpenRouter API keys found in environment variables.")
-
-api_key_cycle = itertools.cycle(api_keys)
-
-def get_next_api_key() -> str:
-    """Get the next API key in rotation."""
-    return next(api_key_cycle)
-
-class RotatingChatOpenAI(ChatOpenAI):
-    """ChatOpenAI wrapper that automatically rotates OpenRouter API keys on failure."""
-
-    def invoke(self, *args, **kwargs):
-        for _ in range(len(api_keys)):  # try each key once per call
-            current_key = get_next_api_key()
-            self.openai_api_key = current_key  # ✅ Correct for ChatOpenAI
-            try:
-                return super().invoke(*args, **kwargs)
-            except Exception as e:
-                # Handle rate-limits or auth errors
-                if any(code in str(e) for code in ["429", "401", "403"]):
-                    print(f"[API Key Rotation] Key {current_key[:5]}... failed, trying next key...")
-                    continue
-                raise  # Re-raise unexpected errors
-        raise RuntimeError("All OpenRouter API keys failed or were rate-limited.")
-
-# --------------------------------------------------------------------
-# 2. Initialize LLM with API Key Rotation
-# --------------------------------------------------------------------
-llm = RotatingChatOpenAI(
-    base_url="https://openrouter.ai/api/v1",
-    openai_api_key=get_next_api_key(),  # ✅ start with the first key
-    model="qwen/qwen3-coder:free",      # must support tool/function calling
-    temperature=1
-)
+from langchain_core.runnables import RunnableConfig  # for LangSmith tracking
+
+# LangSmith to observe the agent
+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]
+    last_exception = None
+    for key in api_keys:
+        if not key:
+            continue
+        try:
+            llm = ChatOpenAI(
+                base_url="https://openrouter.ai/api/v1",
+                api_key=key,
+                model="qwen/qwen3-coder:free",
+                temperature=1
+            )
+            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()
 
-# --------------------------------------------------------------------
-# 3. Tools Setup
-# --------------------------------------------------------------------
+# 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
 python_tool = PythonAstREPLTool()
 search_tool = BraveSearch.from_api_key(
-    api_key=os.getenv("BRAVE_SEARCH_API"),
-    search_kwargs={"count": 4},
-    description="Web search using Brave"
+                    api_key=os.getenv("BRAVE_SEARCH_API"),
+                    search_kwargs={"count": 4},  # returns the 4 best results and their URL
+                    description="Web search using Brave"
 )
 
 community_tools = [search_tool, python_tool]
-custom_tools = calculator_basic + datetime_tools + [
-    transcribe_audio, transcribe_youtube, query_image, webpage_content, read_excel
-]
+custom_tools = calculator_basic + datetime_tools + [transcribe_audio, transcribe_youtube, query_image, webpage_content, read_excel]
 
 tools = community_tools + custom_tools
 llm_with_tools = llm.bind_tools(tools)
+
+# Prepare tools by name
 tools_by_name = {tool.name: tool for tool in tools}
 
-# --------------------------------------------------------------------
-# 4. Define LangGraph State and Nodes
-# --------------------------------------------------------------------
-class MessagesState(TypedDict):
+class MessagesState(TypedDict):  # creates the state (is like the agent's memory at any moment)
     messages: Annotated[list[AnyMessage], add_messages]
 
-# LLM Node
+# LLM node
 def llm_call(state: MessagesState):
     return {
         "messages": [
@@ -96,64 +114,61 @@ def llm_call(state: MessagesState):
         ]
     }
 
-# Tool Node
+# Tool node
 def tool_node(state: MessagesState):
-    """Executes tools requested by the LLM."""
-    results = []
-    last_message = state["messages"][-1]
-
-    for tool_call in getattr(last_message, "tool_calls", []) or []:
-        tool = tools_by_name.get(tool_call["name"])
-        if not tool:
-            results.append(ToolMessage(content=f"Unknown tool: {tool_call['name']}", tool_call_id=tool_call["id"]))
-            continue
-
-        args = tool_call["args"]
-        # Handle dict vs positional args safely
-        try:
-            observation = tool.invoke(**args) if isinstance(args, dict) else tool.invoke(args)
-        except Exception as e:
-            observation = f"[Tool Error] {str(e)}"
-
-        results.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
-
-    return {"messages": results}
-
-# Conditional Routing
+    """Executes the tools"""
+    
+    result = []
+    for tool_call in state["messages"][-1].tool_calls:   # gives a list of the tools the LLM decided to call
+        tool = tools_by_name[tool_call["name"]]   # look up the actual tool function using a dictionary
+        observation = tool.invoke(tool_call["args"])   # executes the tool
+        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))   # the result from the tool is added to the memory
+    return {"messages": result}   # thanks to add_messages, LangGraph will automatically append the result to the agent's message history
+
+# Conditional edge function to route to the tool node or end based upon whether the LLM made a tool call
 def should_continue(state: MessagesState) -> Literal["Action", END]:
-    """Route to tools if LLM made a tool call, else end."""
-    last_message = state["messages"][-1]
-    return "Action" if getattr(last_message, "tool_calls", None) else END
+    """Decide if we should continue the loop or stop based upon whether the LLM made a tool call"""
+
+    last_message = state["messages"][-1]  # looks at the last message (usually from the LLM)
+    
+    # If the LLM makes a tool call, then perform an action
+    if last_message.tool_calls:
+        return "Action"
+    # Otherwise, we stop (reply to the user)
+    return END
 
-# --------------------------------------------------------------------
-# 5. Build LangGraph Agent
-# --------------------------------------------------------------------
+# Build workflow
 builder = StateGraph(MessagesState)
 
+# Add nodes
 builder.add_node("llm_call", llm_call)
 builder.add_node("environment", tool_node)
 
+# Add edges to connect nodes
 builder.add_edge(START, "llm_call")
 builder.add_conditional_edges(
     "llm_call",
     should_continue,
-    {"Action": "environment", END: END}
+    {"Action": "environment",  # name returned by should_continue : Name of the next node
+     END: END}
 )
-builder.add_edge("environment", "llm_call")
+    # If tool calls -> "Action" -> environment (executes the tool)
+    # If no tool calls -> END
 
-gaia_agent = builder.compile()
+builder.add_edge("environment", "llm_call")  # after running the tools go back to the LLM for another round of reasoning
 
-# --------------------------------------------------------------------
-# 6. Agent Wrapper
-# --------------------------------------------------------------------
+gaia_agent = builder.compile()  # converts my builder into a runnable agent by using gaia_agent.invoke()
+
+# Wrapper class to initialize and call the LangGraph agent with a user question
 class LangGraphAgent:
     def __init__(self):
-        print("LangGraphAgent initialized with API key rotation.")
+        print("LangGraphAgent initialized.")
 
     def __call__(self, question: str) -> str:
-        input_state = {"messages": [HumanMessage(content=question)]}
+        input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
         print(f"Running LangGraphAgent with input: {question[:150]}...")
-
+        
+        # tracing configuration for LangSmith
         config = RunnableConfig(
             config={
                 "run_name": "GAIA Agent",
@@ -163,12 +178,11 @@ class LangGraphAgent:
                 "tracing": True
             }
         )
-
-        result = gaia_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
-
-        # Extract "FINAL ANSWER" if present
-        if isinstance(final_response, str):
-            return final_response.split("FINAL ANSWER:")[-1].strip()
-        else:
-            return str(final_response)
+        
+        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