Update app.py

app.py

@@ -1,11 +1,14 @@
-"""Simple Agent Evaluation Runner"""
+"""Enhanced Agent Evaluation Runner with improved capabilities"""
 import os
 import re
+import time
 import gradio as gr
 import requests
 import pandas as pd
 import google.generativeai as genai
 from dotenv import load_dotenv
+from urllib.parse import urlparse, parse_qs
+import json
 
 # Load environment variables
 load_dotenv()
@@ -16,66 +19,237 @@ genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 # Constants
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-class SimpleAgent:
-    """A simple agent using Google Gemini."""
+class EnhancedAgent:
+    """An enhanced agent using Google Gemini with improved capabilities."""
     
     def __init__(self):
-        print("SimpleAgent initialized.")
-        self.model = genai.GenerativeModel('gemini-1.5-flash')
+        print("EnhancedAgent initialized.")
+        # Use gemini-1.5-pro for better performance, fallback to flash
+        try:
+            self.model = genai.GenerativeModel('gemini-1.5-pro')
+        except:
+            self.model = genai.GenerativeModel('gemini-1.5-flash')
         
-    def __call__(self, question: str) -> str:
-        """Process a question and return an answer."""
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        # Rate limiting
+        self.last_request_time = 0
+        self.min_request_interval = 1.0  # 1 second between requests
         
-        # Simple system prompt
-        system_prompt = """You are a helpful assistant. Answer questions as accurately as possible.
-
-IMPORTANT: Your final answer should be:
-- A number (without commas, $ signs, or % signs unless specifically requested)
-- A few words as possible 
-- A comma-separated list if multiple items are requested
+    def _rate_limit(self):
+        """Simple rate limiting to avoid quota issues."""
+        current_time = time.time()
+        time_since_last = current_time - self.last_request_time
+        if time_since_last < self.min_request_interval:
+            time.sleep(self.min_request_interval - time_since_last)
+        self.last_request_time = time.time()
+    
+    def _extract_youtube_info(self, question: str) -> str:
+        """Extract information about YouTube videos mentioned in questions."""
+        youtube_patterns = [
+            r'youtube\.com/watch\?v=([a-zA-Z0-9_-]+)',
+            r'youtu\.be/([a-zA-Z0-9_-]+)'
+        ]
+        
+        for pattern in youtube_patterns:
+            match = re.search(pattern, question)
+            if match:
+                video_id = match.group(1)
+                return f"YouTube video ID: {video_id}. Note: Cannot access video content directly, but can make educated guesses based on context."
+        return ""
+    
+    def _analyze_question_type(self, question: str) -> str:
+        """Analyze the type of question and provide specific guidance."""
+        question_lower = question.lower()
+        
+        # Different question types and their handling strategies
+        if any(word in question_lower for word in ['youtube', 'video', 'watch']):
+            return "VIDEO_ANALYSIS"
+        elif any(word in question_lower for word in ['excel', 'spreadsheet', 'file', 'csv']):
+            return "FILE_ANALYSIS"
+        elif any(word in question_lower for word in ['how many', 'count', 'number of']):
+            return "COUNTING"
+        elif any(word in question_lower for word in ['who', 'what', 'where', 'when']):
+            return "FACTUAL"
+        elif any(word in question_lower for word in ['calculate', 'compute', 'math']):
+            return "CALCULATION"
+        elif any(word in question_lower for word in ['list', 'name', 'identify']):
+            return "LIST"
+        else:
+            return "GENERAL"
+    
+    def _get_enhanced_prompt(self, question: str, question_type: str) -> str:
+        """Generate an enhanced system prompt based on question type."""
+        
+        base_prompt = """You are an expert assistant with broad knowledge across many domains including:
+- Music, entertainment, and media
+- Sports statistics and history  
+- Science and mathematics
+- Geography and world facts
+- Technology and computing
+- Literature and culture
 
-Always end your response with: FINAL ANSWER: [your answer]
+CRITICAL INSTRUCTIONS:
+1. Always provide your best educated guess even if you're not 100% certain
+2. For numerical answers, provide ONLY the number (no commas, currency symbols, or units unless specified)
+3. For names/words, provide the exact spelling
+4. For lists, use comma-separated format
+5. End with: FINAL ANSWER: [your concise answer]
 
-Examples:
-- For "How many albums did X release?" → FINAL ANSWER: 5
-- For "What city is the capital?" → FINAL ANSWER: Paris
-- For "List the top 3 countries" → FINAL ANSWER: USA, China, Japan
 """
         
-        # Combine system prompt with question
-        full_prompt = f"{system_prompt}\n\nQuestion: {question}"
+        if question_type == "VIDEO_ANALYSIS":
+            base_prompt += """
+For video-related questions:
+- If you cannot access the video content, make educated guesses based on:
+  - Video title/URL context
+  - Common knowledge about the topic
+  - Typical content patterns
+- Provide your best estimate rather than saying "cannot access"
+"""
+        elif question_type == "FILE_ANALYSIS":
+            base_prompt += """
+For file-related questions:
+- If you cannot access files directly, make reasonable assumptions
+- Use general knowledge about typical data in such contexts
+- Provide educated estimates based on the question context
+"""
+        elif question_type == "COUNTING":
+            base_prompt += """
+For counting questions:
+- Provide specific numbers when possible
+- If exact count unknown, provide reasonable estimates
+- Consider historical data and typical ranges
+"""
+        elif question_type == "FACTUAL":
+            base_prompt += """
+For factual questions:
+- Use your knowledge base to provide accurate information
+- If multiple possibilities exist, choose the most likely one
+- Be specific with names, dates, and details
+"""
         
-        try:
-            # Generate response using Gemini
-            response = self.model.generate_content(full_prompt)
-            answer = response.text
-            
-            # Extract final answer if it exists
-            final_answer_match = re.search(r'FINAL ANSWER:\s*(.+?)(?:\n|$)', answer, re.IGNORECASE)
-            if final_answer_match:
-                final_answer = final_answer_match.group(1).strip()
-                return final_answer
-            else:
-                # If no "FINAL ANSWER:" format, try to extract a simple answer
-                # Look for numbers, short phrases, or lists
-                lines = answer.strip().split('\n')
-                for line in reversed(lines):  # Start from the end
-                    line = line.strip()
-                    if line and not line.startswith('FINAL'):
-                        # Simple heuristic: if it's short, likely an answer
-                        if len(line) < 100:
-                            return line
+        return base_prompt
+    
+    def _make_api_call_with_retry(self, prompt: str, max_retries: int = 3) -> str:
+        """Make API call with retry logic and error handling."""
+        
+        for attempt in range(max_retries):
+            try:
+                self._rate_limit()  # Apply rate limiting
                 
-                return answer.strip()[:100]  # Fallback to first 100 chars
+                # Generate response using Gemini
+                response = self.model.generate_content(
+                    prompt,
+                    generation_config=genai.types.GenerationConfig(
+                        temperature=0.1,  # Lower temperature for more consistent answers
+                        max_output_tokens=1000,
+                    )
+                )
                 
-        except Exception as e:
-            print(f"Error calling Gemini API: {e}")
-            return f"Error: {str(e)}"
+                if response.text:
+                    return response.text
+                else:
+                    raise Exception("Empty response from API")
+                    
+            except Exception as e:
+                error_msg = str(e).lower()
+                
+                if "quota" in error_msg or "429" in error_msg:
+                    if attempt < max_retries - 1:
+                        wait_time = (2 ** attempt) * 5  # Exponential backoff
+                        print(f"Quota exceeded, waiting {wait_time} seconds...")
+                        time.sleep(wait_time)
+                        continue
+                    else:
+                        return "Error: API quota exceeded"
+                elif "safety" in error_msg:
+                    return "Error: Content safety filter triggered"
+                else:
+                    if attempt < max_retries - 1:
+                        time.sleep(2)  # Wait before retry
+                        continue
+                    else:
+                        return f"Error: {str(e)}"
+        
+        return "Error: Max retries exceeded"
+    
+    def __call__(self, question: str) -> str:
+        """Process a question and return an answer."""
+        print(f"Agent processing: {question[:100]}...")
+        
+        # Analyze question type
+        question_type = self._analyze_question_type(question)
+        print(f"Question type identified: {question_type}")
+        
+        # Extract additional context
+        youtube_info = self._extract_youtube_info(question)
+        
+        # Build enhanced prompt
+        system_prompt = self._get_enhanced_prompt(question, question_type)
+        
+        # Add context if available
+        context = ""
+        if youtube_info:
+            context += f"\nContext: {youtube_info}\n"
+        
+        # Combine everything
+        full_prompt = f"{system_prompt}\n{context}\nQuestion: {question}\n\nProvide your best answer:"
+        
+        # Make API call with retry
+        response = self._make_api_call_with_retry(full_prompt)
+        
+        # Extract final answer
+        return self._extract_final_answer(response, question_type)
+    
+    def _extract_final_answer(self, response: str, question_type: str) -> str:
+        """Extract the final answer from the response."""
+        if response.startswith("Error:"):
+            return response
+        
+        # Look for FINAL ANSWER: pattern
+        final_answer_match = re.search(r'FINAL ANSWER:\s*(.+?)(?:\n|$)', response, re.IGNORECASE)
+        if final_answer_match:
+            answer = final_answer_match.group(1).strip()
+            return self._clean_answer(answer, question_type)
+        
+        # Fallback: extract from end of response
+        lines = response.strip().split('\n')
+        for line in reversed(lines):
+            line = line.strip()
+            if line and len(line) < 200:  # Reasonable answer length
+                return self._clean_answer(line, question_type)
+        
+        # Last resort: return first part of response
+        return self._clean_answer(response[:100], question_type)
+    
+    def _clean_answer(self, answer: str, question_type: str) -> str:
+        """Clean and format the final answer."""
+        answer = answer.strip()
+        
+        # Remove common prefixes
+        prefixes_to_remove = [
+            "the answer is", "answer:", "final answer:", 
+            "result:", "solution:", "therefore",
+            "in conclusion", "to summarize"
+        ]
+        
+        for prefix in prefixes_to_remove:
+            if answer.lower().startswith(prefix):
+                answer = answer[len(prefix):].strip()
+        
+        # Clean punctuation from the end
+        answer = answer.rstrip('.,;:!')
+        
+        # For counting questions, ensure we return just the number
+        if question_type == "COUNTING":
+            number_match = re.search(r'\b(\d+(?:,\d{3})*(?:\.\d+)?)\b', answer)
+            if number_match:
+                return number_match.group(1).replace(',', '')
+        
+        return answer
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
-    Fetches all questions, runs the SimpleAgent on them, submits all answers,
+    Fetches all questions, runs the EnhancedAgent on them, submits all answers,
     and displays the results.
     """
     # Check if user is logged in
@@ -96,7 +270,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
     # 1. Initialize Agent
     try:
-        agent = SimpleAgent()
+        agent = EnhancedAgent()
     except Exception as e:
         print(f"Error initializing agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -149,6 +323,9 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                 "Submitted Answer": submitted_answer
             })
             
+            # Small delay between questions to avoid rate limiting
+            time.sleep(0.5)
+            
         except Exception as e:
             error_msg = f"ERROR: {str(e)}"
             print(f"Error processing task {task_id}: {e}")
@@ -196,15 +373,20 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         return error_msg, results_df
 
 # Build Gradio Interface
-with gr.Blocks(title="Simple Agent Evaluation") as demo:
-    gr.Markdown("# Simple Agent Evaluation Runner")
+with gr.Blocks(title="Enhanced Agent Evaluation") as demo:
+    gr.Markdown("# Enhanced Agent Evaluation Runner")
     gr.Markdown("""
     **Instructions:**
     1. Make sure you have set up your `GOOGLE_API_KEY` in the environment variables
     2. Log in to your Hugging Face account using the button below
     3. Click 'Run Evaluation & Submit All Answers' to start the evaluation
     
-    **Note:** This is a simplified agent that uses Google Gemini to answer questions.
+    **Enhanced Features:**
+    - Improved question analysis and categorization
+    - Better handling of different question types
+    - Rate limiting to avoid API quota issues
+    - Retry logic for failed requests
+    - Enhanced prompting for better accuracy
     """)
 
     gr.LoginButton()
@@ -229,7 +411,7 @@ with gr.Blocks(title="Simple Agent Evaluation") as demo:
 
 if __name__ == "__main__":
     print("=" * 50)
-    print("🚀 Starting Simple Agent Evaluation Runner")
+    print("🚀 Starting Enhanced Agent Evaluation Runner")
     print("=" * 50)
     
     # Check environment variables