Update app.py

app.py

@@ -25,43 +25,29 @@ TOKEN_BUCKET_REFILL_RATE = MAX_MODEL_CALLS_PER_MINUTE / 60.0  # Tokens per secon
 storage = MemoryStorage()
 token_bucket = Limiter(rate=TOKEN_BUCKET_REFILL_RATE, capacity=TOKEN_BUCKET_CAPACITY, storage=storage)
 
-async def check_n_load_attach(session: aiohttp.ClientSession, task_id
-: str, api_url: str = DEFAULT_API_URL) -> Optional[str]:
+async def check_n_load_attach(session: aiohttp.ClientSession, task_id: str, question: str, api_url: str = "https://agents-course-unit4-scoring.hf.space") -> Optional[str]:
     file_url = f"{api_url}/files/{task_id}"
     try:
         async with session.get(file_url, timeout=15) as response:
             if response.status == 200:
-
-                # Determine file extension from Content-Type
                 content_type = str(response.headers.get("Content-Type", "")).lower()
-                extension = ""
-                if "image/png" in content_type:
-                    extension = ".png"
-                elif "image/jpeg" in content_type:
-                    extension = ".jpg"
-                elif "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" in content_type:
-                    extension = ".xlsx"
-                elif "audio/mpeg" in content_type:
-                    extension = ".mp3"
-                elif "application/pdf" in content_type:
-                    extension = ".pdf"
-                elif "text/x-python" in content_type:
-                    extension = ".py"
+                content = await response.read()  # Read the file content
+                
+                # Determine extension based on content_type, content, or question
+                extension = await determine_extension(content_type, content, question)
+                
+                if extension:
+                    filename = f"{task_id}{extension}"
+                    local_file_path = os.path.join("downloads", filename)
+                    os.makedirs("downloads", exist_ok=True)
+                    
+                    async with aiofiles.open(local_file_path, "wb") as file:
+                        await file.write(content)
+                    print(f"File downloaded successfully: {local_file_path}")
+                    return local_file_path
                 else:
                     print(f"Unsupported content type: {content_type} for task {task_id}")
                     return None
-                
-                # Use task_id as the filename to ensure uniqueness
-                filename = f"{task_id}{extension}"
-                local_file_path = os.path.join("downloads", filename)
-                os.makedirs("downloads", exist_ok=True)
-                
-                # Save the file
-                async with aiofiles.open(local_file_path, "wb") as file:
-                    async for chunk in response.content.iter_chunked(8192):
-                        await file.write(chunk)
-                print(f"File downloaded successfully: {local_file_path}")
-                return local_file_path
             else:
                 print(f"Failed to download file for task {task_id}: HTTP {response.status}")
                 return None
@@ -69,6 +55,35 @@ async def check_n_load_attach(session: aiohttp.ClientSession, task_id
         print(f"Error downloading attachment for task {task_id}: {str(e)}")
         return None
 
+async def determine_extension(content_type: str, content: bytes, question: str) -> Optional[str]:
+    # Check if the question mentions Excel
+    if "excel" in question.lower():
+        # Check for XLS signature
+        if content.startswith(b'\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1'):
+            return ".xls"
+        # Check for XLSX signature (ZIP archive)
+        elif content.startswith(b'\x50\x4B\x03\x04'):
+            return ".xlsx"
+        else:
+            return ".xlsx"  # Default to XLSX if unsure
+    # Standard MIME type checks
+    if "image/png" in content_type:
+        return ".png"
+    elif "jpeg" in content_type or "jpg" in content_type:
+        return ".jpg"
+    elif "spreadsheetml.sheet" in content_type:
+        return ".xlsx"
+    elif "vnd.ms-excel" in content_type:
+        return ".xls"
+    elif "audio/mpeg" in content_type:
+        return ".mp3"
+    elif "application/pdf" in content_type:
+        return ".pdf"
+    elif "text/x-python" in content_type:
+        return ".py"
+    else:
+        return None
+
 
 async def fetch_questions(session: aiohttp.ClientSession, questions_url: str) -> list:
     """Fetch questions asynchronously."""