File size: 2,002 Bytes
22a5e4a
 
a810880
 
ad06e3b
ad11720
712cf2c
 
6a16db0
 
22a5e4a
e2bfbed
22a5e4a
a810880
ad11720
c73fe98
 
a810880
 
ad11720
17d9955
 
a810880
 
 
 
 
f8abc87
ad11720
a810880
 
 
 
 
 
 
 
 
 
 
 
 
 
17d9955
6a16db0
a810880
 
 
 
 
 
 
ad11720
 
a810880
ad11720
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import os

from smolagents import CodeAgent, LiteLLMModel, WikipediaSearchTool

from prompt import SYSTEM_PROMPT
from tools import download_gaia_file, read_file_content


class BasicAgent:
    def __init__(self):
        self.model = LiteLLMModel(
            model_id="openrouter/anthropic/claude-sonnet-4",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            temperature=0.0,
            max_tokens=4096,
        )

        # add_base_tools=True gives us python_interpreter, web_search
        # (DuckDuckGo) and visit_webpage for free, so we only add what's
        # missing: Wikipedia search plus our own GAIA file tools.
        self.agent = CodeAgent(
            model=self.model,
            tools=[
                WikipediaSearchTool(user_agent="GAIA-BasicAgent (contact: example@example.com)"),
                download_gaia_file,
                read_file_content,
            ],
            add_base_tools=True,
            max_steps=10,
            instructions=SYSTEM_PROMPT,
            additional_authorized_imports=[
                "pandas",
                "numpy",
                "requests",
                "json",
                "re",
                "math",
                "statistics",
                "datetime",
                "collections",
                "itertools",
                "PIL",
            ],
        )

    def __call__(self, question: str, task_id: str | None = None) -> str:
        prompt = question
        if task_id:
            prompt = (
                f"task_id: {task_id}\n"
                f"Question: {question}\n\n"
                "If this question mentions or implies an attached file, call "
                "download_gaia_file with this exact task_id first, then call "
                "read_file_content on the path it returns before answering."
            )
        try:
            result = self.agent.run(prompt)
        except Exception as e:
            return f"AGENT ERROR: {e}"
        return str(result).strip()