app.py

"""🚀 Optimized GAIA Agent Evaluation Runner for Hugging Face Space"""
import os
import gradio as gr
import requests
import pandas as pd
import time
import warnings

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore", category=UserWarning)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Try to import LangChain dependencies
try:
    from langchain_core.messages import HumanMessage
    from agent import build_graph
    LANGCHAIN_AVAILABLE = True
    print("✅ LangChain dependencies loaded successfully")
except ImportError as e:
    print(f"❌ LangChain import failed: {e}")
    print("🔄 The Space will need to restart after dependencies are installed")
    LANGCHAIN_AVAILABLE = False
    
    # Create fallback classes
    class HumanMessage:
        def __init__(self, content):
            self.content = content
    
    def build_graph(provider="groq"):
        return None

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

# --- Optimized Agent Definition ---
class OptimizedGAIAAgent:
    """🧠 Advanced GAIA Agent with enhanced capabilities"""
    
    def __init__(self, provider="groq"):
        print("🤖 Initializing Optimized GAIA Agent...")
        if not LANGCHAIN_AVAILABLE:
            print("❌ LangChain not available - using fallback mode")
            self.graph = None
            self.provider = provider
            return
            
        try:
            self.graph = build_graph(provider=provider)
            self.provider = provider
            print(f"✅ Agent initialized successfully with {provider} provider!")
        except Exception as e:
            print(f"❌ Error initializing agent: {e}")
            self.graph = None

    def extract_final_answer(self, content: str) -> str:
        """Extract clean final answer from agent response"""
        if not content:
            return "No answer provided"
        
        # Remove common prefixes
        prefixes_to_remove = [
            "FINAL ANSWER:", "Final Answer:", "Answer:", 
            "The answer is:", "The final answer is:"
        ]
        
        clean_content = content.strip()
        for prefix in prefixes_to_remove:
            if clean_content.upper().startswith(prefix.upper()):
                clean_content = clean_content[len(prefix):].strip()
                break
        
        return clean_content if clean_content else content.strip()

    def __call__(self, question: str) -> str:
        """Process question and return clean answer"""
        print(f"🔍 Processing: {question[:50]}{'...' if len(question) > 50 else ''}")
        
        # Check if agent is available
        if not self.graph:
            return f"LangChain agent not available. Question: {question[:100]}..."
        
        try:
            start_time = time.time()
            
            # Wrap question in HumanMessage
            messages = [HumanMessage(content=question)]
            
            # Invoke graph
            result = self.graph.invoke({"messages": messages})
            
            # Extract answer from the last message
            if result and "messages" in result and result["messages"]:
                last_message = result["messages"][-1]
                raw_answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
                
                # Clean the answer
                clean_answer = self.extract_final_answer(raw_answer)
                
                elapsed = time.time() - start_time
                print(f"✅ Answered in {elapsed:.2f}s: {clean_answer[:50]}{'...' if len(clean_answer) > 50 else ''}")
                
                return clean_answer
            else:
                return "No valid response generated"
                
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            print(f"❌ {error_msg}")
            return error_msg


def run_and_submit_all(profile: gr.OAuthProfile | None, progress=gr.Progress()):
    """
    🚀 Optimized function to fetch questions, run agent, and submit answers
    """
    # Check authentication
    if not profile:
        return "🔐 Please login to Hugging Face first!", None
    
    username = profile.username
    print(f"👤 User logged in: {username}")
    
    # Get space information
    space_id = os.getenv("SPACE_ID", "unknown-space")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    
    # API endpoints
    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"
    
    # Initialize agent with progress
    progress(0.05, "🤖 Initializing GAIA Agent...")
    try:
        # Try Groq first (faster), fallback to Google if needed
        try:
            agent = OptimizedGAIAAgent(provider="groq")
        except:
            print("⚠️ Groq failed, trying Google provider...")
            agent = OptimizedGAIAAgent(provider="google")
    except Exception as e:
        error_msg = f"❌ Failed to initialize agent: {e}"
        print(error_msg)
        return error_msg, None
    
    # Fetch questions
    progress(0.1, "📥 Fetching questions from GAIA benchmark...")
    try:
        response = requests.get(questions_url, timeout=30)
        response.raise_for_status()
        questions_data = response.json()
        
        if not questions_data:
            return "📭 No questions received from server", None
            
        print(f"📋 Fetched {len(questions_data)} questions")
        
    except Exception as e:
        error_msg = f"📥 Error fetching questions: {e}"
        print(error_msg)
        return error_msg, None
    
    # Process questions with progress tracking
    results_log = []
    answers_payload = []
    total_questions = len(questions_data)
    
    print(f"🧠 Processing {total_questions} questions...")
    
    for i, item in enumerate(questions_data):
        task_id = item.get("task_id")
        question_text = item.get("question")
        
        if not task_id or not question_text:
            continue
        
        # Update progress
        progress_val = 0.1 + (i / total_questions) * 0.7  # 10% to 80%
        progress(progress_val, f"🧠 Processing question {i+1}/{total_questions}")
        
        try:
            # Get answer from agent
            submitted_answer = agent(question_text)
            
            # Add to payload
            answers_payload.append({
                "task_id": task_id, 
                "submitted_answer": submitted_answer
            })
            
            # Add to results log
            results_log.append({
                "Task ID": task_id,
                "Question": question_text[:100] + "..." if len(question_text) > 100 else question_text,
                "Submitted Answer": submitted_answer[:200] + "..." if len(submitted_answer) > 200 else submitted_answer
            })
            
        except Exception as e:
            error_answer = f"Agent Error: {e}"
            results_log.append({
                "Task ID": task_id,
                "Question": question_text[:100] + "..." if len(question_text) > 100 else question_text,
                "Submitted Answer": error_answer
            })
            print(f"❌ Error on {task_id}: {e}")
    
    if not answers_payload:
        return "⚠️ No answers generated to submit", pd.DataFrame(results_log)
    
    # Submit answers
    progress(0.85, f"📤 Submitting {len(answers_payload)} answers...")
    submission_data = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers_payload
    }
    
    try:
        print(f"📤 Submitting to: {submit_url}")
        response = requests.post(submit_url, json=submission_data, timeout=120)
        response.raise_for_status()
        result_data = response.json()
        
        # Format results
        score = result_data.get('score', 0)
        correct_count = result_data.get('correct_count', 0)
        total_attempted = result_data.get('total_attempted', 0)
        message = result_data.get('message', 'No message')
        
        # Determine performance level
        if score >= 30:
            performance = "🎉 EXCELLENT! Certification Ready!"
            emoji = "🏆"
        elif score >= 20:
            performance = "🎯 GOOD! Almost there!"
            emoji = "💪"
        elif score >= 10:
            performance = "📈 IMPROVING! Keep going!"
            emoji = "🔥"
        else:
            performance = "💡 LEARNING! Room for improvement!"
            emoji = "🌱"
        
        progress(1.0, "✅ Submission completed!")
        
        final_status = f"""
{emoji} GAIA BENCHMARK RESULTS {emoji}

👤 User: {username}
📊 Score: {score}% ({correct_count}/{total_attempted} correct)
🎯 Target: 30% for Unit 4 certification
{performance}

💬 Server Message: {message}
🔗 Code Repository: {agent_code}

{"🎓 CONGRATULATIONS! You've achieved the certification threshold!" if score >= 30 else "💪 Keep optimizing your agent for better results!"}
        """.strip()
        
        print(f"🎉 Submission successful! Score: {score}%")
        return final_status, pd.DataFrame(results_log)
        
    except Exception as e:
        error_msg = f"📤 Submission failed: {e}"
        print(error_msg)
        return error_msg, pd.DataFrame(results_log)


# --- Enhanced Gradio Interface ---
with gr.Blocks(
    title="🤖 Advanced GAIA Agent - Unit 4 Certification",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .main-header {
        text-align: center;
        background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: bold;
    }
    .feature-box {
        background: #f8f9fa;
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
        border-left: 4px solid #4ecdc4;
    }
    """
) as demo:
    
    gr.Markdown("""
    <div class="main-header">
    <h1>🤖 Advanced GAIA Agent - Unit 4 Certification</h1>
    </div>
    
    <div style="text-align: center; margin: 20px 0;">
    <h3>🎯 Goal: Score 30+ on GAIA Level 1 Questions</h3>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("""
            <div class="feature-box">
            <h4>🧠 Agent Features:</h4>
            <ul>
            <li><strong>Multi-Model Support:</strong> Groq Llama 3.1 70B + Google Gemini 2.0</li>
            <li><strong>Advanced Tools:</strong> Web search, Wikipedia, arXiv, Math calculations</li>
            <li><strong>Vector Database:</strong> Supabase-powered similarity search</li>
            <li><strong>Optimized Processing:</strong> Fast response times with error handling</li>
            </ul>
            </div>
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("""
            <div class="feature-box">
            <h4>📋 Instructions:</h4>
            <ol>
            <li><strong>Login:</strong> Use your Hugging Face account below</li>
            <li><strong>Run:</strong> Click the evaluation button</li>
            <li><strong>Wait:</strong> Processing ~20 questions takes 3-5 minutes</li>
            <li><strong>Results:</strong> View your score and detailed answers</li>
            </ol>
            </div>
            """)
    
    gr.Markdown("---")
    
    with gr.Row():
        with gr.Column(scale=2):
            login_button = gr.LoginButton(
                value="🔐 Login with Hugging Face",
                size="lg"
            )
        with gr.Column(scale=3):
            run_button = gr.Button(
                "🚀 Run GAIA Evaluation & Submit All Answers",
                variant="primary",
                size="lg"
            )
    
    with gr.Row():
        status_output = gr.Textbox(
            label="📊 Evaluation Status & Results",
            lines=12,
            interactive=False,
            placeholder="Click the evaluation button above to start...",
            show_copy_button=True
        )
    
    with gr.Row():
        results_table = gr.DataFrame(
            label="📝 Detailed Question Results",
            wrap=True,
            interactive=False,
            column_widths=["15%", "60%", "25%"]
        )
    
    # Enhanced footer with information
    gr.Markdown("""
    ---
    <div style="text-align: center; color: #666; font-size: 14px;">
    <p><strong>💡 Tips for Success:</strong></p>
    <ul style="text-align: left; max-width: 600px; margin: 0 auto;">
    <li>Ensure your API keys are set in the Space secrets</li>
    <li>The agent uses multiple tools automatically based on question type</li>
    <li>Complex questions may take longer to process</li>
    <li>Score ≥30% qualifies for Unit 4 certification</li>
    </ul>
    </div>
    """)
    
    # Connect the button with progress
    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table],
        show_progress=True
    )

if __name__ == "__main__":
    print("\n" + "🚀" * 20 + " GAIA AGENT STARTING " + "🚀" * 20)
    
    # Environment check
    space_host = os.getenv("SPACE_HOST")
    space_id = os.getenv("SPACE_ID")
    
    if space_host:
        print(f"✅ Running on Hugging Face Space")
        print(f"   🌐 URL: https://{space_host}.hf.space")
    else:
        print("ℹ️  Running locally (SPACE_HOST not found)")
    
    if space_id:
        print(f"📂 Space ID: {space_id}")
        print(f"🔗 Repository: https://huggingface.co/spaces/{space_id}")
    else:
        print("⚠️  SPACE_ID not found - using fallback")
    
    # API Keys check
    api_keys = {
        "🤖 Groq": os.getenv("GROQ_API_KEY"),
        "🔍 Google": os.getenv("GOOGLE_API_KEY"), 
        "🌐 Tavily": os.getenv("TAVILY_API_KEY"),
        "🗄️ Supabase URL": os.getenv("SUPABASE_URL"),
        "🔑 Supabase Key": os.getenv("SUPABASE_SERVICE_KEY")    }
    
    print("\n🔐 API Configuration:")
    for name, key in api_keys.items():
        status = "✅ Configured" if key else "❌ Missing"
        print(f"   {name}: {status}")
    
    # Warnings for missing keys
    if not api_keys["🤖 Groq"] and not api_keys["🔍 Google"]:
        print("\n⚠️  WARNING: No LLM API key found! Agent will fail.")
        print("   Please set GROQ_API_KEY or GOOGLE_API_KEY in Space secrets")
    
    if not api_keys["🌐 Tavily"]:
        print("\n💡 INFO: Tavily API key missing - web search will be limited")
    
    print("\n" + "🎬" * 15 + " LAUNCHING GRADIO INTERFACE " + "🎬" * 15)
    
    try:
        demo.launch(
            debug=True,
            share=False,
            show_error=True,
            server_name="0.0.0.0",  # Important for HF Spaces
            server_port=7860        # Standard HF Spaces port
        )
    except Exception as e:
        print(f"❌ Error launching app: {e}")
        raise