Final_Assignment_Template

Sleeping

Support DeepSeek Api Key

3ec4862 over 1 year ago

4.64 kB

"""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,
)

 
# ---------------------------------------------------------------------------
# Load the added system prompt from system_prompt.txt (located in the same directory)
# ---------------------------------------------------------------------------
ADDED_PROMPT_PATH = os.path.join(os.path.dirname(__file__), "added_prompt.txt")
with open(ADDED_PROMPT_PATH, "r", encoding="utf-8") as f:
    ADDED_PROMPT = f.read().strip()


# ---------------------------------------------------------------------------
# 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 == "together":
        from smolagents import InferenceClientModel
        hf_model_id = os.getenv("TOGETHER_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct")
        hf_token = os.getenv("HF_API_KEY")
        return InferenceClientModel(
            model_id=hf_model_id,
            provider=provider,
            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")
        base_url = os.getenv("OPENAI_BASE_URL", "https://api.deepseek.com/v1")
        return OpenAIServerModel(
            model_id=openai_model_id,
            api_key=openai_token,
            api_base=base_url
        )

    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()
        )
        # Append the additional prompt to the existing system prompt
        self.prompt_templates["system_prompt"] += f"\n\n{ADDED_PROMPT}"

    # Convenience so the object itself can be *called* directly
    def __call__(self, question: str, **kwargs: Any) -> str:
        steps = self.run(question, **kwargs)
        # If steps is a primitive, just return it
        if isinstance(steps, (int, float, str)):
            return str(steps).strip()
        last_step = None
        for step in steps:
            last_step = step
        # Defensive: handle int/float/str directly
        if isinstance(last_step, (int, float, str)):
            return str(last_step).strip()
        answer = getattr(last_step, "answer", None)
        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"]

	"""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,
	)


	# ---------------------------------------------------------------------------
	# Load the added system prompt from system_prompt.txt (located in the same directory)
	# ---------------------------------------------------------------------------
	ADDED_PROMPT_PATH = os.path.join(os.path.dirname(__file__), "added_prompt.txt")
	with open(ADDED_PROMPT_PATH, "r", encoding="utf-8") as f:
	ADDED_PROMPT = f.read().strip()


	# ---------------------------------------------------------------------------
	# 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 == "together":
	from smolagents import InferenceClientModel
	hf_model_id = os.getenv("TOGETHER_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct")
	hf_token = os.getenv("HF_API_KEY")
	return InferenceClientModel(
	model_id=hf_model_id,
	provider=provider,
	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")
	base_url = os.getenv("OPENAI_BASE_URL", "https://api.deepseek.com/v1")
	return OpenAIServerModel(
	model_id=openai_model_id,
	api_key=openai_token,
	api_base=base_url
	)

	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()
	)
	# Append the additional prompt to the existing system prompt
	self.prompt_templates["system_prompt"] += f"\n\n{ADDED_PROMPT}"

	# Convenience so the object itself can be called directly
	def __call__(self, question: str, **kwargs: Any) -> str:
	steps = self.run(question, **kwargs)
	# If steps is a primitive, just return it
	if isinstance(steps, (int, float, str)):
	return str(steps).strip()
	last_step = None
	for step in steps:
	last_step = step
	# Defensive: handle int/float/str directly
	if isinstance(last_step, (int, float, str)):
	return str(last_step).strip()
	answer = getattr(last_step, "answer", None)
	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"]