Update app.py

app.py

@@ -10,30 +10,24 @@ from google.generativeai.types import HarmCategory, HarmBlockThreshold
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# --- New Native Gemini Agent ---
+# --- User's Corrected NativeGeminiAgent Class ---
+# This is the superior implementation provided by you.
 class NativeGeminiAgent:
-    """
-    An agent that leverages Gemini's native multi-modal capabilities,
-    including grounding, video, and file understanding.
-    """
     def __init__(self, gemini_api_key: str, api_url: str):
-        print("Initializing NativeGeminiAgent...")
+        print("Initializing NativeGeminiAgent with corrected configuration...")
         genai.configure(api_key=gemini_api_key)
         
         self.api_url = api_url
+        self.model_name = 'gemini-2.5-pro-preview-06-05' # Using the stable, powerful model
         
-        # --- CORRECTED INITIALIZATION ---
-        # Enable native grounding with Google Search. No parameters are needed.
-        google_search_retrieval = genai.protos.Tool(
-            google_search_retrieval=genai.protos.GoogleSearchRetrieval()
-        )
-        
-        # Configure the model with the native tool
-        self.model_name = 'gemini-2.5-pro-preview-06-05' # Using the best stable model
+        # Correct tool configuration using the recommended string-based method
         self.model = genai.GenerativeModel(
             model_name=self.model_name,
-            tools=[google_search_retrieval],
-            system_instruction="You are a world-class problem solver. Your goal is to answer the user's question accurately. Use your tools and reasoning abilities to provide a definitive answer.",
+            tools=['google_search_retrieval'],
+            system_instruction="""You are a world-class problem solver and researcher. 
+            Analyze the question carefully, use available tools to gather information, 
+            and provide accurate, concise answers. Focus on factual information and 
+            avoid speculation.""",
             safety_settings={
                 HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                 HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
@@ -41,14 +35,28 @@ class NativeGeminiAgent:
                 HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
             }
         )
-        print(f"Agent equipped with {self.model_name} and native Google Search grounding.")
+        print(f"Agent initialized with {self.model_name} and Google Search grounding.")
+
+    def _get_mime_type(self, url: str) -> str:
+        """Enhanced MIME type detection."""
+        url_lower = url.lower()
+        if url_lower.endswith(('.jpg', '.jpeg')): return "image/jpeg"
+        elif url_lower.endswith('.png'): return "image/png"
+        elif url_lower.endswith('.gif'): return "image/gif"
+        elif url_lower.endswith('.pdf'): return "application/pdf"
+        elif url_lower.endswith('.txt'): return "text/plain"
+        elif url_lower.endswith('.csv'): return "text/csv"
+        elif url_lower.endswith(('.mp4', '.avi', '.mov')): return "video/mp4"
+        elif url_lower.endswith('.json'): return "application/json"
+        else: return "application/octet-stream"
 
     def _check_if_file_exists(self, url: str) -> bool:
-        """Checks if a remote file exists before sending it to Gemini."""
+        """Enhanced file existence check."""
         try:
-            response = requests.head(url, timeout=10)
+            response = requests.head(url, timeout=15, allow_redirects=True)
             return response.status_code == 200
-        except requests.exceptions.RequestException:
+        except requests.exceptions.RequestException as e:
+            print(f"File check failed for {url}: {e}")
             return False
 
     def __call__(self, question: str, task_id: str) -> str:
@@ -56,39 +64,51 @@ class NativeGeminiAgent:
         
         prompt_parts = [question]
         
-        # 1. Automatically find and add any URLs from the question text
-        urls_in_question = re.findall(r'https?://\S+', question)
-        if urls_in_question:
-            for url in urls_in_question:
-                print(f"Found URL in question: {url}")
-                # Use a general MIME type that Gemini can often auto-detect for sources like YouTube
-                prompt_parts.append(genai.Part.from_uri(uri=url, mime_type="video/mp4"))
+        # Enhanced URL detection
+        urls_in_question = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', question)
+        for url in urls_in_question:
+            try:
+                mime_type = self._get_mime_type(url)
+                prompt_parts.append(genai.Part.from_uri(uri=url, mime_type=mime_type))
+                print(f"Added URL: {url} (MIME: {mime_type})")
+            except Exception as e:
+                print(f"Failed to add URL {url}: {e}")
 
-        # 2. Check for and add any associated files from the GAIA server
+        # Check for associated files
         file_url = f"{self.api_url}/files/{task_id}"
         if self._check_if_file_exists(file_url):
-            print(f"Found associated file, adding URL: {file_url}")
-            # Simple MIME type guessing for common GAIA file types
-            mime_type = "image/jpeg" # Default
-            if file_url.endswith('.pdf'): mime_type = "application/pdf"
-            if file_url.endswith('.txt'): mime_type = "text/plain"
-            prompt_parts.append(genai.Part.from_uri(uri=file_url, mime_type=mime_type))
-        else:
-            print("No associated file found for this task.")
-            
-        print(f"Sending {len(prompt_parts)} parts to the model.")
+            try:
+                mime_type = self._get_mime_type(file_url)
+                prompt_parts.append(genai.Part.from_uri(uri=file_url, mime_type=mime_type))
+                print(f"Added file: {file_url} (MIME: {mime_type})")
+            except Exception as e:
+                print(f"Failed to add file {file_url}: {e}")
 
         try:
-            response = self.model.generate_content(prompt_parts, request_options={'timeout': 120})
-            
-            # The grounding feature may add citations. We remove them for the final answer.
-            final_answer = re.sub(r'\[\d+\]', '', response.text).strip()
-            print(f"Model generated answer: {final_answer}")
-            return final_answer
+            # Use the specified generation config for more stable outputs
+            response = self.model.generate_content(
+                prompt_parts,
+                request_options={'timeout': 120},
+                generation_config=genai.types.GenerationConfig(
+                    temperature=0.1,
+                    top_p=0.8,
+                    max_output_tokens=2048
+                )
+            )
             
+            if response.text:
+                # Thoroughly clean the response text
+                final_answer = response.text.strip()
+                final_answer = re.sub(r'\[\d+\]', '', final_answer) # Remove citations
+                final_answer = re.sub(r'\s+', ' ', final_answer).strip() # Normalize whitespace
+                return final_answer
+            else:
+                return "AGENT_ERROR: Empty response from model"
+                
         except Exception as e:
-            print(f"An error occurred while calling the Gemini API: {e}")
-            return f"AGENT_ERROR: Could not get a response from the model. Details: {e}"
+            error_msg = f"AGENT_ERROR: {str(e)}"
+            print(error_msg)
+            return error_msg
 
 # --- Main run_and_submit_all function ---
 def run_and_submit_all(profile: gr.OAuthProfile | None):
@@ -140,8 +160,8 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Native Multi-Modal GAIA Agent")
-    gr.Markdown("This agent uses Gemini 2.5 Pro with native Google Search grounding and direct multi-modal understanding (video, images, files).")
+    gr.Markdown("# Native Multi-Modal GAIA Agent (Corrected)")
+    gr.Markdown("This agent uses the improved architecture with proper tool configuration, MIME type detection, and error handling.")
     gr.LoginButton()
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)