Update agent.py

agent.py

@@ -1,51 +1,92 @@
 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
-from langgraph.graph import StateGraph, START, END
-from typing import TypedDict, Annotated, Literal
+from langgraph.graph import MessagesState, StateGraph, START, END
+from langchain_core.runnables import RunnableConfig  # for LangSmith tracking
 
 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
 
-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")
-
-llm = ChatOpenAI(
+# --------------------------------------------------------------------
+# 1. API Key Rotation Setup
+# --------------------------------------------------------------------
+api_keys = [
+    os.getenv("OPENROUTER_API_KEY"),
+    os.getenv("OPENROUTER_API_KEY_1")
+]
+
+if not any(api_keys):
+    raise EnvironmentError("No OpenRouter API keys found in environment variables.")
+
+api_key_cycle = itertools.cycle([k for k in api_keys if k])
+
+def get_next_api_key():
+    """Get the next API key in rotation."""
+    return next(api_key_cycle)
+
+class RotatingChatOpenAI(ChatOpenAI):
+    """ChatOpenAI wrapper that automatically rotates API keys on failure."""
+
+    def invoke(self, *args, **kwargs):
+        # Try each key once per call
+        for _ in range(len(api_keys)):
+            self.api_key = get_next_api_key()
+            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 {self.api_key[:5]}... failed, trying next key...")
+                    continue
+                raise  # Re-raise other unexpected errors
+        raise RuntimeError("All OpenRouter API keys failed or rate-limited.")
+
+# --------------------------------------------------------------------
+# 2. Initialize LLM with API Key Rotation
+# --------------------------------------------------------------------
+llm = RotatingChatOpenAI(
     base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENROUTER_API_KEY"),
-    model="qwen/qwen3-coder:free", # Model must support function calling in OpenRouter
+    api_key=get_next_api_key(),  # Start with the first key
+    model="qwen/qwen3-coder:free",  # Model must support function calling
     temperature=1
 )
 
+# --------------------------------------------------------------------
+# 3. Tools Setup
+# --------------------------------------------------------------------
 python_tool = PythonAstREPLTool()
 search_tool = BraveSearch.from_api_key(
-                    api_key=os.getenv("BRAVE_SEARCH_API"),
-                    search_kwargs={"count": 4},  # returns the 4 best results and their URL
-                    description="Web search using Brave"
+    api_key=os.getenv("BRAVE_SEARCH_API"),
+    search_kwargs={"count": 4},
+    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}
 
-class MessagesState(TypedDict):  # creates the state (is like the agent's memory at any moment)
+# --------------------------------------------------------------------
+# 4. Define LangGraph State and Nodes
+# --------------------------------------------------------------------
+class MessagesState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
-# LLM node
+# LLM Node
 def llm_call(state: MessagesState):
     return {
         "messages": [
@@ -55,61 +96,51 @@ def llm_call(state: MessagesState):
         ]
     }
 
-# Tool node
+# Tool Node
 def tool_node(state: MessagesState):
-    """Executes the tools"""
-    
+    """Executes tools requested by the LLM."""
     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
+    for tool_call in state["messages"][-1].tool_calls:
+        tool = tools_by_name[tool_call["name"]]
+        observation = tool.invoke(tool_call["args"])
+        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
+    return {"messages": result}
 
-# Conditional edge function to route to the tool node or end based upon whether the LLM made a tool call
+# Conditional Routing
 def should_continue(state: MessagesState) -> Literal["Action", 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
+    """Route to tools if LLM made a tool call, else end."""
+    last_message = state["messages"][-1]
+    return "Action" if last_message.tool_calls else END
 
-# Build workflow
+# --------------------------------------------------------------------
+# 5. Build LangGraph Agent
+# --------------------------------------------------------------------
 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",  # name returned by should_continue : Name of the next node
-     END: END}
+    {"Action": "environment", END: END}
 )
-    # If tool calls -> "Action" -> environment (executes the tool)
-    # If no tool calls -> END
-
-builder.add_edge("environment", "llm_call")  # after running the tools go back to the LLM for another round of reasoning
+builder.add_edge("environment", "llm_call")
 
-gaia_agent = builder.compile()  # converts my builder into a runnable agent by using gaia_agent.invoke()
+gaia_agent = builder.compile()
 
-# Wrapper class to initialize and call the LangGraph agent with a user question
+# --------------------------------------------------------------------
+# 6. Agent Wrapper
+# --------------------------------------------------------------------
 class LangGraphAgent:
     def __init__(self):
-        print("LangGraphAgent initialized.")
+        print("LangGraphAgent initialized with API key rotation.")
 
     def __call__(self, question: str) -> str:
-        input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
+        input_state = {"messages": [HumanMessage(content=question)]}
         print(f"Running LangGraphAgent with input: {question[:150]}...")
-        
-        # tracing configuration for LangSmith
+
         config = RunnableConfig(
             config={
                 "run_name": "GAIA Agent",
@@ -119,11 +150,12 @@ class LangGraphAgent:
                 "tracing": True
             }
         )
-        result = gaia_agent.invoke(input_state, config)  # prevents infinite looping when the LLM keeps calling tools over and over
+
+        result = gaia_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
+            return final_response