app.py

from smolagents import CodeAgent, HfApiModel, OpenAIServerModel

from Gradio_UI import GradioUI
import yaml
import os
import requests
import pandas as pd
import gradio as gr
import time

# Import tool CLASSES from the src directory
from src.final_answer_tool import FinalAnswerTool
from src.web_browsing_tool import WebBrowser
from src.file_processing_tool import FileIdentifier
from src.image_processing_tool import ImageProcessor
from src.markdown_table_parser import MarkdownTableParserTool # Updated
from src.python_tool import CodeExecutionTool
from src.speech_to_text import SpeechToTextTool # Updated
from src.spreadsheet_tool import SpreadsheetTool
from src.text_reversal_tool import TextReversalTool
from src.video_processing_tool import VideoProcessingTool

# (Keep Constants as is)
# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# --- Basic Agent Definition ---

# Enhanced Phase 1: Lightweight Model and Token Management for HF Spaces
try:
    # Try OpenAI first (if API key available) - Use mini version for better token management
    model = OpenAIServerModel(
        model_id="gpt-4o-mini",  # Use mini version for better token management
        api_base="https://api.openai.com/v1",
        api_key=os.environ.get("OPENAI_API_KEY"),
        max_tokens=2000,  # Increased from 1000 for better reasoning capability
        temperature=0.1,  # Lower temperature for more consistent outputs
    )
    print("Using OpenAI gpt-4o-mini model")
except Exception as e:
    print(f"OpenAI model initialization failed: {e}")
    # Fallback to HF model - More capable than DialoGPT-medium
    try:
        model = HfApiModel(
            model_id="microsoft/DialoGPT-large",  # Upgraded from medium for better capability
            max_tokens=2000,
            temperature=0.1,
            custom_role_conversions=None,
        )
        print("Using fallback HF DialoGPT-large model")
    except Exception as fallback_error:
        print(f"Fallback model initialization failed: {fallback_error}")
        # Final fallback to basic HF model
        model = HfApiModel(
            max_tokens=2000,
            temperature=0.1,
        )
        print("Using basic HF model as final fallback")

# Instantiate Tools
final_answer_tool = FinalAnswerTool()
web_browsing_tool = WebBrowser() 
file_processing_tool = FileIdentifier() 
image_processing_tool = ImageProcessor() 
markdown_parser_tool = MarkdownTableParserTool() # Updated
python_tool = CodeExecutionTool() 
speech_to_text_tool = SpeechToTextTool() # Updated
spreadsheet_tool = SpreadsheetTool() 
text_reversal_tool = TextReversalTool()
video_processing_tool = VideoProcessingTool() 

# Add debug prints for file paths
print("Current directory:", os.getcwd())
print("prompts.yaml exists:", os.path.exists("prompts.yaml"))

# Load Prompts
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
        print("Loaded prompts.yaml successfully. Structure:", type(prompt_templates))  # Debug
        if isinstance(prompt_templates, dict):
            print("Keys:", prompt_templates.keys())  # Debug
        else:
            print("Loaded prompt_templates is not a dictionary.")
except FileNotFoundError:
    print("Error: prompts.yaml not found. Using default templates.")
    prompt_templates = {
        "system_prompt": { # This was a single string, now a dict
            "base": "You are an expert assistant...", # Default value
            "with_tools": "At each step...", # Default value
        },
        "system": { # This section was already a dict, kept for consistency
            "base": "You are a GAIA benchmark agent running in HF Spaces. Be concise and efficient in your responses.",
            "with_tools": "Think briefly, act decisively. Use tools efficiently to solve GAIA benchmark tasks."
        },
        "human": {
            "base": "Here is your task: {{task}}\\\\nProvide exact answer. Be concise and efficient.", # Updated base
            "with_tools": "Here is your task: {{task}}\\\\nUse available tools strategically. Be direct and resource-conscious: {{tools}}" # Updated with_tools
        },
        "planning": {
            "initial_facts": "Task: {{task}}. Identify key facts and missing information concisely.",
            "initial_plan": "Develop an efficient 3-5 step plan for this GAIA task using available tools."
            # etc...
        },
        "managed_agent": {
            "task": "Managed agent task: {{task}}",
            "report": "Managed agent report: {{final_answer}}"
        },
        "final_answer": { 
            "base": "The final answer is: {{answer}}"
        }
        # Include all other required sections as per your YAML structure if they exist
    }
except yaml.YAMLError as e:
    print(f"Error parsing prompts.yaml: {e}")
    print("Using default templates optimized for HF Spaces")
    prompt_templates = {
        "system_prompt": "You are a helpful AI assistant. Please be concise and efficient.",
        "system": {
            "base": "You are a GAIA benchmark agent running in HF Spaces. Be concise and efficient in your responses.",
            "with_tools": "Think briefly, act decisively. Use tools efficiently to solve GAIA benchmark tasks."
        },
        "human": {
            "base": "GAIA Task: {{task}}\\\\nProvide exact answer. Be concise and efficient.",
            "with_tools": "GAIA Task: {{task}}\\\\nUse available tools strategically. Be direct and resource-conscious: {{tools}}"
        },
        "planning": {
            "initial_facts": "Task: {{task}}. Identify key facts and missing information concisely.",
            "initial_plan": "Develop an efficient 3-5 step plan for this GAIA task using available tools."
        },
        "managed_agent": {
            "task": "Managed agent task: {{task}}",
            "report": "Managed agent report: {{final_answer}}"
        },
        "final_answer": { # Placeholder, structure might need refinement based on agent's specific use
            "base": "The final answer is: {{answer}}"
        }
    }


# Enhanced agent configuration for HF Spaces optimization
class EnhancedCodeAgent(CodeAgent):
    def __call__(self, question: str) -> str:
        try:
            response = self.run(question)
            return response
        except Exception as e:
            print(f"Agent execution error: {e}")
            # Provide a graceful fallback response
            return f"I encountered an issue while processing your request. Here's what I know: {str(e)}"
            
# Create the Agent
agent_tools = [
    final_answer_tool,
    web_browsing_tool,
    file_processing_tool,
    image_processing_tool,
    markdown_parser_tool, # Updated
    python_tool,
    speech_to_text_tool, # Updated
    spreadsheet_tool,
    text_reversal_tool,
    video_processing_tool
]

# Flatten system_prompt if it's a dict (e.g., from YAML)
if isinstance(prompt_templates.get("system_prompt"), dict):
    # Use the 'main' variant by default
    prompt_templates["system_prompt"] = prompt_templates["system_prompt"].get("main", "")


# Enhanced agent configuration for HF Spaces optimization
agent = EnhancedCodeAgent(
    model=model,
    tools=agent_tools,
    max_steps=8,  # Increased from 5 to handle multi-step reasoning while staying efficient
    verbosity_level=1,  # Keep some verbosity for debugging in HF Spaces
    name="GAIAAgent",  # Updated name to reflect GAIA benchmark focus
    description="Efficient GAIA benchmark agent optimized for HF Spaces with enhanced token management",
    prompt_templates=prompt_templates
)

def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the agent on them, submits answers,
    and displays the results.
    """
    space_id = os.getenv("SPACE_ID")

    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. Use existing agent
    try:
        # agent is already instantiated globally
        if not agent:
            return "Error: Agent not initialized", None
    except Exception as e:
        print(f"Error accessing agent: {e}")
        return f"Error accessing agent: {e}", None

    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)
            # Ensure submitted_answer is a simple string/number/float
            if isinstance(submitted_answer, dict):
                # Extract meaningful value or convert to string
                if len(submitted_answer) == 1:
                    submitted_answer = list(submitted_answer.values())[0]
                else:
                    submitted_answer = str(submitted_answer)
            elif not isinstance(submitted_answer, (str, int, float)):
                submitted_answer = str(submitted_answer)
            
            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

# Launch the Gradio UI
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")

    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(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")

    # Build Gradio Interface using Blocks
    with gr.Blocks() as demo:
        gr.Markdown("# Enhanced 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).
            """
        )

        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]
        )

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