Update tools.py

tools.py

@@ -4,7 +4,7 @@ import json
 import base64
 import tempfile
 import subprocess
-from pathlib import Path
+import mimetypes
 
 import requests
 import pandas as pd
@@ -13,17 +13,9 @@ from bs4 import BeautifulSoup
 from groq import Groq
 
 
-# ============================================================
-# CONFIG
-# ============================================================
+SCORING_API = "https://agents-course-unit4-scoring.hf.space"
 
-SCORING_API = (
-    "https://agents-course-unit4-scoring.hf.space"
-)
-
-GROQ_API_KEY = os.getenv(
-    "GROQ_API_KEY"
-)
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
 VISION_MODEL = os.getenv(
     "GROQ_VISION_MODEL",
@@ -40,237 +32,368 @@ groq_client = Groq(
 )
 
 
-# ============================================================
-# WEB SEARCH
-# ============================================================
+# ---------------------------------------------------------
+# Web search state
+# ---------------------------------------------------------
+
+LAST_SEARCH_RESULTS = []
+
+
+# ---------------------------------------------------------
+# Web search
+# ---------------------------------------------------------
 
 def web_search(query):
 
-    try:
+    global LAST_SEARCH_RESULTS
 
+    try:
         from ddgs import DDGS
 
-    except ImportError:
+        results = []
 
-        from duckduckgo_search import DDGS
+        with DDGS() as ddgs:
+            search_results = ddgs.text(
+                query,
+                max_results=6
+            )
 
-    results = []
+            for item in search_results:
 
-    with DDGS() as ddgs:
+                result = {
+                    "title": item.get("title", ""),
+                    "url": item.get("href", ""),
+                    "snippet": item.get("body", "")
+                }
 
-        for item in ddgs.text(
-            query,
-            max_results=6
-        ):
+                results.append(result)
 
-            results.append(
-                {
-                    "title": item.get(
-                        "title",
-                        ""
-                    ),
-                    "url": item.get(
-                        "href",
-                        ""
-                    ),
-                    "snippet": item.get(
-                        "body",
-                        ""
-                    )
-                }
+        LAST_SEARCH_RESULTS = results
+
+        if not results:
+            return "No search results found."
+
+        output = []
+
+        for i, item in enumerate(results):
+
+            output.append(
+                f"[{i}]\n"
+                f"TITLE: {item['title']}\n"
+                f"URL: {item['url']}\n"
+                f"SNIPPET: {item['snippet']}"
             )
 
-    return json.dumps(
-        results,
-        ensure_ascii=False,
-        indent=2
-    )
+        return "\n\n".join(output)
 
+    except Exception as e:
 
-# ============================================================
-# WEBPAGE
-# ============================================================
+        return f"Web search error: {type(e).__name__}: {e}"
 
-def read_webpage(url):
 
-    response = requests.get(
-        url,
-        timeout=30,
-        headers={
-            "User-Agent":
-                "Mozilla/5.0 GAIA-Agent"
-        }
-    )
+# ---------------------------------------------------------
+# Read webpage
+# ---------------------------------------------------------
 
-    response.raise_for_status()
+def read_webpage(
+    url=None,
+    cursor=None,
+    pattern=None
+):
 
-    soup = BeautifulSoup(
-        response.text,
-        "html.parser"
-    )
+    global LAST_SEARCH_RESULTS
 
-    for tag in soup(
-        [
-            "script",
-            "style",
-            "noscript",
-            "svg",
-            "nav",
-            "footer"
-        ]
-    ):
+    # -----------------------------------------------------
+    # Resolve missing URL using latest search results
+    # -----------------------------------------------------
 
-        tag.decompose()
+    if not url:
 
-    text = soup.get_text(
-        "\n",
-        strip=True
-    )
+        if cursor is None:
+            cursor = 0
 
-    # Keep context manageable.
-    return text[:25000]
+        try:
+            cursor = int(cursor)
+        except Exception:
+            cursor = 0
 
+        if (
+            cursor < 0
+            or cursor >= len(LAST_SEARCH_RESULTS)
+        ):
+            return (
+                "ERROR: No valid URL was supplied and the requested "
+                "search-result cursor does not exist."
+            )
 
-# ============================================================
-# DOWNLOAD GAIA FILE
-# ============================================================
+        url = LAST_SEARCH_RESULTS[cursor].get("url", "")
 
-def download_file(task_id):
+    if not url:
+        return "ERROR: Could not resolve webpage URL."
 
-    url = (
-        f"{SCORING_API}/files/{task_id}"
-    )
+    # -----------------------------------------------------
+    # Fetch webpage
+    # -----------------------------------------------------
 
-    response = requests.get(
-        url,
-        timeout=60
-    )
+    try:
 
-    response.raise_for_status()
+        response = requests.get(
+            url,
+            timeout=20,
+            headers={
+                "User-Agent": (
+                    "Mozilla/5.0 "
+                    "(compatible; GAIA-Agent/1.0)"
+                )
+            }
+        )
 
-    content_type = response.headers.get(
-        "content-type",
-        ""
-    ).lower()
+        response.raise_for_status()
 
-    disposition = response.headers.get(
-        "content-disposition",
-        ""
-    )
+        soup = BeautifulSoup(
+            response.text,
+            "html.parser"
+        )
 
-    filename = None
+        for tag in soup(
+            [
+                "script",
+                "style",
+                "noscript",
+                "svg",
+                "nav",
+                "footer",
+                "header"
+            ]
+        ):
+            tag.decompose()
 
-    match = re.search(
-        r'filename="?([^"]+)"?',
-        disposition
-    )
+        text = soup.get_text(
+            "\n",
+            strip=True
+        )
+
+        # -------------------------------------------------
+        # If model supplied a pattern, return context
+        # around occurrences rather than the whole page.
+        # -------------------------------------------------
+
+        if pattern:
+
+            pattern_lower = pattern.lower()
+            text_lower = text.lower()
+
+            positions = []
+
+            start = 0
+
+            while True:
 
-    if match:
+                position = text_lower.find(
+                    pattern_lower,
+                    start
+                )
 
-        filename = match.group(1)
+                if position == -1:
+                    break
 
-    if not filename:
+                positions.append(position)
 
-        extension_map = {
-            "audio/mpeg": ".mp3",
-            "audio/mp3": ".mp3",
-            "audio/wav": ".wav",
-            "image/png": ".png",
-            "image/jpeg": ".jpg",
-            "application/pdf": ".pdf",
-            "text/plain": ".txt",
-            "text/x-python": ".py",
-            "application/python": ".py",
-            "text/csv": ".csv",
-            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx"
-        }
+                start = position + len(pattern_lower)
 
-        extension = ".bin"
+                if len(positions) >= 5:
+                    break
 
-        for key, value in extension_map.items():
+            if positions:
 
-            if key in content_type:
+                chunks = []
 
-                extension = value
-                break
+                for position in positions:
 
-        filename = (
-            f"{task_id}{extension}"
+                    left = max(
+                        0,
+                        position - 1800
+                    )
+
+                    right = min(
+                        len(text),
+                        position + 3500
+                    )
+
+                    chunks.append(
+                        text[left:right]
+                    )
+
+                text = "\n\n--- MATCH ---\n\n".join(
+                    chunks
+                )
+
+        if len(text) > 18000:
+            text = text[:18000]
+
+        return (
+            f"URL: {url}\n\n"
+            f"CONTENT:\n{text}"
         )
 
-    filename = re.sub(
-        r"[^A-Za-z0-9_.-]",
-        "_",
-        filename
-    )
+    except Exception as e:
 
-    folder = (
-        Path(tempfile.gettempdir())
-        / "gaia_files"
-    )
+        return (
+            f"Webpage read error: "
+            f"{type(e).__name__}: {e}"
+        )
 
-    folder.mkdir(
-        parents=True,
-        exist_ok=True
-    )
 
-    path = (
-        folder
-        / filename
-    )
+# ---------------------------------------------------------
+# Download GAIA attachment
+# ---------------------------------------------------------
 
-    path.write_bytes(
-        response.content
-    )
+def download_file(task_id):
 
-    print(
-        f"Downloaded: {path}"
-    )
+    try:
+
+        url = f"{SCORING_API}/files/{task_id}"
 
-    return str(path)
+        response = requests.get(
+            url,
+            timeout=60
+        )
 
+        response.raise_for_status()
 
-# ============================================================
-# TEXT FILE
-# ============================================================
+        content_type = response.headers.get(
+            "content-type",
+            ""
+        ).lower()
+
+        disposition = response.headers.get(
+            "content-disposition",
+            ""
+        )
+
+        filename = None
+
+        # -------------------------------------------------
+        # Try Content-Disposition
+        # -------------------------------------------------
+
+        match = re.search(
+            r'filename="?([^";]+)"?',
+            disposition,
+            re.IGNORECASE
+        )
+
+        if match:
+            filename = match.group(1)
+
+        # -------------------------------------------------
+        # Try URL filename
+        # -------------------------------------------------
+
+        if not filename:
+
+            url_name = url.rstrip("/").split("/")[-1]
+
+            if "." in url_name:
+                filename = url_name
+
+        # -------------------------------------------------
+        # Content-type fallback
+        # -------------------------------------------------
+
+        if not filename:
+
+            extension_map = {
+                "audio/mpeg": ".mp3",
+                "audio/mp3": ".mp3",
+                "audio/wav": ".wav",
+                "audio/x-wav": ".wav",
+                "image/png": ".png",
+                "image/jpeg": ".jpg",
+                "image/jpg": ".jpg",
+                "application/pdf": ".pdf",
+                "text/plain": ".txt",
+                "text/x-python": ".py",
+                "application/json": ".json",
+                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
+                "application/vnd.ms-excel": ".xls"
+            }
+
+            filename = (
+                "gaia_attachment"
+                + extension_map.get(
+                    content_type,
+                    ".bin"
+                )
+            )
+
+        # -------------------------------------------------
+        # Save
+        # -------------------------------------------------
+
+        safe_filename = os.path.basename(
+            filename
+        )
+
+        path = os.path.join(
+            tempfile.gettempdir(),
+            safe_filename
+        )
+
+        with open(
+            path,
+            "wb"
+        ) as f:
+            f.write(response.content)
+
+        return path
+
+    except Exception as e:
+
+        raise RuntimeError(
+            f"Could not download task file: "
+            f"{type(e).__name__}: {e}"
+        )
+
+
+# ---------------------------------------------------------
+# Inspect text
+# ---------------------------------------------------------
 
 def inspect_text(path):
 
-    return Path(path).read_text(
+    with open(
+        path,
+        "r",
         encoding="utf-8",
-        errors="ignore"
-    )[:30000]
+        errors="replace"
+    ) as f:
 
+        text = f.read()
 
-# ============================================================
-# EXCEL
-# ============================================================
+    if len(text) > 20000:
+        text = text[:20000]
 
-def inspect_excel(path):
+    return text
 
-    workbook = pd.ExcelFile(
-        path
-    )
 
-    output = []
+# ---------------------------------------------------------
+# Inspect Excel
+# ---------------------------------------------------------
 
-    output.append(
-        "SHEETS:"
-    )
+def inspect_excel(path):
 
-    output.append(
-        str(workbook.sheet_names)
+    sheets = pd.read_excel(
+        path,
+        sheet_name=None
     )
 
-    for sheet in workbook.sheet_names:
+    output = []
 
-        df = pd.read_excel(
-            path,
-            sheet_name=sheet
-        )
+    for sheet_name, df in sheets.items():
 
         output.append(
-            f"\n===== {sheet} ====="
+            f"=== SHEET: {sheet_name} ==="
         )
 
         output.append(
@@ -281,188 +404,198 @@ def inspect_excel(path):
             f"Columns: {list(df.columns)}"
         )
 
-        # Avoid flooding the model.
         output.append(
             df.head(100).to_string(
                 index=False
             )
         )
 
-    return "\n".join(
-        output
-    )[:35000]
+    return "\n\n".join(output)
 
 
-# ============================================================
-# PDF
-# ============================================================
+# ---------------------------------------------------------
+# Inspect PDF
+# ---------------------------------------------------------
 
 def inspect_pdf(path):
 
     from PyPDF2 import PdfReader
 
-    reader = PdfReader(
-        path
-    )
+    reader = PdfReader(path)
 
     pages = []
 
     for page in reader.pages:
 
-        text = (
-            page.extract_text()
-            or ""
-        )
+        text = page.extract_text() or ""
 
         pages.append(text)
 
-    return "\n".join(
-        pages
-    )[:35000]
+        if len("\n".join(pages)) > 20000:
+            break
+
+    text = "\n".join(pages)
+
+    return text[:20000]
+
+
+# ---------------------------------------------------------
+# Inspect Python / source
+# ---------------------------------------------------------
 
+def inspect_source(path):
 
-# ============================================================
-# GENERIC FILE INSPECTION
-# ============================================================
+    return inspect_text(path)
 
-def inspect_file(path):
+
+# ---------------------------------------------------------
+# Generic file inspection
+# ---------------------------------------------------------
+
+def inspect_file(task_id):
+
+    path = download_file(task_id)
 
     extension = (
-        Path(path)
-        .suffix
+        os.path.splitext(path)[1]
         .lower()
     )
 
+    if extension in [
+        ".xlsx",
+        ".xls"
+    ]:
+        return inspect_excel(path)
+
+    if extension == ".pdf":
+        return inspect_pdf(path)
+
     if extension in [
         ".txt",
         ".py",
-        ".md",
-        ".csv",
-        ".tsv",
         ".json",
-        ".xml"
+        ".csv",
+        ".md",
+        ".xml",
+        ".html"
     ]:
-
-        return inspect_text(
-            path
-        )
+        return inspect_source(path)
 
     if extension in [
-        ".xlsx",
-        ".xls"
+        ".png",
+        ".jpg",
+        ".jpeg"
     ]:
-
-        return inspect_excel(
-            path
+        return (
+            f"Image file downloaded successfully: {path}\n"
+            "Use analyze_image for visual analysis."
         )
 
-    if extension == ".pdf":
-
-        return inspect_pdf(
-            path
+    if extension in [
+        ".mp3",
+        ".wav",
+        ".m4a",
+        ".ogg"
+    ]:
+        return (
+            f"Audio file downloaded successfully: {path}\n"
+            "Use transcribe_audio for speech analysis."
         )
 
     return (
-        f"Unsupported file type: "
-        f"{extension}"
+        f"Downloaded file: {path}\n"
+        f"Extension: {extension}"
     )
 
 
-# ============================================================
-# PYTHON
-# ============================================================
+# ---------------------------------------------------------
+# Run Python
+# ---------------------------------------------------------
 
 def run_python(code):
 
-    folder = (
-        Path(tempfile.gettempdir())
-        / "gaia_python"
+    path = os.path.join(
+        tempfile.gettempdir(),
+        "gaia_agent_code.py"
     )
 
-    folder.mkdir(
-        parents=True,
-        exist_ok=True
-    )
+    with open(
+        path,
+        "w",
+        encoding="utf-8"
+    ) as f:
+        f.write(code)
 
-    script = (
-        folder
-        / "script.py"
-    )
+    try:
 
-    script.write_text(
-        code,
-        encoding="utf-8"
-    )
+        result = subprocess.run(
+            [
+                "python",
+                path
+            ],
+            capture_output=True,
+            text=True,
+            timeout=30
+        )
 
-    process = subprocess.run(
-        [
-            "python",
-            str(script)
-        ],
-        capture_output=True,
-        text=True,
-        timeout=30
-    )
+        output = (
+            result.stdout
+            + "\n"
+            + result.stderr
+        ).strip()
 
-    output = (
-        process.stdout
-        + "\n"
-        + process.stderr
-    )
+        if len(output) > 12000:
+            output = output[:12000]
 
-    return output[-20000:]
+        return output
 
+    except subprocess.TimeoutExpired:
 
-# ============================================================
-# AUDIO TRANSCRIPTION
-# ============================================================
+        return "Python execution timed out after 30 seconds."
 
-def transcribe_audio(path):
+    except Exception as e:
+
+        return (
+            f"Python execution error: "
+            f"{type(e).__name__}: {e}"
+        )
+
+
+# ---------------------------------------------------------
+# Audio transcription
+# ---------------------------------------------------------
+
+def transcribe_audio(task_id):
+
+    path = download_file(task_id)
 
     with open(
         path,
         "rb"
-    ) as audio:
-
-        result = (
-            groq_client.audio.transcriptions.create(
-                file=audio,
-                model=AUDIO_MODEL,
-                response_format="text",
-                temperature=0
-            )
-        )
-
-    if hasattr(
-        result,
-        "text"
-    ):
+    ) as audio_file:
 
-        return result.text
+        transcription = groq_client.audio.transcriptions.create(
+            file=audio_file,
+            model=AUDIO_MODEL,
+            response_format="text"
+        )
 
-    return str(result)
+    return str(transcription)
 
 
-# ============================================================
-# IMAGE ANALYSIS
-# ============================================================
+# ---------------------------------------------------------
+# Image analysis
+# ---------------------------------------------------------
 
 def analyze_image(
-    path,
+    task_id,
     question
 ):
 
-    raw = Path(
-        path
-    ).read_bytes()
-
-    encoded = base64.b64encode(
-        raw
-    ).decode("utf-8")
+    path = download_file(task_id)
 
     extension = (
-        Path(path)
-        .suffix
+        os.path.splitext(path)[1]
         .lower()
     )
 
@@ -473,104 +606,110 @@ def analyze_image(
         ".webp": "image/webp"
     }
 
-    mime = mime_map.get(
+    mime_type = mime_map.get(
         extension,
         "image/png"
     )
 
-    data_url = (
-        f"data:{mime};base64,{encoded}"
-    )
-
-    response = (
-        groq_client.chat.completions.create(
-            model=VISION_MODEL,
-            messages=[
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": question
-                        },
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": data_url
-                            }
+    with open(
+        path,
+        "rb"
+    ) as f:
+
+        encoded = base64.b64encode(
+            f.read()
+        ).decode("utf-8")
+
+    response = groq_client.chat.completions.create(
+
+        model=VISION_MODEL,
+
+        temperature=0,
+
+        max_completion_tokens=1500,
+
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": (
+                            "Analyze this image carefully.\n\n"
+                            f"Question: {question}\n\n"
+                            "Return the factual answer needed for the question."
+                        )
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": (
+                                f"data:{mime_type};base64,"
+                                f"{encoded}"
+                            )
                         }
-                    ]
-                }
-            ],
-            temperature=0.1,
-            max_tokens=1500
-        )
+                    }
+                ]
+            }
+        ]
     )
 
-    return (
-        response
-        .choices[0]
-        .message
-        .content
-    )
+    return response.choices[0].message.content
 
 
-# ============================================================
-# TOOL DISPATCH
-# ============================================================
+# ---------------------------------------------------------
+# Tool dispatcher
+# ---------------------------------------------------------
 
 def execute_tool(
-    tool_name,
-    tool_input
+    name,
+    args,
+    task_id=None
 ):
 
-    if tool_name == "web_search":
+    if name == "web_search":
 
         return web_search(
-            str(tool_input)
+            args.get("query", "")
         )
 
-    if tool_name == "read_webpage":
+    if name == "read_webpage":
 
         return read_webpage(
-            str(tool_input)
+            url=args.get("url"),
+            cursor=args.get("cursor"),
+            pattern=args.get("pattern")
         )
 
-    if tool_name == "inspect_file":
+    if name == "download_file":
+
+        return download_file(
+            args.get("task_id") or task_id
+        )
+
+    if name == "inspect_file":
 
         return inspect_file(
-            str(tool_input)
+            args.get("task_id") or task_id
         )
 
-    if tool_name == "run_python":
+    if name == "run_python":
 
         return run_python(
-            str(tool_input)
+            args.get("code", "")
         )
 
-    if tool_name == "transcribe_audio":
+    if name == "transcribe_audio":
 
         return transcribe_audio(
-            str(tool_input)
+            args.get("task_id") or task_id
         )
 
-    if tool_name == "analyze_image":
-
-        if not isinstance(
-            tool_input,
-            dict
-        ):
-
-            return (
-                "analyze_image requires "
-                "path and question."
-            )
+    if name == "analyze_image":
 
         return analyze_image(
-            tool_input["path"],
-            tool_input["question"]
+            task_id=args.get("task_id") or task_id,
+            question=args.get("question", "")
         )
 
-    return (
-        f"Unknown tool: {tool_name}"
-    )
+    return f"Unknown tool: {name}"