Blame for tools.py

"""Tools for the GAIA agent: web search, arithmetic, and image recognition."""

import ast
import base64
import math
import operator
import re
from typing import Any

import requests
from ddgs import DDGS
from langchain_core.messages import HumanMessage
from prompts import IMAGE_ANALYSIS_PROMPT

MAX_IMAGES_PER_QUESTION = 3
_IMAGE_EXTENSIONS = r"(?:jpg|jpeg|png|gif|webp|bmp|svg)"
_URL_TAIL = r"(?:\?[^\s\)\]\"\'<>]*)?"

# Safe operators for GAIA-style numeric questions
_SAFE_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.FloorDiv: operator.floordiv,
    ast.Pow: operator.pow,
    ast.Mod: operator.mod,
    ast.USub: operator.neg,
    ast.UAdd: operator.pos,
}


def _eval_ast(node: ast.AST) -> float:
    if isinstance(node, ast.Constant):
        if isinstance(node.value, (int, float)):
            return float(node.value)
        raise ValueError(f"Unsupported constant: {node.value}")
    if isinstance(node, ast.Num):  # pragma: no cover - py<3.8
        return float(node.n)
    if isinstance(node, ast.BinOp):
        op = _SAFE_OPERATORS.get(type(node.op))
        if not op:
            raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
        return op(_eval_ast(node.left), _eval_ast(node.right))
    if isinstance(node, ast.UnaryOp):
        op = _SAFE_OPERATORS.get(type(node.op))
        if not op:
            raise ValueError(f"Unsupported unary op: {type(node.op).__name__}")
        return op(_eval_ast(node.operand))
    if isinstance(node, ast.Call):
        if isinstance(node.func, ast.Name):
            name = node.func.id
            args = [_eval_ast(a) for a in node.args]
            if name == "sqrt" and len(args) == 1:
                return math.sqrt(args[0])
            if name == "round" and len(args) in (1, 2):
                return float(round(args[0], int(args[1]) if len(args) == 2 else 0))
            if name == "abs" and len(args) == 1:
                return abs(args[0])
        raise ValueError("Unsupported function call")
    raise ValueError(f"Unsupported expression node: {type(node).__name__}")


def _normalize_expression(expr: str) -> str:
    expr = expr.strip()
    expr = expr.replace("×", "*").replace("÷", "/").replace("^", "**")
    expr = re.sub(r"\s+", "", expr)
    return expr


def calculate(expression: str) -> str:
    """
    Evaluate a numeric expression safely (numbers, + - * / // ** %, sqrt, abs, round).
    Returns a string result for the agent to use as evidence.
    """
    expr = _normalize_expression(expression)
    if not expr:
        return "Error: empty expression"
    try:
        tree = ast.parse(expr, mode="eval")
        result = _eval_ast(tree.body)
        if result == int(result) and "." not in str(expr).split("/")[-1]:
            return str(int(result))
        return str(round(result, 10)).rstrip("0").rstrip(".")
    except Exception as e:
        return f"Error: {e}"


def run_calculations(expressions: list[str]) -> list[str]:
    """Run multiple calculations; skip empty or duplicate expressions."""
    seen: set[str] = set()
    outputs: list[str] = []
    for raw in expressions:
        expr = str(raw).strip()
        if not expr or expr in seen:
            continue
        seen.add(expr)
        result = calculate(expr)
        outputs.append(f"{expr} = {result}")
    return outputs


def _message_content(response: Any) -> str:
    content = response.content
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
            elif isinstance(block, dict) and block.get("text"):
                parts.append(str(block["text"]))
        return "\n".join(parts).strip()
    return str(content).strip()


def extract_image_urls(text: str) -> list[str]:
    """Find image URLs in question text (markdown, HTML, plain links)."""
    candidates: list[str] = []
    patterns = [
        rf"https?://[^\s\)\]\"\'<>]+\.{_IMAGE_EXTENSIONS}{_URL_TAIL}",
        r"!\[[^\]]*\]\((https?://[^\)]+)\)",
        r"https?://huggingface\.co/datasets/[^\s\)\]\"\'<>]+",
        r"https?://[^\s\)\]\"\'<>]*(?:/image|/images|/file|/files)[^\s\)\]\"\'<>]*",
    ]
    for pattern in patterns:
        for match in re.finditer(pattern, text, re.IGNORECASE):
            url = match.group(1) if match.lastindex else match.group(0)
            url = url.rstrip(".,;:")
            if url.startswith("http"):
                candidates.append(url)

    seen: set[str] = set()
    deduped: list[str] = []
    for url in candidates:
        key = url.lower()
        if key not in seen:
            seen.add(key)
            deduped.append(url)
        if len(deduped) >= MAX_IMAGES_PER_QUESTION:
            break
    return deduped


def _fetch_image_base64(url: str) -> tuple[str, str]:
    """Download image and return (media_type, base64_data)."""
    headers = {"User-Agent": "Mozilla/5.0 (compatible; GAIA-Agent/1.0)"}
    resp = requests.get(url, timeout=45, headers=headers)
    resp.raise_for_status()
    content_type = (resp.headers.get("content-type") or "image/jpeg").split(";")[0].strip()
    if not content_type.startswith("image/"):
        content_type = "image/jpeg"
    return content_type, base64.standard_b64encode(resp.content).decode("ascii")


def describe_image(url: str, question: str, vision_model: Any) -> str:
    """
    Analyze an image URL with Claude vision and return a text description.
    """
    try:
        media_type, data = _fetch_image_base64(url)
        prompt = IMAGE_ANALYSIS_PROMPT.format(question=question, url=url)
        message = HumanMessage(
            content=[
                {"type": "text", "text": prompt},
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": media_type,
                        "data": data,
                    },
                },
            ]
        )
        return _message_content(vision_model.invoke([message]))
    except Exception as e:
        return f"Vision error: {e}"


def analyze_images(urls: list[str], question: str, vision_model: Any) -> list[str]:
    """Run vision on up to MAX_IMAGES_PER_QUESTION URLs."""
    results: list[str] = []
    for url in urls[:MAX_IMAGES_PER_QUESTION]:
        description = describe_image(url, question, vision_model)
        results.append(f"[Image: {url}]\n{description}")
    return results


def search_web(query: str) -> str:
    """Search the web for information related to the query."""
    results = DDGS().text(query, max_results=8)
    if not results:
        return "No results found."
    return "\n\n".join(
        f"{r.get('title', '')}: {r.get('body', r.get('href', ''))}"
        for r in results
    )