File size: 3,351 Bytes
52d1305
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bb49a20
75a272e
bb49a20
 
 
 
 
 
52d1305
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4a751ef
52d1305
4a751ef
 
52d1305
 
 
 
 
 
bb49a20
 
52d1305
bb49a20
 
52d1305
 
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
"""agent.py – GAIA benchmark agent using *smolagents*.

This module exposes:

* ``gaia_agent()`` – factory returning a ready‑to‑use agent instance.
* ``GAIAAgent``  – subclass of ``smolagents.CodeAgent``.

The LLM backend is chosen at runtime via the ``MODEL_PROVIDER``
environment variable (``hf`` or ``openai``) exactly like *example.py*.
"""

import os
from typing import Any, Sequence

from dotenv import load_dotenv

# SmolAgents Tools
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    Tool,
)

# Custom Tools from tools.py
from tools import (
    PythonRunTool,
    ExcelLoaderTool,
    YouTubeTranscriptTool,
    AudioTranscriptionTool,
    SimpleOCRTool,
)

# ---------------------------------------------------------------------------
# Model selection helper
# ---------------------------------------------------------------------------

load_dotenv()  # Make sure we read credentials from .env when running locally

def _select_model():
    """Return a smolagents *model* as configured by the ``MODEL_PROVIDER`` env."""

    provider = os.getenv("MODEL_PROVIDER", "hf").lower()

    if provider == "hf":
        from smolagents import InferenceClientModel
        hf_model_id = os.getenv("HF_MODEL", "HuggingFaceH4/zephyr-7b-beta")
        hf_token = os.getenv("HF_API_KEY")
        return InferenceClientModel(model_id=hf_model_id, token=hf_token)

    if provider == "openai":
        from smolagents import OpenAIServerModel
        openai_model_id = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
        openai_token = os.getenv("OPENAI_API_KEY")
        return OpenAIServerModel(model_id=openai_model_id, api_key=openai_token)

    raise ValueError(
        f"Unsupported MODEL_PROVIDER: {provider!r}. "
        "Use 'hf' (default) or 'openai'."
    )

# ---------------------------------------------------------------------------
# Core Agent implementation
# ---------------------------------------------------------------------------

DEFAULT_TOOLS = [
    DuckDuckGoSearchTool(),
    PythonRunTool(),
    ExcelLoaderTool(),
    YouTubeTranscriptTool(),
    AudioTranscriptionTool(),
    SimpleOCRTool(),
]

class GAIAAgent(CodeAgent):
    def __init__(self, tools=None):
        super().__init__(tools=tools or DEFAULT_TOOLS, model=_select_model())

    # Convenience so the object itself can be *called* directly
    def __call__(self, question: str, **kwargs: Any) -> str:
        steps = self.run(question, **kwargs)
        last_step = None
        for step in steps:
            last_step = step
        # If last_step is a FinalAnswerStep with .answer, return it
        answer = getattr(last_step, "answer", None)
        # Always return as a string, with stripping for safety
        if answer is not None:
            return str(answer).strip()
        return str(last_step).strip()

# ---------------------------------------------------------------------------
# Factory helpers expected by app.py
# ---------------------------------------------------------------------------

def gaia_agent(*, extra_tools: Sequence[Tool] | None = None) -> GAIAAgent:
    # Compose the toolset: always include all default tools, plus any extras
    toolset = list(DEFAULT_TOOLS)
    if extra_tools:
        toolset.extend(extra_tools)
    return GAIAAgent(tools=toolset)

__all__ = ["GAIAAgent", "gaia_agent"]