tools.py

import base64
import math
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

import requests
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain_groq import ChatGroq

from agents import FileAnalysisAgent, MediaAgent

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

WEB_QUERY_MAX_CHARS = 400
WEB_RESULT_MAX_CHARS = 1_000
WEB_ANSWER_MAX_CHARS = 1_200
WEB_OUTPUT_MAX_CHARS = 5_000
WEB_MAX_RESULTS = 3

_tavily_client = None


def _truncate(text: str, max_chars: int) -> str:
    """Keep tool payloads bounded while making truncation visible."""
    text = str(text or "").strip()
    if len(text) <= max_chars:
        return text
    marker = "\n[truncated]"
    if max_chars <= len(marker):
        return text[:max_chars]
    return f"{text[: max_chars - len(marker)].rstrip()}{marker}"


def _compact_query(query: str) -> str:
    """Normalize and cap a natural-language query before sending it to Tavily."""
    compact = " ".join(str(query or "").split())
    if not compact:
        raise ValueError("The web search query is empty.")
    if len(compact) <= WEB_QUERY_MAX_CHARS:
        return compact

    shortened = compact[:WEB_QUERY_MAX_CHARS]
    if " " in shortened:
        shortened = shortened.rsplit(" ", 1)[0]
    return shortened


def _get_tavily_client():
    """Create the client only when web search is actually requested."""
    global _tavily_client
    if _tavily_client is None:
        api_key = os.getenv("TAVILY_API_KEY")
        if not api_key:
            raise RuntimeError("TAVILY_API_KEY is not configured.")
        from tavily import TavilyClient

        _tavily_client = TavilyClient(api_key=api_key)
    return _tavily_client


def _format_search_results(response: dict) -> str:
    """Return a small, predictable search payload for the manager model."""
    blocks = []
    answer = response.get("answer")
    if answer:
        blocks.append(f"Search summary: {_truncate(answer, WEB_ANSWER_MAX_CHARS)}")

    for index, result in enumerate(response.get("results", [])[:WEB_MAX_RESULTS], 1):
        title = _truncate(result.get("title", "Untitled"), 200)
        url = str(result.get("url", "")).strip()
        content = _truncate(result.get("content", ""), WEB_RESULT_MAX_CHARS)
        blocks.append(
            f"Source {index}\nTitle: {title}\nURL: {url}\nExcerpt: {content}"
        )

    if not blocks:
        return "No web results found."
    return _truncate("\n\n".join(blocks), WEB_OUTPUT_MAX_CHARS)


@tool
def search_tool(query: str) -> str:
    """
    Search the web for current information.

    Args:
        query: Search query.
    """
    compact_query = _compact_query(query)
    response = _get_tavily_client().search(
        query=compact_query,
        max_results=WEB_MAX_RESULTS,
        search_depth="basic",
        include_answer=True,
        include_raw_content=False,
        include_images=False,
    )
    return _format_search_results(response)


@tool
def fetch_url(url: str) -> str:
    """Fetch a web page and extract its main readable text.

    Use this when a search snippet is not enough and you need the full page
    content.

    Args:
        url: HTTP(S) URL.
    """
    import trafilatura

    downloaded = trafilatura.fetch_url(url)
    if not downloaded:
        return f"Failed to fetch {url}"
    text = trafilatura.extract(downloaded)
    return text or "(no extractable content)"


@tool
def wikipedia_lookup(title: str) -> str:
    """Fetch the summary of an English Wikipedia article.

    More reliable than web search for historical facts, dates, and biographies.

    Args:
        title: Article title, e.g. "Marie Curie".
    """
    import wikipediaapi

    wiki = wikipediaapi.Wikipedia(
        user_agent="hf-agents-course-final/1.0 (batberni@gmail.com)",
        language="en",
    )
    page = wiki.page(title)
    if not page.exists():
        return f"No Wikipedia page found for '{title}'"
    return f"{page.summary}\n\nURL: {page.fullurl}"


@tool
def download_task_file(task_id: str) -> str:
    """Download the file attached to a task from the HF agents course API.

    Many evaluation questions come with an attached file (image, audio, pdf,
    xlsx, csv, ...). Call this first to get a local path, then pass that path
    to read_pdf / read_excel / analyze_image / transcribe_audio, etc.

    Args:
        task_id: The task_id from the questions endpoint.

    Returns:
        The local file path where the file was saved.
    """
    url = f"{DEFAULT_API_URL}/files/{task_id}"
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()

    cd = resp.headers.get("Content-Disposition", "")
    m = re.search(r'filename="?([^"]+)"?', cd)
    filename = m.group(1) if m else task_id

    path = Path(tempfile.gettempdir()) / filename
    path.write_bytes(resp.content)
    return str(path)


@tool
def read_pdf(path: str) -> str:
    """Extract text content from a PDF file.

    Args:
        path: Local file path to the PDF.
    """
    from pypdf import PdfReader

    reader = PdfReader(path)
    return "\n".join(page.extract_text() or "" for page in reader.pages)


@tool
def read_excel(path: str, sheet: str | None = None) -> str:
    """Read an Excel workbook and return its contents as text.

    Args:
        path: Local path to .xlsx / .xls file.
        sheet: Optional sheet name. Omit to dump every sheet.
    """
    import pandas as pd

    if sheet:
        return pd.read_excel(path, sheet_name=sheet).to_string()
    sheets = pd.read_excel(path, sheet_name=None)
    return "\n\n".join(
        f"=== {name} ===\n{df.to_string()}" for name, df in sheets.items()
    )


@tool
def read_csv(path: str) -> str:
    """Read a CSV file and return its contents as text.

    Args:
        path: Local path to .csv file.
    """
    import pandas as pd

    return pd.read_csv(path).to_string()


@tool
def read_text_file(path: str) -> str:
    """Read a plain-text file (txt, json, md, py, ...) and return its contents.

    Args:
        path: Local file path.
    """
    return Path(path).read_text(encoding="utf-8", errors="replace")


@tool
def python_exec(code: str) -> str:
    """Execute Python code in a subprocess and return stdout (and stderr).

    Use this for arithmetic on lists, date math, pandas manipulations, parsing,
    or anything that is easier to *compute* than to reason about. You must
    `print(...)` any value you want to see back.

    Available: the standard library plus pandas, numpy (if installed in the env).

    Args:
        code: Python source to execute.
    """
    try:
        result = subprocess.run(
            [sys.executable, "-c", code],
            capture_output=True,
            text=True,
            timeout=30,
        )
        out = result.stdout
        if result.stderr:
            out += f"\n[stderr]\n{result.stderr}"
        return out.strip() or "(no output)"
    except subprocess.TimeoutExpired:
        return "(timeout after 30s)"


_CALC_ENV = {
    name: getattr(math, name) for name in dir(math) if not name.startswith("_")
}
_CALC_ENV.update({"abs": abs, "round": round, "min": min, "max": max, "sum": sum})


@tool
def calculator(expression: str) -> str:
    """Evaluate a math expression.

    Supports +, -, *, /, //, %, **, parentheses, and the math module
    (sin, cos, sqrt, log, pi, e, ...). No variables, no assignment.

    Args:
        expression: e.g. "2 * (3 + 4) ** 0.5" or "sqrt(2) + log(10)".
    """
    try:
        return str(eval(expression, {"__builtins__": {}}, _CALC_ENV))
    except Exception as e:
        return f"Error: {e}"


_IMAGE_MIME = {
    "jpg": "jpeg",
    "jpeg": "jpeg",
    "png": "png",
    "gif": "gif",
    "webp": "webp",
    "bmp": "bmp",
}


@tool
def analyze_image(image_path_or_url: str, question: str) -> str:
    """Answer a question about an image using a vision-language model.

    Handles both local files and http(s) URLs. Use this for chess positions,
    charts, screenshots, photos, diagrams, etc.

    Args:
        image_path_or_url: Local file path or HTTP(S) URL.
        question: What you want to know about the image.
    """
    vlm_model = os.getenv("GROQ_VLM_MODEL")
    if not vlm_model:
        return (
            "No vision model configured. Set env var GROQ_VLM_MODEL to a "
            "Groq VLM id (verify with the /v1/models endpoint on your account)."
        )

    if image_path_or_url.startswith(("http://", "https://")):
        image_url = image_path_or_url
    else:
        with open(image_path_or_url, "rb") as f:
            b64 = base64.b64encode(f.read()).decode()
        ext = image_path_or_url.rsplit(".", 1)[-1].lower()
        mime = _IMAGE_MIME.get(ext, "jpeg")
        image_url = f"data:image/{mime};base64,{b64}"

    vlm = ChatGroq(
        model=vlm_model,
        temperature=0.1,
        api_key=os.getenv("GROQ_API_KEY"),
    )
    msg = HumanMessage(
        content=[
            {"type": "text", "text": question},
            {"type": "image_url", "image_url": {"url": image_url}},
        ]
    )
    return vlm.invoke([msg]).content


@tool
def transcribe_audio(audio_path: str) -> str:
    """Transcribe an audio file to text with Whisper (Groq).

    Args:
        audio_path: Local file path to the audio (mp3, wav, m4a, flac, ...).
    """
    from groq import Groq

    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
    with open(audio_path, "rb") as f:
        transcription = client.audio.transcriptions.create(
            file=(os.path.basename(audio_path), f.read()),
            model="whisper-large-v3-turbo",
        )
    return transcription.text


@tool
def analyze_video(url_or_path: str, question: str, num_frames: int = 5) -> str:
    """Answer a question about a video by sampling frames and asking a VLM.

    Downloads the video if a URL is given (YouTube supported), samples frames
    evenly along its duration, and asks the vision model. Pair with
    ``youtube_transcript`` for full audio+visual understanding.

    Args:
        url_or_path: YouTube URL, other video URL, or local file path.
        question: What you want to know about the video.
        num_frames: How many frames to sample (default 5, max 5).
    """
    from imageio_ffmpeg import get_ffmpeg_exe

    num_frames = max(1, min(num_frames, 5))
    tmpdir = Path(tempfile.mkdtemp())

    if url_or_path.startswith(("http://", "https://")):
        import yt_dlp

        video_path = str(tmpdir / "video.mp4")
        ydl_opts = {
            "outtmpl": video_path,
            "format": "worst[ext=mp4]/worst",
            "quiet": True,
            "no_warnings": True,
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url_or_path])
    else:
        video_path = url_or_path

    ffmpeg = get_ffmpeg_exe()
    probe = subprocess.run([ffmpeg, "-i", video_path], capture_output=True, text=True)
    dur_match = re.search(r"Duration: (\d+):(\d+):(\d+)", probe.stderr)
    if dur_match:
        h, m, s = map(int, dur_match.groups())
        duration = max(1, h * 3600 + m * 60 + s)
    else:
        duration = 30

    frames_b64 = []
    for i in range(num_frames):
        t = duration * (i + 0.5) / num_frames
        frame_path = tmpdir / f"frame_{i}.jpg"
        subprocess.run(
            [
                ffmpeg,
                "-y",
                "-ss",
                str(t),
                "-i",
                video_path,
                "-vframes",
                "1",
                "-q:v",
                "3",
                str(frame_path),
            ],
            capture_output=True,
        )
        if frame_path.exists():
            b64 = base64.b64encode(frame_path.read_bytes()).decode()
            frames_b64.append(f"data:image/jpeg;base64,{b64}")

    if not frames_b64:
        return "Failed to extract any frames from the video."

    vlm_model = os.getenv("GROQ_VLM_MODEL")
    if not vlm_model:
        return (
            "No vision model configured. Set env var GROQ_VLM_MODEL to a "
            "Groq VLM id (verify with the /v1/models endpoint on your account)."
        )
    vlm = ChatGroq(
        model=vlm_model,
        temperature=0.1,
        api_key=os.getenv("GROQ_API_KEY"),
    )
    content = [
        {
            "type": "text",
            "text": f"Below are {len(frames_b64)} frames sampled evenly from a "
            f"video (duration ~{duration}s). {question}",
        }
    ]
    for url in frames_b64:
        content.append({"type": "image_url", "image_url": {"url": url}})
    return vlm.invoke([HumanMessage(content=content)]).content


@tool
def youtube_transcript(url_or_id: str) -> str:
    """Fetch the transcript of a YouTube video.

    Falls back gracefully when captions are unavailable.

    Args:
        url_or_id: Full YouTube URL or the 11-char video id.
    """
    from youtube_transcript_api import YouTubeTranscriptApi

    m = re.search(r"(?:v=|youtu\.be/|/shorts/|/embed/)([0-9A-Za-z_-]{11})", url_or_id)
    video_id = m.group(1) if m else url_or_id
    try:
        entries = YouTubeTranscriptApi.get_transcript(video_id)
    except Exception as e:
        return f"No transcript available for {video_id}: {e}"
    return " ".join(e["text"] for e in entries)


_file_agent = None
_media_agent = None


def _get_file_agent() -> FileAnalysisAgent:
    global _file_agent
    if _file_agent is None:
        _file_agent = FileAnalysisAgent()
    return _file_agent


def _get_media_agent() -> MediaAgent:
    global _media_agent
    if _media_agent is None:
        _media_agent = MediaAgent()
    return _media_agent


@tool
def web_research(question: str) -> str:
    """Search the web and return a compact summary with source excerpts.

    This is deliberately a direct, bounded Tavily call rather than a nested
    agent. Use a short query about one fact at a time. The query is capped at
    400 characters and the result payload at 5,000 characters.

    Args:
        question: Concise natural-language search query.
    """
    return search_tool.invoke({"query": question})


@tool
def file_analysis(question: str) -> str:
    """Answer a question about an attached document (PDF, XLSX, CSV, TXT, JSON).

    Delegates to a specialist that downloads the file (given a task_id, URL,
    or local path), routes to the right reader by extension, and answers.

    Include the file reference in your question, e.g.:
    "Given task_id abc-123, what is the total revenue in Q3?"

    Args:
        question: Natural-language question mentioning the task_id / URL / path.
    """
    return _get_file_agent()(question)


@tool
def media_analysis(question: str) -> str:
    """Answer a question about media (image, audio, or video).

    Delegates to a specialist that downloads if needed, then routes to VLM /
    Whisper / video frame sampling / YouTube captions.

    Include the media reference in your question, e.g.:
    "For task_id abc-123, what animal is shown?"
    "Watch https://youtube.com/watch?v=... and tell me the speaker's main claim."

    Args:
        question: Natural-language question mentioning the task_id / URL / path.
    """
    return _get_media_agent()(question)


ALL_TOOLS = [
    web_research,
    file_analysis,
    media_analysis,
    python_exec,
    calculator,
]