Blame for agent.py

"""
agent.py: GAIA agent for the Agents Course, Unit 4 (single file: app.py imports BasicAgent,
check_model and MODEL_ID from here).
Groq (openai/gpt-oss-120b) does the reasoning; Gemini is used only for images, audio,
PDFs and YouTube videos; Tavily is the web search (DuckDuckGo as fallback).

Secrets (Space Settings -> Variables and secrets):
    GROQ_API_KEY     required   console.groq.com
    TAVILY_API_KEY   optional   app.tavily.com    (not needed: without it, search uses free keyless
                                                   sources: the ddgs metasearch, DuckDuckGo, Wikipedia)
    GEMINI_API_KEY   optional   aistudio.google.com (without it, image/audio/video questions are skipped)
"""
import os
import re
import time
from typing import Optional
from urllib.parse import urlparse

import requests
from bs4 import BeautifulSoup

from google import genai
from google.genai import types

from langchain_core.rate_limiters import InMemoryRateLimiter
from langchain_core.tools import StructuredTool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_groq import ChatGroq
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_classic.agents import AgentExecutor, create_tool_calling_agent  # moved here in LangChain 1.x


class DailyLimitReached(Exception):
    """Groq's tokens-per-day budget for the current model is used up."""


class WebSearchTool:
    """DuckDuckGo HTML search with rate-limit backoff.
    Note: from a datacenter IP (like a HF Space) DuckDuckGo rate-limits
    aggressively, so this is best-effort. Swap in an API-key search
    (Tavily, Serper, Brave) if you need reliability.
    """

    def __init__(self):
        self.last_request_time = 0
        self.min_request_interval = 2.0
        self.max_retries = 5

    def search(self, query: str, domain: Optional[str] = None) -> str:
        for attempt in range(self.max_retries):
            now = time.time()
            gap = now - self.last_request_time
            if gap < self.min_request_interval:
                time.sleep(self.min_request_interval - gap)
            try:
                results = self._do_search(query, domain)
                self.last_request_time = time.time()
                return results
            except Exception as e:
                if "202 Ratelimit" in str(e) and attempt < self.max_retries - 1:
                    time.sleep((2 ** attempt) * self.min_request_interval)
                    continue
                return f"Search failed: {str(e)}"
        return "Search failed due to rate limiting."

    def _do_search(self, query: str, domain: Optional[str] = None) -> str:
        base_url = "https://html.duckduckgo.com/html"
        params = {"q": query + (f" site:{domain}" if domain else "")}
        headers = {"User-Agent": "Mozilla/5.0 (compatible; research-agent/1.0)"}
        response = requests.get(base_url, params=params, headers=headers, timeout=10)
        if response.status_code == 202:
            raise Exception("202 Ratelimit")
        response.raise_for_status()

        soup = BeautifulSoup(response.text, "html.parser")
        results = []
        for result in soup.find_all("div", {"class": "result"}):
            title = result.find("a", {"class": "result__a"})
            snippet = result.find("a", {"class": "result__snippet"})
            if title and snippet:
                results.append(
                    f"[{title.get_text()}]({title.get('href')})\n{snippet.get_text()}\n"
                )
        if not results:
            return "No search results found."
        return "## Search Results\n\n" + "\n".join(results[:8])


class GeminiAgent:
    def __init__(self, api_key: Optional[str] = None, model_name: str = "gemini-2.5-flash",
                 groq_model: str = "openai/gpt-oss-120b"):
        self.groq_model = groq_model
        import warnings
        warnings.filterwarnings("ignore", category=UserWarning)
        warnings.filterwarnings("ignore", category=DeprecationWarning)

        self.api_key = api_key
        self.model_name = model_name

        # Gemini is optional: without a key, the media tools say so instead of crashing.
        self.genai_client = genai.Client(api_key=api_key) if api_key else None

        self.web_search_tool = WebSearchTool()

        # Primary search: Tavily (reliable from datacenter IPs). Reads the
        # TAVILY_API_KEY secret from the environment. Falls back to the
        # DuckDuckGo scraper if the key or package is missing.
        self.tavily_client = None
        tavily_key = os.environ.get("TAVILY_API_KEY")
        if tavily_key:
            try:
                from tavily import TavilyClient
                self.tavily_client = TavilyClient(api_key=tavily_key)
            except Exception as e:
                print(f"Tavily unavailable, falling back to DuckDuckGo: {e}")

        self.llm = self._setup_llm()

        self.tools = [
            StructuredTool.from_function(
                func=self._web_search,
                name="web_search",
                description=(
                    "Search the web for current or factual information. "
                    "Pass the search query as 'query'."
                ),
            ),
            StructuredTool.from_function(
                func=self._visit_webpage,
                name="visit_webpage",
                description=(
                    "Read the text of a web page found by web_search. "
                    "Pass the full page URL as 'url'."
                ),
            ),
            StructuredTool.from_function(
                func=self._wikipedia_search,
                name="wikipedia",
                description=(
                    "Look up encyclopedic facts on Wikipedia. "
                    "Pass the topic or article title as 'query'."
                ),
            ),
            StructuredTool.from_function(
                func=self._analyze_video,
                name="analyze_video",
                description=(
                    "Analyze a YouTube video. Pass the full YouTube URL as 'url'."
                ),
            ),
            StructuredTool.from_function(
                func=self._analyze_file,
                name="analyze_file",
                description=(
                    "Read and analyze a local file (csv, xlsx, txt, json, py, "
                    "images, audio, pdf). Pass the exact local path as "
                    "'file_path'."
                ),
            ),
        ]

        self.agent = self._setup_agent()

    # ------------------------------------------------------------------ #
    # Public entry point
    # ------------------------------------------------------------------ #
    def run(self, query: str) -> str:
        """Answer a single GAIA question and return a clean, exact-match answer."""
        response = ""
        for attempt in range(3):
            try:
                result = self.agent.invoke({"input": query})
                response = result.get("output", "") if isinstance(result, dict) else str(result)
                break
            except Exception as e:
                msg = str(e)
                if "tokens per day" in msg or "(TPD)" in msg:
                    raise DailyLimitReached(msg[:300])     # let the caller switch to another model
                response = f"Error: {msg[:200]}"
                # Groq sometimes returns a malformed tool call ("Parsing failed",
                # "tool call validation failed"): just try the question again.
                print(f"attempt {attempt + 1} failed: {msg[:120]}")
        # If the agent ran out of steps without concluding, force a direct
        # best-effort answer instead of submitting the raw stop message.
        low = (response or "").lower()
        if (not response) or low.startswith("error:") or ("stopped due to" in low) or ("max iteration" in low) or (
            "iteration limit" in low
        ) or ("time limit" in low):
            try:
                direct = self.llm.invoke(
                    "Answer the following question with ONLY the final answer and "
                    "nothing else. Format: a number (no thousands separators, no "
                    "units unless explicitly asked), OR as few words as possible, "
                    "OR a comma-separated list. No explanation, no trailing "
                    f"period.\n\nQuestion: {query}"
                )
                response = getattr(direct, "content", str(direct))
            except Exception:
                pass
        return self._format_answer(response)

    # ------------------------------------------------------------------ #
    # Answer formatting (the part that wins/loses leaderboard points)
    # ------------------------------------------------------------------ #
    def _format_answer(self, text) -> str:
        if not isinstance(text, str):
            text = str(text)
        # If the model echoed the template, keep only what follows it.
        m = re.search(r"FINAL ANSWER:\s*(.*)", text, flags=re.IGNORECASE | re.DOTALL)
        if m:
            text = m.group(1)
        text = text.strip().strip('"').strip("'").strip()
        # Drop trailing punctuation/whitespace; GAIA wants the bare answer.
        text = re.sub(r"[.\s]+$", "", text)
        return text

    # ------------------------------------------------------------------ #
    # Tools
    # ------------------------------------------------------------------ #
    def _web_search(self, query: str) -> str:
        # Try Tavily first.
        if self.tavily_client is not None:
            try:
                resp = self.tavily_client.search(
                    query=query,
                    max_results=4,
                    search_depth="advanced",
                    include_answer=True,
                )
                parts = []
                if resp.get("answer"):
                    parts.append(f"Answer: {resp['answer']}")
                for r in resp.get("results", []):
                    parts.append(
                        f"[{r.get('title', '')}]({r.get('url', '')})\n"
                        f"{(r.get('content', '') or '')[:600]}"
                    )
                if parts:
                    return "\n\n".join(parts)
            except Exception as e:
                print(f"Tavily search failed, falling back to DuckDuckGo: {e}")

        # Keyless option 1: ddgs metasearch (queries several engines, so one blocked engine doesn't matter).
        try:
            from ddgs import DDGS
            hits = DDGS().text(query, max_results=6, backend="auto")
            if hits:
                return "## Search Results\n\n" + "\n".join(
                    f"[{h.get('title', '')}]({h.get('href', '')})\n{(h.get('body', '') or '')[:400]}\n" for h in hits
                )
        except Exception as e:
            print(f"ddgs search failed: {e}")

        # Keyless option 2: DuckDuckGo HTML scrape.
        try:
            result = self.web_search_tool.search(query)
            if result and not result.startswith(("Search failed", "No search results")):
                return result
        except Exception as e:
            print(f"DuckDuckGo search failed: {e}")

        # Keyless option 3: Wikipedia's own search API (works from servers).
        wiki = self._wikipedia_api(query)
        return wiki or "Search found nothing. Try a shorter or different query, or the wikipedia tool."

    def _visit_webpage(self, url: str) -> str:
        """Fetch a page and return its readable text, trimmed to fit the model's budget."""
        try:
            r = requests.get(url, headers={"User-Agent": "Mozilla/5.0 (compatible; research-agent/1.0)"}, timeout=15)
            r.raise_for_status()
            soup = BeautifulSoup(r.text, "html.parser")
            for tag in soup(["script", "style", "nav", "footer", "header", "aside"]):
                tag.decompose()
            text = re.sub(r"\n\s*\n+", "\n\n", soup.get_text("\n")).strip()
            return text[:5000] or "The page has no readable text."
        except Exception as e:
            return f"Could not open {url}: {e}"

    def _wikipedia_api(self, query: str) -> str:
        """Search Wikipedia via its public API and return the best article's text (no key needed)."""
        api = "https://en.wikipedia.org/w/api.php"
        headers = {"User-Agent": "gaia-agent/1.0 (Hugging Face Agents Course)"}
        try:
            hits = requests.get(api, params={"action": "query", "list": "search", "srsearch": query,
                                             "format": "json", "srlimit": 3}, headers=headers, timeout=15
                                ).json().get("query", {}).get("search", [])
            if not hits:
                return ""
            title = hits[0]["title"]
            pages = requests.get(api, params={"action": "query", "prop": "extracts", "explaintext": 1,
                                              "titles": title, "format": "json", "redirects": 1},
                                 headers=headers, timeout=15).json().get("query", {}).get("pages", {})
            text = next(iter(pages.values()), {}).get("extract", "")
            others = ", ".join(h["title"] for h in hits[1:])
            return f"Wikipedia: {title}\n(other matches: {others})\n\n{text[:4000]}"
        except Exception as e:
            print(f"Wikipedia API failed: {e}")
            return ""

    def _wikipedia_search(self, query: str) -> str:
        try:
            wiki = WikipediaAPIWrapper(top_k_results=3, doc_content_chars_max=3000)
            result = wiki.run(query)
            if result and result.strip():
                return result
        except Exception:
            pass
        # The wikipedia library is flaky from datacenter IPs (empty/JSON errors):
        # try Wikipedia's own API, then a general web search.
        return self._wikipedia_api(query) or self._web_search(f"{query} Wikipedia")

    def _analyze_video(self, url: str) -> str:
        """Send a YouTube URL directly to Gemini (server-side fetch, no yt-dlp)."""
        if self.genai_client is None:
            return "Video analysis is not available (no GEMINI_API_KEY). Answer from other sources if possible."
        try:
            parsed = urlparse(url)
            if not all([parsed.scheme, parsed.netloc]):
                return "Please provide a valid URL with http:// or https://."
            if "youtube.com" not in url and "youtu.be" not in url:
                return "Only YouTube videos are supported."

            prompt = (
                "Watch this video and describe it in detail. Transcribe what is "
                "said and report exactly what is shown, including any numbers, "
                "names, on-screen text, counts, or specific facts that could be "
                "needed to answer a question about it."
            )
            self.rate_limiter.acquire(blocking=True)
            resp = self.genai_client.models.generate_content(
                model=self.model_name,
                contents=types.Content(
                    parts=[
                        types.Part(file_data=types.FileData(file_uri=url)),
                        types.Part(text=prompt),
                    ]
                ),
            )
            return resp.text if getattr(resp, "text", None) else str(resp)
        except Exception as e:
            return f"Error analyzing video: {str(e)}"

    def _analyze_file(self, file_path: str) -> str:
        """Dispatch a local file to the right reader based on its extension."""
        file_path = file_path.strip().strip('"').strip("'").strip()
        if not os.path.exists(file_path):
            return f"File not found: {file_path}"

        ext = os.path.splitext(file_path)[1].lower()
        try:
            if ext == ".csv":
                return self._read_table(file_path, kind="csv")
            if ext in (".xlsx", ".xls"):
                return self._read_table(file_path, kind="excel")
            if ext in (".txt", ".json", ".py", ".md", ".tsv", ".xml", ".html"):
                with open(file_path, "r", errors="ignore") as f:
                    return f.read()[:4000]
            if ext in (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"):
                return self._gemini_media(
                    file_path,
                    "Describe this image precisely. Report any text, numbers, "
                    "objects, chess positions, or data shown.",
                )
            if ext in (".mp3", ".wav", ".m4a", ".flac", ".ogg", ".aac"):
                return self._gemini_media(
                    file_path,
                    "Transcribe this audio completely and accurately, "
                    "including every detail mentioned.",
                )
            if ext == ".pdf":
                return self._gemini_media(
                    file_path, "Extract all text and relevant content from this PDF."
                )
            # Fallback: try plain text.
            with open(file_path, "r", errors="ignore") as f:
                return f.read()[:4000]
        except Exception as e:
            return f"Error reading file {file_path}: {str(e)}"

    def _read_table(self, file_path: str, kind: str) -> str:
        import pandas as pd

        df = pd.read_csv(file_path) if kind == "csv" else pd.read_excel(file_path)
        header = f"Rows: {len(df)}, Columns: {list(df.columns)}\n\n"
        if len(df) <= 100:
            body = df.to_string()
        else:
            body = (
                "First 50 rows:\n"
                + df.head(50).to_string()
                + "\n\nSummary statistics:\n"
                + df.describe(include="all").to_string()
            )
        return (header + body)[:4000]

    def _gemini_media(self, file_path: str, prompt: str) -> str:
        """Use Gemini's native multimodal support for images/audio/pdf."""
        if self.genai_client is None:
            return "Media analysis is not available (no GEMINI_API_KEY)."
        try:
            uploaded = self.genai_client.files.upload(file=file_path)
            # Audio/PDF may need server-side processing before use.
            for _ in range(30):
                state = getattr(getattr(uploaded, "state", None), "name", "ACTIVE")
                if state == "ACTIVE":
                    break
                if state == "FAILED":
                    return "Media processing failed."
                time.sleep(2)
                uploaded = self.genai_client.files.get(name=uploaded.name)
            self.rate_limiter.acquire(blocking=True)
            resp = self.genai_client.models.generate_content(
                model=self.model_name,
                contents=[uploaded, prompt],
            )
            return resp.text if getattr(resp, "text", None) else str(resp)
        except Exception as e:
            return f"Error analyzing media file: {str(e)}"

    # ------------------------------------------------------------------ #
    # Model + agent setup
    # ------------------------------------------------------------------ #
    def _setup_llm(self):
        # Main reasoning runs on Groq's gpt-oss-120b: a strong reasoning model
        # with a 200K tokens/day budget and native function-calling (which is
        # why this agent uses a tool-calling agent, not text ReAct). Its
        # per-minute cap is ~8K tokens, so throttle to ~2.4 calls/min; ChatGroq
        # retries any throttled call. Stored on self so Gemini media calls pace
        # through it too.
        self.rate_limiter = InMemoryRateLimiter(
            requests_per_second=0.04,
            check_every_n_seconds=0.1,
            max_bucket_size=1,
        )
        return ChatGroq(
            model=self.groq_model,
            temperature=0,
            max_tokens=1536,
            max_retries=2,
            rate_limiter=self.rate_limiter,
        )

    def _setup_agent(self) -> AgentExecutor:
        system = (
            "You are a precise research assistant answering GAIA benchmark "
            "questions. Think step by step. Many questions can be solved by "
            "reasoning or from your own knowledge alone — answer those directly "
            "without calling any tool. Use a tool only when you genuinely need "
            "outside information, and never repeat the same tool call. After a "
            "search, open the most relevant result with visit_webpage when the "
            "snippet is not enough. If a "
            "question includes the content of an attached file, use it directly; "
            "if it only gives a local path, call analyze_file with that exact path. "
            "Never ask the user to provide a file. For YouTube links, use "
            "analyze_video.\n\n"
            "When you give your final answer, it MUST contain ONLY the answer "
            "and nothing else, formatted to GAIA rules: a number (no thousands "
            "separators, no units like $ or % unless the question explicitly "
            "asks for them), OR as few words as possible, OR a comma-separated "
            "list. No explanation, no sentence, no article, no trailing period."
        )
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system),
                ("human", "{input}"),
                MessagesPlaceholder("agent_scratchpad"),
            ]
        )

        agent = create_tool_calling_agent(self.llm, self.tools, prompt)

        return AgentExecutor(
            agent=agent,
            tools=self.tools,
            max_iterations=5,
            verbose=False,
            handle_parsing_errors=True,
        )


# ====================================================================== #
# Connection to app.py
# ====================================================================== #
import tempfile

API_URL = "https://agents-course-unit4-scoring.hf.space"   # same as DEFAULT_API_URL in app.py
import threading

# Groq models in order. Each has its own tokens-per-day budget, so when one runs out
# the next one takes over. The current list is at console.groq.com/docs/models.
GROQ_MODELS = ["openai/gpt-oss-120b", "openai/gpt-oss-20b", "qwen/qwen3.8-27b"]
MODEL_ID = "groq/" + GROQ_MODELS[0]          # shown in the result box; updated if the model changes

_AGENT = None    # one shared agent, so its rate limiter covers every question
_MODEL_INDEX = 0
_LOCK = threading.Lock()


def _get_agent():
    global _AGENT
    with _LOCK:
        if _AGENT is None:
            _AGENT = GeminiAgent(api_key=os.environ.get("GEMINI_API_KEY"), groq_model=GROQ_MODELS[_MODEL_INDEX])
        return _AGENT


def _next_model(failed_agent):
    """Switch to the next Groq model. Returns False if none are left."""
    global _AGENT, _MODEL_INDEX, MODEL_ID
    with _LOCK:
        if _AGENT is not failed_agent:          # another thread already switched
            return True
        if _MODEL_INDEX + 1 >= len(GROQ_MODELS):
            return False
        _MODEL_INDEX += 1
        MODEL_ID = "groq/" + GROQ_MODELS[_MODEL_INDEX]
        print(f"Daily token limit reached: switching to {GROQ_MODELS[_MODEL_INDEX]}")
        _AGENT = GeminiAgent(api_key=os.environ.get("GEMINI_API_KEY"), groq_model=GROQ_MODELS[_MODEL_INDEX])
        return True


def check_model():
    """One tiny call to Groq before the run. Returns None if it works, or the error text."""
    if not os.environ.get("GROQ_API_KEY"):
        return "GROQ_API_KEY secret is missing. Add it in Settings -> Variables and secrets, then restart the Space."
    try:
        _get_agent().llm.invoke("Reply with the word OK.")
        return None
    except Exception as e:
        return f"{type(e).__name__}: {e}"


class BasicAgent:
    def __init__(self):
        self.last_error = {}     # task_id -> error text, shown in the results table

    @staticmethod
    def _download_file(task_id):
        """Some questions come with a file (xlsx, csv, py, png, mp3...). Returns its local path or None."""
        if not task_id:
            return None
        try:
            r = requests.get(f"{API_URL}/files/{task_id}", timeout=20)
        except requests.RequestException:
            return None
        if r.status_code != 200 or not r.content:
            return None
        name = "task_file"
        m = re.search(r'filename="?([^";]+)"?', r.headers.get("content-disposition", ""))
        if m:
            name = m.group(1)
        path = os.path.join(tempfile.gettempdir(), f"{task_id}_{name}")
        with open(path, "wb") as f:
            f.write(r.content)
        return path

    @staticmethod
    def _file_context(path):
        """Put the attached file's content straight into the question, so the model
        doesn't have to decide to open it (it often didn't)."""
        agent = _get_agent()
        ext = os.path.splitext(path)[1].lower()
        if ext in (".py", ".txt", ".json", ".md", ".csv", ".tsv", ".xml", ".html", ".xlsx", ".xls"):
            content = agent._analyze_file(path)
            return (f"\n\nThe attached file ({os.path.basename(path)}) is below. Use it directly.\n"
                    f"----- FILE START -----\n{content}\n----- FILE END -----")
        if agent.genai_client is not None:      # images, audio, pdf: let Gemini read it first
            content = agent._analyze_file(path)
            return (f"\n\nThe attached file ({os.path.basename(path)}) was analysed for you:\n"
                    f"----- ANALYSIS START -----\n{content}\n----- ANALYSIS END -----")
        return (f"\n\nThe attached file is saved at this local path: {path} "
                f"(images and audio can only be read if a GEMINI_API_KEY is set).")

    def __call__(self, question, task_id=None):
        query = question
        path = self._download_file(task_id)
        if path:
            query += self._file_context(path)
        elif "attach" in question.lower():
            self.last_error[task_id] = "The question mentions an attachment, but no file could be downloaded."

        for _ in range(len(GROQ_MODELS)):
            agent = _get_agent()
            try:
                answer = agent.run(query)
                break
            except DailyLimitReached as e:
                if not _next_model(agent):
                    self.last_error[task_id] = f"All Groq models used up their daily tokens: {e}"[:300]
                    return "unknown"
            except Exception as e:
                self.last_error[task_id] = f"{type(e).__name__}: {e}"[:300]
                return "unknown"
        else:
            return "unknown"
        if answer.startswith("Error:"):
            self.last_error[task_id] = answer[:300]
            return "unknown"
        return answer or "unknown"