Update agent.py

agent.py

@@ -20,44 +20,44 @@ from prompt import system_prompt
 # --------------------------------------------------------------------
 # 1. API Key Rotation Setup
 # --------------------------------------------------------------------
-api_keys = [
+api_keys = [k for k in [
     os.getenv("OPENROUTER_API_KEY"),
     os.getenv("OPENROUTER_API_KEY_1")
-]
+] if k]
 
-if not any(api_keys):
+if not 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])
+api_key_cycle = itertools.cycle(api_keys)
 
-def get_next_api_key():
+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 API keys on failure."""
+    """ChatOpenAI wrapper that automatically rotates OpenRouter 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()
+        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 {self.api_key[:5]}... failed, trying next key...")
+                    print(f"[API Key Rotation] Key {current_key[:5]}... failed, trying next key...")
                     continue
-                raise  # Re-raise other unexpected errors
-        raise RuntimeError("All OpenRouter API keys failed or rate-limited.")
+                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",
-    api_key=get_next_api_key(),  # Start with the first key
-    model="qwen/qwen3-coder:free",  # Model must support function calling
+    openai_api_key=get_next_api_key(),  # ✅ start with the first key
+    model="qwen/qwen3-coder:free",      # must support tool/function calling
     temperature=1
 )
 
@@ -99,18 +99,31 @@ def llm_call(state: MessagesState):
 # Tool Node
 def tool_node(state: MessagesState):
     """Executes tools requested by the LLM."""
-    result = []
-    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}
+    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
 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 last_message.tool_calls else END
+    return "Action" if getattr(last_message, "tool_calls", None) else END
 
 # --------------------------------------------------------------------
 # 5. Build LangGraph Agent
@@ -154,8 +167,8 @@ class LangGraphAgent:
         result = gaia_agent.invoke(input_state, config)
         final_response = result["messages"][-1].content
 
-        try:
+        # Extract "FINAL ANSWER" if present
+        if isinstance(final_response, str):
             return final_response.split("FINAL ANSWER:")[-1].strip()
-        except Exception:
-            print("Could not split on 'FINAL ANSWER:'")
-            return final_response
+        else:
+            return str(final_response)