Update tools.py

tools.py

@@ -1,61 +1,323 @@
+import os
+import re
+import json
+import subprocess
+from pathlib import Path
+from typing import Optional
+
+import requests
 import pandas as pd
+from langchain_core.tools import tool
+from youtube_transcript_api import YouTubeTranscriptApi
+
+
+DEFAULT_API_URL = os.getenv(
+    "AGENT_COURSE_API_URL",
+    "https://agents-course-unit4-scoring.hf.space",
+)
+
+CACHE_DIR = Path("/tmp/agent_course_files")
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _safe_filename(name: str) -> str:
+    """Make a filename safe for local storage."""
+    return re.sub(r"[^a-zA-Z0-9._-]+", "_", name).strip("_") or "file"
+
+
+def _filename_from_headers(task_id: str, response: requests.Response) -> str:
+    """Infer filename from Content-Disposition or fallback to task_id."""
+    content_disposition = response.headers.get("content-disposition", "")
+
+    match = re.search(r'filename="?([^";]+)"?', content_disposition)
+    if match:
+        return _safe_filename(match.group(1))
+
+    content_type = response.headers.get("content-type", "").lower()
+
+    ext = ""
+    if "spreadsheet" in content_type or "excel" in content_type:
+        ext = ".xlsx"
+    elif "csv" in content_type:
+        ext = ".csv"
+    elif "python" in content_type or "text" in content_type:
+        ext = ".txt"
+    elif "audio" in content_type:
+        ext = ".mp3"
+    elif "image/png" in content_type:
+        ext = ".png"
+    elif "image/jpeg" in content_type:
+        ext = ".jpg"
+
+    return _safe_filename(task_id) + ext
+
+
+def _download_file_by_task_id(task_id: str) -> Optional[Path]:
+    """Download attached task file from the official scoring API."""
+    if not task_id:
+        return None
+
+    url = f"{DEFAULT_API_URL}/files/{task_id}"
+
+    try:
+        response = requests.get(url, timeout=30)
+
+        if response.status_code == 404:
+            return None
+
+        response.raise_for_status()
+
+        filename = _filename_from_headers(task_id, response)
+        file_path = CACHE_DIR / filename
+
+        file_path.write_bytes(response.content)
+        return file_path
 
-def answer_excel_question(self, question: str, file_path: str) -> str:
-    df = pd.read_excel(file_path)
+    except Exception:
+        return None
 
-    prompt = f"""
-    Question:
-    {question}
-    
-    Excel data:
-    {df.to_csv(index=False)}
-    
-    Answer directly. Return only the final answer.
+
+def _resolve_file(task_id: str = "", file_path: str = "") -> Optional[Path]:
+    """Resolve either a local file path or a task_id into a real local Path."""
+    if file_path:
+        path = Path(file_path)
+        if path.exists():
+            return path
+
+    if task_id:
+        return _download_file_by_task_id(task_id)
+
+    return None
+
+
+@tool
+def download_task_file(task_id: str) -> str:
     """
+    Download the attached file for a Hugging Face Agent Course task_id.
+    Use this when a question mentions an attached file, image, audio, spreadsheet, or code file.
+    Returns the local file path and basic file information.
+    """
+    path = _download_file_by_task_id(task_id)
 
-    response = self.model.invoke([
-        SystemMessage(content="You are a precise data analysis assistant."),
-        HumanMessage(content=prompt)
-    ])
+    if path is None:
+        return f"No attached file found for task_id={task_id}."
 
-    return response.content.strip()
+    info = {
+        "task_id": task_id,
+        "file_path": str(path),
+        "filename": path.name,
+        "suffix": path.suffix,
+        "size_bytes": path.stat().st_size,
+    }
 
+    return json.dumps(info, ensure_ascii=False)
 
-import subprocess
 
-def answer_python_question(self, question: str, file_path: str) -> str:
+@tool
+def read_attached_text_file(task_id: str = "", file_path: str = "", max_chars: int = 12000) -> str:
+    """
+    Read the text content of an attached file.
+    Use this for .txt, .py, .csv, .md, .json, or other plain-text attachments.
+    Provide either task_id or file_path.
+    """
+    path = _resolve_file(task_id=task_id, file_path=file_path)
+
+    if path is None:
+        return "No file could be resolved from the given task_id or file_path."
+
+    try:
+        text = path.read_text(encoding="utf-8", errors="replace")
+        return text[:max_chars]
+    except Exception as e:
+        return f"Failed to read file {path}: {e}"
+
+
+@tool
+def answer_python_question(task_id: str = "", file_path: str = "") -> str:
+    """
+    Execute an attached Python file and return its final output.
+    Use this when the question asks for the final numeric output from attached Python code.
+    Provide either task_id or file_path.
+    """
+    path = _resolve_file(task_id=task_id, file_path=file_path)
+
+    if path is None:
+        return "No Python file could be resolved from the given task_id or file_path."
+
     try:
         result = subprocess.run(
-            ["python", file_path],
+            ["python", str(path)],
+            cwd=str(path.parent),
             capture_output=True,
             text=True,
             timeout=10,
         )
-        output = result.stdout.strip()
-        if output:
-            return output.splitlines()[-1].strip()
-        return result.stderr.strip()
+
+        stdout = result.stdout.strip()
+        stderr = result.stderr.strip()
+
+        if stdout:
+            lines = [line.strip() for line in stdout.splitlines() if line.strip()]
+            return lines[-1] if lines else stdout
+
+        if stderr:
+            return f"Python execution produced no stdout. stderr:\n{stderr[-2000:]}"
+
+        return "Python execution finished with no output."
+
+    except subprocess.TimeoutExpired:
+        return "Python execution timed out."
     except Exception as e:
-        return f""
+        return f"Failed to execute Python file: {e}"
 
 
-from youtube_transcript_api import YouTubeTranscriptApi
-from urllib.parse import urlparse, parse_qs
-
-def extract_youtube_id(self, url: str) -> str | None:
-    parsed = urlparse(url)
-    if parsed.hostname in ["www.youtube.com", "youtube.com"]:
-        return parse_qs(parsed.query).get("v", [None])[0]
-    if parsed.hostname == "youtu.be":
-        return parsed.path.lstrip("/")
-    return None
+@tool
+def answer_excel_question(task_id: str = "", file_path: str = "", question: str = "") -> str:
+    """
+    Read an attached Excel or CSV file and return compact spreadsheet data for answering the question.
+    Use this for questions about attached .xlsx, .xls, or .csv files.
+    Provide task_id when available. Also include the original question.
+    """
+    path = _resolve_file(task_id=task_id, file_path=file_path)
+
+    if path is None:
+        return "No spreadsheet file could be resolved from the given task_id or file_path."
+
+    try:
+        suffix = path.suffix.lower()
+
+        if suffix == ".csv":
+            sheets = {"sheet1": pd.read_csv(path)}
+        else:
+            sheets = pd.read_excel(path, sheet_name=None)
+
+        outputs = []
+
+        for sheet_name, df in sheets.items():
+            df = df.copy()
+
+            outputs.append(f"Sheet: {sheet_name}")
+            outputs.append(f"Shape: {df.shape[0]} rows x {df.shape[1]} columns")
+            outputs.append(f"Columns: {list(df.columns)}")
+
+            preview_csv = df.head(80).to_csv(index=False)
+            outputs.append("Preview CSV:")
+            outputs.append(preview_csv)
+
+            numeric_summary = df.select_dtypes(include="number").sum(numeric_only=True)
+            if not numeric_summary.empty:
+                outputs.append("Numeric column sums:")
+                outputs.append(numeric_summary.to_string())
+
+            # Heuristic helper for common “food excluding drinks” sales question.
+            q = question.lower()
+            if "food" in q and ("drink" in q or "drinks" in q):
+                possible_category_cols = [
+                    col for col in df.columns
+                    if any(key in str(col).lower() for key in ["category", "type", "item", "name", "menu"])
+                ]
+                possible_value_cols = [
+                    col for col in df.columns
+                    if any(key in str(col).lower() for key in ["sales", "revenue", "amount", "total", "usd", "price"])
+                ]
+
+                if possible_category_cols and possible_value_cols:
+                    category_col = possible_category_cols[0]
+                    value_col = possible_value_cols[0]
+
+                    values = (
+                        df[value_col]
+                        .astype(str)
+                        .str.replace(r"[$,]", "", regex=True)
+                    )
+                    values = pd.to_numeric(values, errors="coerce").fillna(0)
+
+                    drink_pattern = r"drink|drinks|beverage|soda|coffee|tea|juice|water|milkshake|shake|smoothie"
+                    food_mask = ~df[category_col].astype(str).str.lower().str.contains(
+                        drink_pattern,
+                        regex=True,
+                        na=False,
+                    )
+
+                    total_food_sales = values[food_mask].sum()
+                    outputs.append(
+                        f"Heuristic food-not-drinks total using category column "
+                        f"'{category_col}' and value column '{value_col}': "
+                        f"${total_food_sales:.2f}"
+                    )
+
+        return "\n\n".join(outputs)[:20000]
+
+    except Exception as e:
+        return f"Failed to read spreadsheet {path}: {e}"
+
+
+def _extract_youtube_id(text: str) -> Optional[str]:
+    """Extract YouTube video ID from a URL or text containing a URL."""
+    urls = re.findall(r"https?://[^\s)]+", text)
 
-def get_youtube_transcript(self, question: str) -> str:
-    import re
-    urls = re.findall(r"https?://\S+", question)
     for url in urls:
-        video_id = self.extract_youtube_id(url)
-        if video_id:
-            transcript = YouTubeTranscriptApi.get_transcript(video_id)
-            return " ".join([x["text"] for x in transcript])
-    return ""
+        if "youtube.com/watch" in url:
+            match = re.search(r"[?&]v=([^&\s]+)", url)
+            if match:
+                return match.group(1)
+
+        if "youtu.be/" in url:
+            match = re.search(r"youtu\.be/([^?&\s]+)", url)
+            if match:
+                return match.group(1)
+
+    return None
+
+
+def _transcript_to_text(transcript) -> str:
+    """Convert transcript result to plain text, compatible with old/new youtube-transcript-api."""
+    pieces = []
+
+    for item in transcript:
+        if isinstance(item, dict):
+            text = item.get("text", "")
+        else:
+            text = getattr(item, "text", "")
+
+        if text:
+            pieces.append(text.replace("\n", " ").strip())
+
+    return " ".join(pieces)
+
+
+@tool
+def get_youtube_transcript(url_or_question: str) -> str:
+    """
+    Get the transcript text of a YouTube video from a URL or a question containing a YouTube URL.
+    Use this when the question asks what someone says in a YouTube video.
+    This may not help for purely visual questions about video frames.
+    """
+    video_id = _extract_youtube_id(url_or_question)
+
+    if not video_id:
+        return "No YouTube video ID found."
+
+    try:
+        # Compatibility with older youtube-transcript-api versions.
+        try:
+            transcript = YouTubeTranscriptApi.get_transcript(
+                video_id,
+                languages=["en", "en-US", "en-GB"],
+            )
+        except AttributeError:
+            api = YouTubeTranscriptApi()
+            transcript = api.fetch(
+                video_id,
+                languages=["en", "en-US", "en-GB"],
+            )
+
+        text = _transcript_to_text(transcript)
+
+        if not text:
+            return "Transcript was found but is empty."
+
+        return text[:20000]
+
+    except Exception as e:
+        return f"Failed to fetch YouTube transcript for video_id={video_id}: {e}"