app.py

import os
import gradio as gr
import requests
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, TimeoutError
import re
import base64
import subprocess
import tempfile
from bs4 import BeautifulSoup
from urllib.parse import quote
import time

import wikipediaapi

from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    LiteLLMModel,
    OpenAIModel,
    Tool,
    FinalAnswerTool,
    WikipediaSearchTool,
    PythonInterpreterTool,
    InferenceClientModel
)

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# --- System Prompt ---
SYSTEM_PROMPT = """
You are a GAIA benchmark agent. Your sole objective is to produce the correct final answer.

=== OUTPUT RULES ===
- Return ONLY the final answer — no explanation, no preamble, no formatting
- Numbers: digits only (e.g. 42, not "forty-two")
- Names: exact name only
- Lists: comma-separated if required by the question
- Dates: use the format explicitly requested, or ISO (YYYY-MM-DD) if unspecified

=== REASONING STRATEGY ===
1. PLAN FIRST: Before calling any tool, reason about what information is needed and which tool is best suited.
2. DECOMPOSE: Break multi-step questions into sub-tasks. Solve each sub-task before combining into a final answer.
3. VERIFY: Cross-check answers using a second tool or source when feasible. Prefer primary/authoritative sources.
4. DEDUCE: If full content is unavailable (paywalls, broken links, restricted video), use titles, descriptions, metadata, and search snippets to logically infer the answer.

=== TOOL USAGE RULES ===
1. NO REPEAT CALLS: Never call the same tool with the same arguments twice. If a call fails or returns empty, move on.
2. FALLBACK CHAIN: Specialized tool fails → general web search → page fetch → deduce from context.
3. ONE FALLBACK: Use each fallback strategy exactly once per sub-task.
4. STOP LOOPING: If after 3 distinct tool attempts the answer is still unclear, make your best-reasoned guess and output it immediately.

=== FAILURE RECOVERY ===
- Stuck or hitting max steps? Output your single best answer NOW and stop.
- Partial information is enough — reason from what you have.
- An educated, well-reasoned guess beats silence or an infinite loop.
"""

WEB_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/122 Safari/537.36",
    "Accept-Language": "en-US,en;q=0.9",
    "Referer": "https://www.google.com/"
}

# MODEL_NAME = "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free"
# API_BASE="https://openrouter.ai/api/v1"

MODEL_NAME = "gemini-2.5-flash"
API_BASE = "https://generativelanguage.googleapis.com/v1beta/openai/"

# ─────────────────────────────────────────────────────────────
# Tool 2 – YouTube transcript
# ─────────────────────────────────────────────────────────────
class YouTubeTranscriptTool(Tool):
    name = "youtube_transcript"
    description = (
        "Fetch the transcript (captions) of a YouTube video. "
        "Use whenever a question references a YouTube URL and asks about "
        "video content, dialogue, or what is shown."
    )
    inputs = {"url": {"type": "string", "description": "Full YouTube video URL"}}
    output_type = "string"

    def forward(self, url: str) -> str:
        m = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{11})", url)
        if not m:
            return f"Could not extract video ID from: {url}"
        vid = m.group(1)
        headers = {"User-Agent": "Mozilla/5.0"}

        # Primary: youtubetotranscript.com
        try:
            r = requests.get(
                f"https://youtubetotranscript.com/transcript?v={vid}&lang=en",
                headers=headers, timeout=20,
            )
            if r.status_code == 200 and len(r.text.strip()) > 50:
                return r.text[:12000]
        except Exception:
            pass

        # Fallback: public proxy JSON API
        try:
            r = requests.get(
                f"https://api.youtubetranscript.com/?video_id={vid}",
                headers=headers, timeout=20,
            )
            if r.status_code == 200:
                segs = r.json()
                text = " ".join(s.get("text", "") for s in segs if isinstance(s, dict))
                if text:
                    return text[:12000]
        except Exception:
            pass

        # Fallback 2: yt-dlp subtitle fetch (if installed in the Space)
        try:
            result = subprocess.run(
                ["yt-dlp", "--skip-download", "--write-auto-sub",
                 "--sub-lang", "en", "--sub-format", "vtt",
                 "-o", "/tmp/ytsub", url],
                capture_output=True, text=True, timeout=30,
            )
            for f in ["/tmp/ytsub.en.vtt", "/tmp/ytsub.en-US.vtt"]:
                if os.path.exists(f):
                    raw = open(f).read()
                    # Strip VTT formatting
                    clean = re.sub(r"<[^>]+>", "", raw)
                    clean = re.sub(r"\d{2}:\d{2}.*\n", "", clean)
                    clean = re.sub(r"\n{2,}", "\n", clean).strip()
                    return clean[:12000]
        except Exception:
            pass

        return (
            f"Could not retrieve transcript for video {vid}. "
            "Try searching for the video title + 'transcript' via web_search."
        )


# ─────────────────────────────────────────────────────────────
# Tool 3 – Safe webpage visitor (strips HTML tags)
# ─────────────────────────────────────────────────────────────
class SafeVisitWebpageTool(Tool):
    name = "visit_webpage"
    description = "Fetch a webpage and return its readable plain-text content."
    inputs = {"url": {"type": "string", "description": "URL to fetch"}}
    output_type = "string"

    def forward(self, url: str) -> str:

        try:
            r = requests.get(url, headers=WEB_HEADERS, timeout=15, allow_redirects=True)
            if r.status_code != 200:
                return f"HTTP {r.status_code}"
            text = re.sub(r"<[^>]+>", " ", r.text)
            text = re.sub(r"[ \t]{2,}", " ", text)
            text = re.sub(r"\n{3,}", "\n\n", text).strip()
            return text[:10000]
        except Exception as e:
            return f"ERROR: {e}"


# ─────────────────────────────────────────────────────────────
# Tool 4 – Image analyser (vision via Gemini multimodal)
# Handles: chess boards, diagrams, any attached image
# ─────────────────────────────────────────────────────────────
class AnalyzeImageTool(Tool):
    name = "analyze_image"
    description = (
        "Download an image file (by URL or local path) and analyse its content "
        "using a vision-capable model. Use for chess positions, diagrams, photos, etc."
    )
    inputs = {
        "source": {
            "type": "string",
            "description": "URL or local file path of the image (PNG/JPG/GIF/WEBP)",
        },
        "question": {
            "type": "string",
            "description": "What to ask about the image, e.g. 'What is the best next move for black?'",
        },
    }
    output_type = "string"

    def forward(self, source: str, question: str) -> str:
        # Load image bytes
        try:
            if source.startswith("http"):
                r = requests.get(source, timeout=15)
                r.raise_for_status()
                img_bytes = r.content
                mime = r.headers.get("Content-Type", "image/png").split(";")[0]
            else:
                with open(source, "rb") as f:
                    img_bytes = f.read()
                ext = source.rsplit(".", 1)[-1].lower()
                mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg",
                        "png": "image/png", "gif": "image/gif",
                        "webp": "image/webp"}.get(ext, "image/png")
        except Exception as e:
            return f"Could not load image: {e}"

        b64 = base64.b64encode(img_bytes).decode()

        # Call Gemini vision via OpenAI-compatible endpoint
        api_key = os.getenv("API_KEY", "")
        payload = {
            "model": MODEL_NAME,
            "max_tokens": 1024,
            "messages": [{
                "role": "user",
                "content": [
                    {"type": "image_url",
                     "image_url": {"url": f"data:{mime};base64,{b64}"}},
                    {"type": "text", "text": question},
                ],
            }],
        }
        try:
            resp = requests.post(
                "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
                headers={"Authorization": f"Bearer {api_key}",
                         "Content-Type": "application/json"},
                json=payload, timeout=30,
            )
            resp.raise_for_status()
            return resp.json()["choices"][0]["message"]["content"].strip()
        except Exception as e:
            return f"Vision API error: {e}"


# ─────────────────────────────────────────────────────────────
# Tool 5 – Audio transcription (MP3/WAV → text via Gemini)
# Handles: voice memos, lecture recordings, pie recipes, etc.
# ─────────────────────────────────────────────────────────────
class TranscribeAudioTool(Tool):
    name = "transcribe_audio"
    description = (
        "Download an audio file (MP3/WAV/M4A) and transcribe its speech to text. "
        "Use for questions that reference attached audio recordings or voice memos."
    )
    inputs = {
        "source": {
            "type": "string",
            "description": "URL or local file path of the audio file",
        },
    }
    output_type = "string"

    def forward(self, source: str) -> str:
        # Load audio bytes
        try:
            if source.startswith("http"):
                r = requests.get(source, timeout=30)
                r.raise_for_status()
                audio_bytes = r.content
                ext = source.split("?")[0].rsplit(".", 1)[-1].lower() or "mp3"
            else:
                with open(source, "rb") as f:
                    audio_bytes = f.read()
                ext = source.rsplit(".", 1)[-1].lower()
        except Exception as e:
            return f"Could not load audio: {e}"

        mime_map = {"mp3": "audio/mpeg", "wav": "audio/wav",
                    "m4a": "audio/mp4", "ogg": "audio/ogg",
                    "flac": "audio/flac"}
        mime = mime_map.get(ext, "audio/mpeg")
        b64 = base64.b64encode(audio_bytes).decode()

        # Primary: Gemini multimodal audio understanding
        api_key = os.getenv("API_KEY", "")
        payload = {
            "model": MODEL_NAME,
            "max_tokens": 2048,
            "messages": [{
                "role": "user",
                "content": [
                    {"type": "input_audio",
                     "input_audio": {"data": b64, "format": ext}},
                    {"type": "text",
                     "text": (
                         "Please transcribe this audio recording completely and accurately. "
                         "Include every word spoken. Return only the transcription, nothing else."
                     )},
                ],
            }],
        }
        try:
            resp = requests.post(
                "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
                headers={"Authorization": f"Bearer {api_key}",
                         "Content-Type": "application/json"},
                json=payload, timeout=60,
            )
            if resp.status_code == 200:
                return resp.json()["choices"][0]["message"]["content"].strip()
        except Exception:
            pass

        return "Could not transcribe audio"


# ─────────────────────────────────────────────────────────────
# Tool 7 – Excel file reader
# Handles: questions with an attached .xlsx file
# ─────────────────────────────────────────────────────────────
class ReadExcelFileTool(Tool):
    name = "read_excel_file"
    description = (
        "Download an Excel (.xlsx) file and return its contents as a markdown table. "
        "Use when a question involves data from an attached spreadsheet."
    )
    inputs = {
        "source": {
            "type": "string",
            "description": "URL or local file path of the .xlsx file",
        },
        "sheet": {
            "type": "string",
            "description": "Sheet name or index to read (default: first sheet)",
            "default": "0","nullable": True,
        },
    }
    output_type = "string"

    def forward(self, source: str, sheet: str = "0") -> str:
        try:
            if source.startswith("http"):
                r = requests.get(source, timeout=15)
                r.raise_for_status()
                with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as tmp:
                    tmp.write(r.content)
                    tmp_path = tmp.name
            else:
                tmp_path = source

            # Parse sheet argument
            try:
                sheet_arg = int(sheet)
            except ValueError:
                sheet_arg = sheet

            df = pd.read_excel(tmp_path, sheet_name=sheet_arg)

            if source.startswith("http"):
                os.unlink(tmp_path)

            # Return as markdown table (truncated if huge)
            md = df.to_markdown(index=False)
            if len(md) > 10_000:
                # Also return column summary
                summary = f"Shape: {df.shape}\nColumns: {list(df.columns)}\n\n"
                return summary + md[:10_000] + "\n... (truncated)"
            return md

        except Exception as e:
            return f"Excel read error: {e}"


# ─────────────────────────────────────────────────────────────
# Tool 8 – Task file downloader (fetches attached files from scoring API)
# ─────────────────────────────────────────────────────────────
class DownloadTaskFileTool(Tool):
    name = "download_task_file"
    description = (
        "Download an attached file for a GAIA task from the scoring API and save it locally. "
        "Returns the local file path. Use this first whenever a question mentions an attached file "
        "(image, audio, Python script, Excel spreadsheet, etc.), then pass the path to the "
        "appropriate analysis tool (analyze_image, transcribe_audio, run_python_file, read_excel_file)."
    )
    inputs = {
        "task_id": {
            "type": "string",
            "description": "The task_id of the GAIA question whose file you want to download",
        },
        "file_name": {
            "type": "string",
            "description": "The file_name field from the question metadata (e.g. 'abc123.png')",
        },
    }
    output_type = "string"

    def forward(self, task_id: str, file_name: str) -> str:
        url = f"{DEFAULT_API_URL}/files/{task_id}"
        try:
            r = requests.get(url, timeout=30)
            r.raise_for_status()
        except Exception as e:
            return f"Could not download file for task {task_id}: {e}"

        ext = file_name.rsplit(".", 1)[-1] if "." in file_name else "bin"
        with tempfile.NamedTemporaryFile(
            suffix=f".{ext}", prefix=f"gaia_{task_id}_",
            delete=False, dir="/tmp",
        ) as tmp:
            tmp.write(r.content)
            return tmp.name

class BasicAgent:
    def __init__(self):

        model = InferenceClientModel("Qwen/Qwen2.5-72B-Instruct")
        
        # model = LiteLLMModel(
        #     model_id="groq/llama-3.3-70b-versatile",
        #     api_key=os.environ.get("GROQ_API_KEY"),
        #     max_tokens=1024,
        #     temperature=0.0, 
        # )

        # model = OpenAIModel(
        #     model_id="meta-llama/llama-4-scout-17b-16e-instruct",
        #     api_key=os.getenv("GROQ_API_KEY"),
        #     api_base="https://api.groq.com/openai/v1"
        # )

        # model = OpenAIModel(
        #     model_id=MODEL_NAME,
        #     api_key=os.getenv("API_KEY"),
        #     api_base=API_BASE
        # )

        # ─────────────────────────────
        # SUB-AGENT
        # ─────────────────────────────
        web_agent = ToolCallingAgent(
            tools=[DuckDuckGoSearchTool(),
                    WikipediaSearchTool(
                                user_agent="MySmolAgentApp/1.0 (contact@example.com)",
                                language="en",
                                content_type="text",  # Keep it short to save LLM tokens
                                extract_format="WIKI"
                            ),                    
                    SafeVisitWebpageTool()
            ],
            model=model,
            max_steps=10,
            name="web_agent",
            description="Use ONLY for web_search, wikipedia and browsing. Always prefer wikipedia for factual queries with known entities.")
        

        # ─────────────────────────────
        # MAIN AGENT
        # ─────────────────────────────
        self.agent = CodeAgent(
            tools=[
                    PythonInterpreterTool(),
                    
                    YouTubeTranscriptTool(),
                    
                    SafeVisitWebpageTool(),
                    
                    AnalyzeImageTool(),
                    
                    TranscribeAudioTool(),
                    
                    ReadExcelFileTool(),
                    
                    DownloadTaskFileTool(),
                    WikipediaSearchTool(
                                user_agent="MySmolAgentApp/1.0 (contact@example.com)",
                                language="en",
                                content_type="text",  # Keep it short to save LLM tokens
                                extract_format="WIKI"
                            ),    
            ],
            model=model,
            managed_agents=[web_agent],

            additional_authorized_imports=[
                "re", "json", "math", "collections", "pandas",
                "datetime", "statistics", "base64", "os",
            ],
            max_steps=10,
            verbosity_level=1 
        )

    def __call__(self, question: str, task_id: str = "", file_name: str = ""):

        prompt = SYSTEM_PROMPT + f"\n\nQuestion: {question}"

        if task_id and file_name:
            prompt += f"""

ATTACHED FILE DETECTED:
- task_id: {task_id}
- file_name: {file_name}

INSTRUCTION:
1. First call download_task_file(task_id, file_name)
2. Then route file to correct tool:
   - image → AnalyzeImageTool
   - audio → TranscribeAudioTool
   - xlsx → ReadExcelFileTool
   - py → PythonInterpreterTool
"""

        return self.agent.run(prompt)
        

def run_and_submit_all( profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code

    if profile:
        username= f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Instantiate Agent ( modify this part to create your agent)
    try:
        agent = BasicAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
             print("Fetched questions list is empty.")
             return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
         print(f"Error decoding JSON response from questions endpoint: {e}")
         print(f"Response text: {response.text[:500]}")
         return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    # 3. Run your Agent
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            submitted_answer = agent(question_text)
            time.sleep(15)
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
        except Exception as e:
             print(f"Error running agent on task {task_id}: {e}")
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})

    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Prepare Submission 
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)

    # 5. Submit
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df


# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**
        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
        ---
        **Disclaimers:**
        Once clicking on the "submit button, it can take quite some time ( this is the time for the agent to go through all the questions).
        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance for the delay process of the submit button, a solution could be to cache the answers and submit in a seperate action or even to answer the questions in async.
        """
    )

    gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    # Removed max_rows=10 from DataFrame constructor
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table]
    )

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    # Check for SPACE_HOST and SPACE_ID at startup for information
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")

    if space_id_startup: # Print repo URLs if SPACE_ID is found
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")

    print("-"*(60 + len(" App Starting ")) + "\n")

    print("Launching Gradio Interface for Basic Agent Evaluation...")
    demo.launch(debug=True, share=False)