Update agent.py

agent.py

@@ -116,48 +116,56 @@ class VisitWebpageTool(Tool):
         except RequestException as e:
             return f"Error fetching the webpage: {str(e)}"
         except Exception as e:
-            return f"An unexpected error occurred: {str(e)}"
+            return
+ f"An unexpected error occurred: {str(e)}"
 
     def __init__(self, *args, **kwargs):
         self.is_initialized = False
 
 class DownloadTaskAttachmentTool(Tool):
     name = "download_file"
-    description = "Downloads the file attached to the task ID"
+    description = "Downloads the file attached to the task ID and returns the local file path. Supports Excel (.xlsx), image (.png), and audio (.mp3) files."
     inputs = {'task_id': {'type': 'string', 'description': 'The task id to download attachment from.'}}
     output_type = "string"
 
-    
-    def forward(self, task_id: str) -> str:
-        """
-        Downloads a file associated with the given task ID.
-        Returns the file path where the file is saved locally.
-        """
-        file_url = f"{DEFAULT_API_URL}/files/{task_id}"
-        local_file_path = f"downloads/{task_id}.file"
+    def __init__(self, rate_limiter: Optional[Limiter] = None, default_api_url: str = "https://agents-course-unit4-scoring.hf.space", *args, **kwargs):
+        self.is_initialized = False
+        self.rate_limiter = rate_limiter
+        self.default_api_url = default_api_url
 
+    async def forward(self, task_id: str) -> str:
+        file_url = f"{self.default_api_url}/files/{task_id}"
         print(f"Downloading file for task ID {task_id} from {file_url}...")
         try:
-
-            file_url = f"{DEFAULT_API_URL}/files/{task_id}"
-            local_path = f"downloads/{task_id}.xlsx"
-                        
+            if self.rate_limiter:
+                while not self.rate_limiter.consume(1):
+                    print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
+                    await asyncio.sleep(60 / 15)  # Assuming 15 RPM
             response = requests.get(file_url, stream=True, timeout=15)
             response.raise_for_status()
-
+            
+            # Determine file extension based on Content-Type
+            content_type = response.headers.get('Content-Type', '').lower()
+            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'
+            else:
+                return f"Error: Unsupported file type {content_type} for task {task_id}"
+            
+            local_file_path = f"downloads/{task_id}{extension}"
             os.makedirs("downloads", exist_ok=True)
             with open(local_file_path, "wb") as file:
                 for chunk in response.iter_content(chunk_size=8192):
                     file.write(chunk)
-
             print(f"File downloaded successfully: {local_file_path}")
             return local_file_path
         except requests.exceptions.RequestException as e:
-            print(f"Error downloading file for task {task_id}: {e}")
-            raise
-
-    def __init__(self, *args, **kwargs):
-        self.is_initialized = False
+            return f"Error downloading file for task {task_id}: {str(e)}"
 
 @tool
 def SpeechToTextTool(audio_path: str) -> str: