app.py

import os
import re
import gradio as gr
import requests
import pandas as pd
import torch
from transformers import pipeline
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.prompts import ChatPromptTemplate
from langchain.prompts import PromptTemplate
from langchain_huggingface import HuggingFacePipeline
from langchain_core.output_parsers import StrOutputParser
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated, List
from langchain_community.document_loaders.youtube import YoutubeLoader
import numexpr

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
SYSTEM_PROMPT = """You are a helpful assistant tasked with answering questions.

You have access to a set of tools to help you. The question you receive may require you to use these tools.
When you receive a question, you should first think about what steps you need to take.
Based on your plan, you can then call the necessary tools.
After calling a tool, you will get a result. You should analyze the result and decide if you need to call another tool or if you have enough information to answer the question.

When you have the final answer, you must output it in the following format:
FINAL ANSWER: [YOUR FINAL ANSWER]

YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma-separated list of numbers and/or strings.
- If you are asked for a number, do not use commas for thousands separators or units like '$' or '%' unless specified.
- If you are asked for a string, do not use articles or abbreviations (e.g., for cities).
- If you are asked for a comma-separated list, apply the above rules to each element.

Example:
Question: What is the capital of France?
Your thought process: I need to find the capital of France. I will use the web search tool.
Tool call: web_search("capital of France")
Tool output: Paris is the capital of France.
Your final answer: FINAL ANSWER: Paris
"""

# --- Tool Definitions ---


@tool
def web_search(query: str):
    """Searches the web using DuckDuckGo."""
    print(f"--- Calling Web Search Tool with query: {query} ---")
    search = DuckDuckGoSearchRun()
    return search.run(query)


@tool
def math_calculator(expression: str):
    """Calculates the result of a mathematical expression."""
    print(f"--- Calling Math Calculator Tool with expression: {expression} ---")
    try:
        # Use numexpr for safe evaluation
        result = numexpr.evaluate(expression).item()
        return result
    except Exception as e:
        return f"Error evaluating expression: {e}"


@tool
def image_analyzer(image_url: str):
    """Analyzes an image and returns a description."""
    print(f"--- Calling Image Analyzer Tool with URL: {image_url} ---")
    try:
        # Using a CPU-friendly image-to-text model
        image_to_text = pipeline(
            "image-to-text", model="Salesforce/blip-image-captioning-base"
        )
        description = image_to_text(image_url)[0]["generated_text"]
        return description
    except Exception as e:
        return f"Error analyzing image: {e}"


@tool
def youtube_transcript_reader(youtube_url: str):
    """Reads the transcript of a YouTube video."""
    print(f"--- Calling YouTube Transcript Reader Tool with URL: {youtube_url} ---")
    try:
        loader = YoutubeLoader.from_youtube_url(youtube_url, add_video_info=False)
        docs = loader.load()
        transcript = " ".join([doc.page_content for doc in docs])
        # Return a manageable chunk of the transcript
        return transcript[:4000]
    except Exception as e:
        return f"Error reading YouTube transcript: {e}"


# --- Agent State Definition ---
class AgentState(TypedDict):
    question: str
    messages: Annotated[list, lambda x, y: x + y]
    sender: str


# --- LangGraph Agent Definition ---
class GaiaAgent:
    def __init__(self):
        print("Initializing GaiaAgent...")
        self.tools = [
            web_search,
            math_calculator,
            image_analyzer,
            youtube_transcript_reader,
        ]

        # Initialize the LLM
        print("Loading LLM...")
        llm = HuggingFacePipeline.from_model_id(
            model_id="HuggingFaceH4/zephyr-7b-beta",
            task="text-generation",
            pipeline_kwargs={
                "max_new_tokens": 512,
                "top_k": 50,
                "temperature": 0.1,
                "do_sample": False,
                "torch_dtype": torch.bfloat16,
                "device_map": "auto",
            },
        )
        print("LLM loaded.")

        # Create the agent graph
        prompt = PromptTemplate(
            template=SYSTEM_PROMPT
            + """
Here is the current conversation:
{messages}

Question: {question}
""",
            input_variables=["messages", "question"],
        )

        self.agent = prompt | llm | StrOutputParser()
        self.graph = self._create_graph()
        print("GaiaAgent initialized.")

    def _create_graph(self):
        graph = StateGraph(AgentState)
        graph.add_node("agent", self._call_agent)
        graph.add_node("tools", self._call_tools)
        graph.add_conditional_edges(
            "agent", self._decide_action, {"tools": "tools", END: END}
        )
        graph.add_edge("tools", "agent")
        graph.set_entry_point("agent")
        return graph.compile()

    def _call_agent(self, state: AgentState):
        print("--- Calling Agent ---")
        message_history = "\n".join(state["messages"])
        response = self.agent.invoke(
            {"messages": message_history, "question": state["question"]}
        )
        return {"messages": [response], "sender": "agent"}

    def _decide_action(self, state: AgentState):
        print("--- Deciding Action ---")
        response = state["messages"][-1]
        if "FINAL ANSWER:" in response:
            return END
        else:
            return "tools"

    def _call_tools(self, state: AgentState):
        print("--- Calling Tools ---")
        raw_tool_call = state["messages"][-1]

        # Simple regex to find tool calls like tool_name("argument")
        tool_call_match = re.search(r"(\w+)\((.*?)\)", raw_tool_call)
        if not tool_call_match:
            return {"messages": ["No valid tool call found."], "sender": "tools"}

        tool_name = tool_call_match.group(1).strip()
        tool_input_str = tool_call_match.group(2).strip()

        # Remove quotes from the input string if they exist
        if tool_input_str.startswith('"') and tool_input_str.endswith('"'):
            tool_input = tool_input_str[1:-1]
        else:
            tool_input = tool_input_str

        tool_to_call = next((t for t in self.tools if t.name == tool_name), None)

        if tool_to_call:
            try:
                result = tool_to_call.run(tool_input)
                return {"messages": [str(result)], "sender": "tools"}
            except Exception as e:
                return {
                    "messages": [f"Error executing tool {tool_name}: {e}"],
                    "sender": "tools",
                }
        else:
            return {"messages": [f"Tool '{tool_name}' not found."], "sender": "tools"}

    def __call__(self, question: str) -> str:
        print(f"Agent received question: {question[:100]}...")

        initial_state = {"question": question, "messages": [], "sender": "user"}

        final_state = self.graph.invoke(initial_state, {"recursion_limit": 10})

        final_answer = final_state["messages"][-1]

        # Extract the answer after "FINAL ANSWER:"
        match = re.search(
            r"FINAL ANSWER:\s*(.*)", final_answer, re.IGNORECASE | re.DOTALL
        )
        if match:
            extracted_answer = match.group(1).strip()
            print(f"Agent returning final answer: {extracted_answer}")
            return extracted_answer
        else:
            print("Agent could not find a final answer in the required format.")
            # Return a fallback answer if parsing fails
            return "Could not determine the final answer."


def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the GaiaAgent on them, submits all answers,
    and displays the results.
    """
    if not profile:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    username = profile.username
    print(f"User logged in: {username}")

    space_id = os.getenv("SPACE_ID")
    if not space_id:
        return "SPACE_ID environment variable is not set. Cannot proceed.", None

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

    # 1. Instantiate Agent
    try:
        agent = GaiaAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(f"Agent code URL: {agent_code}")

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=20)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        return f"Error 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:
            continue

        try:
            submitted_answer = agent(question_text)
            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:
        return "Agent did not produce any answers.", pd.DataFrame(results_log)

    # 4. Prepare and Submit
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    print(f"Submitting {len(answers_payload)} answers for user '{username}'...")
    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.")
        return final_status, pd.DataFrame(results_log)
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}. Detail: {e.response.text}"
        return f"Submission Failed: {error_detail}", pd.DataFrame(results_log)
    except Exception as e:
        return f"An unexpected error occurred during submission: {e}", pd.DataFrame(
            results_log
        )


# --- Build Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**
        1.  This Space contains a `langgraph`-based agent equipped with tools for web search, math, image analysis, and YouTube transcript reading.
        2.  Log in to your Hugging Face account using the button below. Your HF username is used for the submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch the questions, run the agent, submit the answers, and see your score.
        ---
        **Disclaimer:**
        -   Once you click the submit button, please be patient. The agent needs time to process all the questions, which can take several minutes depending on the model and hardware.
        """
    )

    gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(
        label="Run Status / Submission Result", lines=5, interactive=False
    )
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

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

if __name__ == "__main__":
    print("\n" + "-" * 30 + " App Starting " + "-" * 30)
    demo.launch(debug=True, share=False)