tools.py

"""tools.py
Reusable LangChain Tool definitions for GAIA evaluation agent.

Each tool complies with the LangChain `Tool` interface so they can be
plugged into a LangGraph-driven agent.  All tools are stateless and rely on
external services or pure-Python execution.
"""
from __future__ import annotations

import os
import requests
from typing import Any, Dict, List

from langchain_community.tools import DuckDuckGoSearchRun
from langchain_experimental.tools import PythonREPLTool
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_core.tools import tool, Tool
from langchain.tools import BaseTool
from langchain.callbacks.manager import CallbackManagerForToolRun

# Fallback API URL (same default as in app.py)
DEFAULT_API_URL = os.getenv("GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space")

# -----------------------------------------------------------------------------
# Simple Search Tool (DuckDuckGo)
# -----------------------------------------------------------------------------

duckduck_wrapper = DuckDuckGoSearchAPIWrapper()
search_tool = DuckDuckGoSearchRun(api_wrapper=duckduck_wrapper)

# -----------------------------------------------------------------------------
# Simple Python REPL tool (for quick calculations / parsing)
# -----------------------------------------------------------------------------
python_tool = PythonREPLTool()

# -----------------------------------------------------------------------------
# Calculator using built-in eval (safe-ish)
# -----------------------------------------------------------------------------
@tool("calculator", return_direct=True)
def calculator(expression: str) -> str:  # noqa: D401
    """Evaluate a basic mathematical expression and return the result.

    Uses Python's `eval` in a restricted namespace for quick arithmetic.
    """
    try:
        allowed_names: Dict[str, Any] = {
            k: v for k, v in vars(__import__("math")).items() if not k.startswith("__")
        }
        result = eval(expression, {"__builtins__": {}}, allowed_names)  # type: ignore[arg-type]
        return str(result)
    except Exception as exc:  # pragma: no cover
        return f"ERROR: {exc}"  # LLM can decide what to do with error

# -----------------------------------------------------------------------------
# File Loader Tool for GAIA Task-specific files
# -----------------------------------------------------------------------------
class GAIAFileLoaderTool(BaseTool):
    """Download auxiliary file for a GAIA task and return its contents or path."""

    name: str = "load_file"
    description: str = (
        "Use this tool to download an auxiliary file linked to a GAIA question. "
        "Input must be a valid task_id (integer). The tool returns text content of "
        "the file if it is textual, otherwise returns the local file path."
    )

    api_url: str = DEFAULT_API_URL  # injected configurable base URL

    def _run(
        self,
        task_id: str,
        *,
        run_manager: CallbackManagerForToolRun | None = None,
    ) -> str:
        """Synchronous run implementation."""
        try:
            int_task_id = int(task_id)
        except ValueError as exc:  # pragma: no cover
            return f"ERROR: task_id should be an integer: {exc}"

        files_endpoint = f"{self.api_url}/files/{int_task_id}"
        try:
            resp = requests.get(files_endpoint, timeout=30)
            resp.raise_for_status()
        except Exception as exc:  # pragma: no cover
            return f"ERROR: failed to download file: {exc}"

        # Determine content-type and return appropriate value
        content_type = resp.headers.get("content-type", "")
        if content_type.startswith("text/") or content_type in {"application/json", "application/xml"}:
            return resp.text[:8000]  # hard cap to keep prompts short

        # Otherwise, save binary file locally
        file_name = resp.headers.get("content-disposition") or f"gaia_task_{int_task_id}_file"
        file_path = os.path.join(os.getcwd(), file_name)
        with open(file_path, "wb") as f:
            f.write(resp.content)
        return file_path

    async def _arun(self, *args: Any, **kwargs: Any) -> str:  # noqa: D401
        """Asynchronous version not implemented (sync fallback)."""
        return self._run(*args, **kwargs)

# -----------------------------------------------------------------------------
# Public helper to create list of tools
# -----------------------------------------------------------------------------

def create_tools(api_url: str | None = None) -> List[Tool]:
    """Return a list of Tool objects for the agent.

    Parameters
    ----------
    api_url : str | None
        Override for GAIA API base url. If None, falls back to DEFAULT_API_URL.
    """

    file_loader = GAIAFileLoaderTool(api_url=api_url or DEFAULT_API_URL)

    # Wrap DuckDuckGo and Python REPL into standard Tool when necessary
    wrapped_search = Tool(
        name="search",
        func=search_tool.run,
        description="Search the web quickly for general information.",
        return_direct=False,
    )

    wrapped_python = Tool(
        name="python",
        func=python_tool.run,
        description="Execute Python code for calculation or data parsing.",
        return_direct=False,
    )

    return [wrapped_search, calculator, wrapped_python, file_loader]