app.py

import os
import gradio as gr
import requests
import pandas as pd
import time
from dotenv import load_dotenv
from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, tool

# Load environment variables
try:
    load_dotenv()
except Exception:
    pass

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

# --- Tool Definitions ---

@tool
def reverse_text(text: str) -> str:
    """Reverses the given text string. Useful for tasks requiring backward reading.

    Args:
        text: The string to be reversed.
    """
    return text[::-1]

@tool
def calculator(expression: str) -> str:
    """Evaluates a mathematical expression (e.g., '2 * 3 + 5') safely.

    Args:
        expression: The mathematical expression string to evaluate.
    """
    try:
        return str(eval(expression, {"__builtins__": {}}))
    except Exception as e:
        return f"Error: {e}"

@tool
def download_file(task_id: str) -> str:
    """Downloads a file associated with a Task ID from the evaluation API.

    Args:
        task_id: The unique identifier for the task to fetch the file for.
    """
    try:
        api_url = "https://agents-course-unit4-scoring.hf.space"
        files_url = f"{api_url}/files/{task_id}"
        response = requests.get(files_url, timeout=30)
        response.raise_for_status()
        return response.text
    except Exception as e:
        return f"Error downloading file: {e}"

# --- Agent Definition ---

class GAIAAgent:
    def __init__(self, verbose=False):
        self.verbose = verbose
        
        # 1. GET GEMINI API KEY
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            print("ERROR: GEMINI_API_KEY not found in environment variables.")
        
        # 2. SET GEMINI MODEL ID
        # FIX: Use the specific stable version 'gemini-1.5-flash-001'
        # The API often rejects 'latest' or base aliases on the v1beta endpoint.
        model_id = "gemini/gemini-2.5-flash-preview-09-2025"
            
        if self.verbose:
            print(f"Initializing Agent with model: {model_id}")

        # 3. INITIALIZE LITELLM MODEL
        self.model = LiteLLMModel(
            model_id=model_id,
            api_key=api_key,
        )

        # Initialize Tools
        self.tools = [
            DuckDuckGoSearchTool(), 
            reverse_text, 
            calculator, 
            download_file
        ]

        # Authorized imports for the CodeAgent
        authorised_imports = [
            "requests", "bs4", "pandas", "numpy", "scipy", "matplotlib", 
            "seaborn", "sklearn", "nltk", "PIL", "cv2", "re", "math", "time"
        ]

        self.agent = CodeAgent(
            tools=self.tools,
            model=self.model, 
            add_base_tools=True,
            planning_interval=3,
            verbosity_level=2 if self.verbose else 0,
            additional_authorized_imports=authorised_imports
        )

    def _is_reversed_text(self, text):
        """Check if the text appears to be reversed"""
        return (text.strip().startswith(".") or 
                "rewsna" in text.lower() or 
                "esaelp" in text.lower())

    def __call__(self, question: str) -> str:
        """Process a question and return the answer"""
        
        # Handle reversed text logic manually before sending to agent
        if self._is_reversed_text(question):
            if self.verbose: print("Detected reversed text. Decoding...")
            question = question[::-1]

        # Basic Prompt Template
        prompt = f"""
        QUESTION: {question}
        
        INSTRUCTIONS:
        1. Solve the problem step-by-step using Python code.
        2. If you need information, search for it.
        3. If you need to calculate something, write code to do it.
        4. If you need a file, use download_file with the task_id provided.
        5. FINAL ANSWER: Print the final result strictly as the last line of your output.
        """

        try:
            # Run the agent
            answer = self.agent.run(prompt)
            return str(answer)
        except Exception as e:
            return f"Error processing question: {e}"

# --- Gradio Logic ---

def run_and_submit_all(profile: gr.OAuthProfile | None, sample_size=0):
    # Determine space ID
    space_id = os.getenv("SPACE_ID") or "generic/agent-runner"
    
    # Check login
    if profile is None:
        return "Please Login to Hugging Face first!", pd.DataFrame()
    
    username = profile.username
    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # Initialize Agent
    try:
        agent = GAIAAgent(verbose=True)
    except Exception as e:
        return f"Error initializing agent: {e}", pd.DataFrame()

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # Fetch Questions
    try:
        print("Fetching questions...")
        response = requests.get(questions_url, timeout=15)
        questions_data = response.json()
    except Exception as e:
        return f"Error fetching questions: {e}", pd.DataFrame()

    # Optional Sampling
    if sample_size > 0:
        questions_data = questions_data[:int(sample_size)]

    results_log = []
    answers_payload = []

    print(f"Running on {len(questions_data)} questions...")
    
    for i, item in enumerate(questions_data):
        task_id = item.get("task_id")
        q_text = item.get("question")
        
        if not task_id or not q_text:
            continue

        try:
            print(f"--- Task {i+1} ---")
            ans = agent(q_text)
            answers_payload.append({"task_id": task_id, "submitted_answer": ans})
            results_log.append({"Task ID": task_id, "Question": q_text, "Answer": ans})
            
            # Rate limit protection (Wait 2s between questions)
            time.sleep(2) 
        except Exception as e:
            results_log.append({"Task ID": task_id, "Error": str(e)})

    # Submit
    submission = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers_payload
    }
    
    try:
        print("Submitting answers...")
        r = requests.post(submit_url, json=submission)
        r.raise_for_status()
        return f"Success! {r.json()}", pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission failed: {e}", pd.DataFrame(results_log)

def test_single(question):
    agent = GAIAAgent(verbose=True)
    return agent(question)

# --- UI ---
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent Runner (Gemini Powered)")
    gr.LoginButton()
    
    with gr.Tab("Test"):
        q_in = gr.Textbox(label="Enter Question")
        btn_test = gr.Button("Run")
        out_test = gr.Textbox(label="Answer")
        btn_test.click(test_single, q_in, out_test)
        
    with gr.Tab("Evaluate"):
        slider = gr.Slider(0, 20, value=0, label="Sample Size (0=All)")
        btn_eval = gr.Button("Run & Submit")
        out_status = gr.Textbox(label="Status")
        out_df = gr.DataFrame(label="Results")
        
        btn_eval.click(run_and_submit_all, inputs=[slider], outputs=[out_status, out_df])

if __name__ == "__main__":
    demo.launch()