Add new gaia agent

agent/gaia.py

@@ -0,0 +1,215 @@
+import logging
+import os
+import asyncio
+import time
+
+from pydantic import BaseModel
+from typing import Annotated
+
+from langchain_anthropic import ChatAnthropic, convert_to_anthropic_tool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool, InjectedToolCallId
+from langchain_tavily import TavilySearch
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState, ToolNode
+
+
+logging.basicConfig(level=os.getenv("LOGLEVEL", "WARNING"))
+
+
+class AgentState(BaseModel):
+    question: dict
+    final_agent_answer: dict | None
+    messages: Annotated[list, add_messages]
+
+
+@tool
+def add(x: float, y: float):
+    """This function adds two numbers."""
+    logging.info(f"Added {x} and {y}")
+    return x + y
+
+
+@tool
+def subtract(x: float, y: float):
+    """This function subtracts two numbers."""
+    logging.info(f"Subtracting {y} from {x}")
+    return x - y
+
+
+@tool
+def multiply(x: float, y: float):
+    """This function multiplies two numbers."""
+    logging.info(f"Multiplying {x} and {y}")
+    return x * y
+
+
+@tool
+def divide(x: float, y: float):
+    """this function divides two numbers. handles division by zero."""
+    logging.info(f"dividing {x} by {y}")
+    if y == 0:
+        return "error: cannot divide by zero."
+    return x / y
+
+
+@tool
+def submit_final_answer(
+    answer: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    state: Annotated[AgentState, InjectedState],
+):
+    """This function should be called to submit your final answer only once you have tetermined it. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
+
+    print(f"Submitting final answer: {answer}")
+
+    answer_data = {
+        "task_id": state.question["task_id"],
+        "agent_answer": answer,
+    }
+
+    logging.info("Final answer written, updating state with final answer...")
+    return Command(
+        update={
+            "final_agent_answer": answer_data,
+            "messages": [
+                ToolMessage(
+                    "You have successfully submitted your final answer. There is nothing left to be done.",
+                    tool_call_id=tool_call_id,
+                )
+            ],
+        }
+    )
+
+
+tavily = TavilySearch(max_results=2)
+tools = [add, subtract, multiply, divide, tavily, submit_final_answer]
+
+anthropic_tools = []
+
+for raw_tool in tools:
+    anthropic_tool = convert_to_anthropic_tool(raw_tool)
+    anthropic_tools.append(anthropic_tool)
+
+# To cache all tools we add the cache control block to the last tool
+anthropic_tools[-1]["cache_control"] = {"type": "ephemeral"}
+
+
+llm = ChatAnthropic(
+    model_name="claude-sonnet-4-20250514",
+    max_tokens=5000,
+    timeout=None,
+    thinking={"type": "enabled", "budget_tokens": 4000},
+    model_kwargs={
+        "extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"}
+    },
+)
+llm_with_tools = llm.bind_tools(anthropic_tools)
+
+
+class GaiaAgent:
+    def __init__(self):
+        self.llm = llm_with_tools
+        self.agent_graph = self.compile_graph()
+
+    def compile_graph(self):
+        graph = StateGraph(AgentState)
+
+        def should_continue(state):
+            logging.info(
+                "Checking for final answer in decide_next_node conditional edge"
+            )
+            logging.info(state.final_agent_answer)
+            if state.final_agent_answer:
+                logging.info("Final answer submitted. Ending agent flow.")
+                return END
+            else:
+                logging.info(
+                    "No final answer submitted yet, proceed to the tool nodes."
+                )
+                return "tools"
+
+        graph.add_node(self.consider_question)
+        graph.add_node("tools", ToolNode(tools))
+
+        graph.add_edge(START, "consider_question")
+        graph.add_edge("tools", "consider_question")
+        graph.add_conditional_edges(
+            "consider_question", should_continue, ["tools", END]
+        )
+        return graph.compile()
+
+    async def consider_question(self, state: AgentState):
+        """Home of the agent. Looks at all the messages so far, generates the next message."""
+        logging.info("Considering question...")
+        time.sleep(5)
+        if state.final_agent_answer is None:
+            messages = state.messages
+            response = await self.llm.ainvoke(messages)
+            if hasattr(response, "content"):
+                for message in response.content:
+                    if "text" in message:
+                        logging.info(message["text"])
+            return {"messages": [response]}
+        else:
+            # If a final answer has been determined no more consideration is required
+            logging.info(
+                "Skipping question consideration because final answer is available"
+            )
+            return state
+
+    async def answer_question(self, question):
+        if question["file_name"]:
+            return "I don't know - I can't handle files yet!"
+        question_text = question["question"]
+
+        logging.debug("Initializing agent state to answer question...")
+        system_prompt = """
+        You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer by calling the submit_final_answer tool. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+        To operate effectively, always remember:
+            1. Before using any math tools for operations, make sure you have thought about the math problem sufficiently and stated the equation that you will solve. Plan the equation first, then use the math tools to solve it precisely.
+        """
+        initial_state = {
+            "question": question,
+            "final_agent_answer": None,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": system_prompt,
+                            "cache_control": {"type": "ephemeral"},
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": question_text,
+                },
+            ],
+        }
+
+        async def get_final_answer(agent):
+            final_output: dict | None = None
+            async for mode, chunk in agent.astream(
+                initial_state,
+                stream_mode=["values", "messages"],
+                config={"recursion_limit": 30},
+            ):
+                if mode == "values":
+                    final_output = chunk
+
+            if final_output is None:
+                return "I don't know!"
+
+            return final_output["final_agent_answer"]["agent_answer"]
+
+        result = await get_final_answer(self.agent_graph)
+        return result
+
+    def __call__(self, question):
+        return asyncio.run(self.answer_question(question))