Create tools.py

tools.py

@@ -0,0 +1,424 @@
+import os
+import re
+import json
+import base64
+import tempfile
+import subprocess
+
+import requests
+import pandas as pd
+
+from bs4 import BeautifulSoup
+from duckduckgo_search import DDGS
+from huggingface_hub import InferenceClient
+
+
+# ============================================================
+# CONFIG
+# ============================================================
+
+SCORING_API = "https://agents-course-unit4-scoring.hf.space"
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+TEXT_MODEL = os.getenv(
+    "TEXT_MODEL",
+    "Qwen/Qwen2.5-72B-Instruct"
+)
+
+VISION_MODEL = os.getenv(
+    "VISION_MODEL",
+    "Qwen/Qwen2.5-VL-72B-Instruct"
+)
+
+ASR_MODEL = os.getenv(
+    "ASR_MODEL",
+    "openai/whisper-large-v3"
+)
+
+client = InferenceClient(
+    token=HF_TOKEN
+)
+
+
+# ============================================================
+# WEB SEARCH
+# ============================================================
+
+def web_search(query):
+    results = []
+
+    with DDGS() as ddgs:
+        for result in ddgs.text(
+            query,
+            max_results=6
+        ):
+            results.append({
+                "title": result.get("title", ""),
+                "url": result.get("href", ""),
+                "snippet": result.get("body", "")
+            })
+
+    return json.dumps(
+        results,
+        ensure_ascii=False,
+        indent=2
+    )
+
+
+# ============================================================
+# WEBPAGE READER
+# ============================================================
+
+def read_webpage(url):
+    response = requests.get(
+        url,
+        timeout=25,
+        headers={
+            "User-Agent": "Mozilla/5.0"
+        }
+    )
+
+    response.raise_for_status()
+
+    soup = BeautifulSoup(
+        response.text,
+        "html.parser"
+    )
+
+    for tag in soup([
+        "script",
+        "style",
+        "nav",
+        "footer",
+        "header"
+    ]):
+        tag.decompose()
+
+    text = soup.get_text(
+        "\n",
+        strip=True
+    )
+
+    # Avoid sending enormous webpages to the model.
+    return text[:40000]
+
+
+# ============================================================
+# DOWNLOAD GAIA ATTACHMENT
+# ============================================================
+
+def download_file(task_id):
+    url = f"{SCORING_API}/files/{task_id}"
+
+    response = requests.get(
+        url,
+        timeout=60
+    )
+
+    response.raise_for_status()
+
+    content_type = response.headers.get(
+        "content-type",
+        ""
+    ).lower()
+
+    extension = ""
+
+    if "png" in content_type:
+        extension = ".png"
+
+    elif "jpeg" in content_type or "jpg" in content_type:
+        extension = ".jpg"
+
+    elif "mp3" in content_type or "audio" in content_type:
+        extension = ".mp3"
+
+    elif "spreadsheet" in content_type or "excel" in content_type:
+        extension = ".xlsx"
+
+    elif "pdf" in content_type:
+        extension = ".pdf"
+
+    elif "python" in content_type or "text" in content_type:
+        extension = ".txt"
+
+    else:
+        # Try to infer from Content-Disposition.
+        disposition = response.headers.get(
+            "content-disposition",
+            ""
+        )
+
+        match = re.search(
+            r'filename="?([^"]+)"?',
+            disposition
+        )
+
+        if match:
+            extension = os.path.splitext(
+                match.group(1)
+            )[1]
+
+    if not extension:
+        extension = ".bin"
+
+    path = os.path.join(
+        tempfile.gettempdir(),
+        f"gaia_{task_id}{extension}"
+    )
+
+    with open(path, "wb") as f:
+        f.write(response.content)
+
+    return path
+
+
+# ============================================================
+# FILE INSPECTION
+# ============================================================
+
+def inspect_file(path):
+    extension = os.path.splitext(
+        path
+    )[1].lower()
+
+    # ---------------- TXT / PY / CSV ----------------
+
+    if extension in [
+        ".txt",
+        ".py",
+        ".md",
+        ".csv",
+        ".tsv"
+    ]:
+
+        with open(
+            path,
+            "r",
+            encoding="utf-8",
+            errors="ignore"
+        ) as f:
+            return f.read()[:50000]
+
+    # ---------------- EXCEL ----------------
+
+    if extension in [
+        ".xlsx",
+        ".xls"
+    ]:
+
+        excel = pd.ExcelFile(path)
+
+        output = []
+
+        for sheet in excel.sheet_names:
+
+            df = pd.read_excel(
+                path,
+                sheet_name=sheet
+            )
+
+            output.append(
+                f"\n--- SHEET: {sheet} ---\n"
+            )
+
+            output.append(
+                f"Columns: {list(df.columns)}\n"
+            )
+
+            output.append(
+                f"Rows: {len(df)}\n"
+            )
+
+            output.append(
+                df.head(50).to_string(
+                    index=False
+                )
+            )
+
+        return "\n".join(output)[:50000]
+
+    # ---------------- PDF ----------------
+
+    if extension == ".pdf":
+
+        try:
+            import PyPDF2
+
+            reader = PyPDF2.PdfReader(path)
+
+            pages = []
+
+            for page in reader.pages:
+                pages.append(
+                    page.extract_text() or ""
+                )
+
+            return "\n".join(pages)[:50000]
+
+        except Exception as e:
+            return f"Could not read PDF: {e}"
+
+    return (
+        f"File downloaded successfully at:\n"
+        f"{path}\n"
+        f"File type: {extension}\n"
+        f"Use an appropriate tool to analyze it."
+    )
+
+
+# ============================================================
+# PYTHON EXECUTION
+# ============================================================
+
+def run_python(code):
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".py",
+        delete=False,
+        encoding="utf-8"
+    ) as f:
+
+        f.write(code)
+
+        path = f.name
+
+    try:
+
+        result = subprocess.run(
+            ["python", path],
+            capture_output=True,
+            text=True,
+            timeout=30
+        )
+
+        return (
+            "STDOUT:\n"
+            + result.stdout
+            + "\nSTDERR:\n"
+            + result.stderr
+        )[:30000]
+
+    finally:
+
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+# ============================================================
+# AUDIO TRANSCRIPTION
+# ============================================================
+
+def transcribe_audio(path):
+
+    result = client.automatic_speech_recognition(
+        path,
+        model=ASR_MODEL
+    )
+
+    return result.text
+
+
+# ============================================================
+# IMAGE ANALYSIS
+# ============================================================
+
+def analyze_image(path, question):
+
+    with open(path, "rb") as f:
+        image_bytes = f.read()
+
+    encoded = base64.b64encode(
+        image_bytes
+    ).decode("utf-8")
+
+    # Determine MIME type.
+    extension = os.path.splitext(
+        path
+    )[1].lower()
+
+    if extension == ".jpg" or extension == ".jpeg":
+        mime = "image/jpeg"
+    else:
+        mime = "image/png"
+
+    response = client.chat.completions.create(
+        model=VISION_MODEL,
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": question
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url":
+                            f"data:{mime};base64,{encoded}"
+                        }
+                    }
+                ]
+            }
+        ],
+        max_tokens=1000,
+        temperature=0
+    )
+
+    return response.choices[0].message.content
+
+
+# ============================================================
+# TOOL DISPATCHER
+# ============================================================
+
+def execute_tool(
+    tool_name,
+    tool_input,
+    task_id,
+    question
+):
+
+    if tool_name == "web_search":
+        return web_search(tool_input)
+
+    if tool_name == "read_webpage":
+        return read_webpage(tool_input)
+
+    if tool_name == "download_file":
+        path = download_file(task_id)
+
+        return (
+            f"FILE_PATH:\n{path}\n\n"
+            f"{inspect_file(path)}"
+        )
+
+    if tool_name == "inspect_file":
+
+        path = tool_input.strip()
+
+        return inspect_file(path)
+
+    if tool_name == "run_python":
+        return run_python(tool_input)
+
+    if tool_name == "transcribe_audio":
+
+        path = tool_input.strip()
+
+        return transcribe_audio(path)
+
+    if tool_name == "analyze_image":
+
+        path = tool_input.strip()
+
+        return analyze_image(
+            path,
+            question
+        )
+
+    return f"Unknown tool: {tool_name}"