Update agent.py

agent.py

@@ -85,255 +85,7 @@ def GoogleSearchTool(query: str) -> str:
 #    
 #    return f"The image description: '{response}'"
 
-class VisitWebpageTool(Tool):
-    name = "visit_webpage"
-    description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
-    inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
-    output_type = "string"
 
-    def forward(self, url: str) -> str:
-        try:
-            response = requests.get(url, timeout=20)
-            response.raise_for_status()
-            markdown_content = markdownify(response.text).strip()
-            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
-            from smolagents.utils import truncate_content
-            return truncate_content(markdown_content, 10000)
-        except requests.exceptions.Timeout:
-            return "The request timed out. Please try again later or check the URL."
-        except requests.exceptions.RequestException as e:
-            return f"Error fetching the webpage: {str(e)}"
-        except Exception as 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 and returns the local file path. Supports Excel (.xlsx), image (.png, .jpg), audio (.mp3), PDF (.pdf), and Python (.py) files."
-    inputs = {'task_id': {'type': 'string', 'description': 'The task id to download attachment from.'}}
-    output_type = "string"
-    DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-
-    def __init__(self, rate_limiter: Optional[Limiter] = None, default_api_url: str = DEFAULT_API_URL, *args, **kwargs):
-        self.is_initialized = False
-        self.rate_limiter = rate_limiter
-        self.default_api_url = default_api_url
-
-    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:
-            if self.rate_limiter:
-                while not self.rate_limiter.consume(1):
-                    print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
-                    time.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'
-            elif 'application/pdf' in content_type:
-                extension = '.pdf'
-            elif 'text/x-python' in content_type:
-                extension = '.py'
-            else:
-                return f"Error: Unsupported file type {content_type} for task {task_id}. Try using visit_webpage or web_search if the content is online."
-            
-            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.HTTPError as e:
-            if e.response.status_code == 429:
-                return f"Error: Rate limit exceeded for task {task_id}. Try again later."
-            return f"Error downloading file for task {task_id}: {str(e)}"
-        except requests.exceptions.RequestException as e:
-            return f"Error downloading file for task {task_id}: {str(e)}"
-
-class SpeechToTextTool(Tool):
-    name = "speech_to_text"
-    description = (
-        "Converts an audio file to text using OpenAI Whisper."
-    )
-    inputs = {
-        "audio_path": {"type": "string", "description": "Path to audio file (.mp3, .wav)"},
-    }
-    output_type = "string"
-
-    def __init__(self):
-        super().__init__()
-        self.model = whisper.load_model("base")
-
-    def forward(self, audio_path: str) -> str:
-        if not os.path.exists(audio_path):
-            return f"Error: File not found at {audio_path}"
-        result = self.model.transcribe(audio_path)
-        return result.get("text", "")
-
-class ExcelReaderTool(Tool):
-    name = "excel_reader"
-
-    description = """
-    This tool reads and processes Excel files (.xlsx, .xls).
-    It can extract data, calculate statistics, and perform data analysis on spreadsheets.
-    """
-    inputs = {
-        "excel_path": {
-            "type": "string"
-,
-            "description": "The path to the Excel file to read",
-        },
-        "sheet_name": {
-            "type": "string",
-
-            "description": "The name of the sheet to read (optional, defaults to first sheet)",
-            "nullable": True
-        }
-    }
-    output_type = "string"
-    
-    def forward(self, excel_path: str, sheet_name: str = None) -> str:
-        """
-        Reads and processes the given Excel file.
-        """
-        try:
-            # Check if the file exists
-            if not os.path.exists(excel_path):
-                return f"Error: Excel file not found at {excel_path}"
-                
-            import pandas as pd
-            
-            # Read the Excel file
-            if sheet_name:
-                df = pd.read_excel(excel_path, sheet_name=sheet_name)
-            else:
-                df = pd.read_excel(excel_path)
-                
-            # Get basic info about the data
-            info = {
-                "shape": df.shape,
-                "columns": list(df.columns),
-                "dtypes": df.dtypes.to_dict(),
-                "head": df.head(5).to_dict()
-            }
-            
-            # Return formatted info
-            result = f"Excel file: {excel_path}\n"
-            result += f"Shape: {info['shape'][0]} rows × {info['shape'][1]} columns\n\n"
-            result += "Columns:\n"
-            for col in info['columns']:
-                result += f"- {col} ({info['dtypes'].get(col)})\n"
-            
-            result += "\nPreview (first 5 rows):\n"
-            result += df.head(5).to_string()
-            
-            return result
-            
-        except Exception as e:
-            return f"Error reading Excel file: {str(e)}"
-
-
-
-
-    class DownloadImageTool(Tool):
-        name = "download_chess_image"
-        description = "Downloads chess position image from task ID"
-        inputs = {'task_id': {'type': 'string'}}
-        output_type = "string"
-
-        def forward(self, task_id: str) -> str:
-            try:
-                response = requests.get(
-                    f"https://agents-course-unit4-scoring.hf.space/files/{task_id}",
-                    stream=True
-                )
-                response.raise_for_status()
-                
-                img_path = f"chess_{task_id}.png"
-                with open(img_path, "wb") as f:
-                    for chunk in response.iter_content(8192):
-                        f.write(chunk)
-                return img_path
-            except Exception as e:
-                raise RuntimeError(f"Image download failed: {str(e)}")
-
-
-    
-    class ChessEngineTool(Tool):
-        import chess
-        import chess.engine
-        name = "stockfish_analysis"
-        description = "Analyzes chess position using Stockfish"
-        inputs = {'fen': {'type': 'string'}}
-        output_type = "string"
-
-        def forward(self, fen: str) -> str:
-            try:
-                board = chess.Board(fen)
-                engine = chess.engine.SimpleEngine.popen_uci("stockfish")
-                result = engine.play(board, chess.engine.Limit(time=2.0))
-                engine.quit()
-                return board.san(result.move)
-            except Exception as e:
-                return f"Engine error: {str(e)}"
-
-    async def analyze_position(self, task_id: str):
-        try:
-            # Step 1: Download image
-            img_path = await self.tools[0](task_id)
-            
-            # Step 2: Get multimodal analysis
-            response = await self.model.acreate(
-                messages=[{
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": """Analyze this chess position. 
-                        It's black's turn. Provide the winning move in algebraic notation.
-                        Respond ONLY with the move, nothing else."""},
-                        {"type": "image_url", "image_url": {"url": f"file://{img_path}"}}
-                    ]
-                }],
-                temperature=0.1
-            )
-            
-            return response.choices[0].message.content
-            
-        except Exception as e:
-            return f"Analysis failed: {str(e)}"
-
-
-
-        
-class PythonCodeReaderTool(Tool):
-    name = "read_python_code"
-    description = "Reads a Python (.py) file and returns its content as a string."
-    inputs = {
-        "file_path": {"type": "string", "description": "The path to the Python file to read"}
-    }
-    output_type = "string"
-
-    def forward(self, file_path: str) -> str:
-        try:
-            if not os.path.exists(file_path):
-                return f"Error: Python file not found at {file_path}"
-            with open(file_path, "r", encoding="utf-8") as file:
-                content = file.read()
-            return content
-        except Exception as e:
-            return f"Error reading Python file: {str(e)}"
             
 class MagAgent:
     def __init__(self, rate_limiter: Optional[Limiter] = None):
@@ -354,12 +106,10 @@ class MagAgent:
             max_tokens=2048
         )
         
-        self.tools = [
-            self.DownloadImageTool(),
-            self.ChessEngineTool()
-        ]
+        # Initialize core tools
+        self.download_tool = self.DownloadTaskAttachmentTool(rate_limiter=rate_limiter)
+        self.chess_engine = self.ChessEngineTool()
 
-                
         # Load prompt templates
         with open("prompts.yaml", 'r') as stream:
             prompt_templates = yaml.safe_load(stream)
@@ -367,55 +117,322 @@ class MagAgent:
         # Initialize rate limiter for DuckDuckGoSearchTool
         search_rate_limiter = Limiter(rate=30/60, capacity=30, storage=MemoryStorage()) if not rate_limiter else rate_limiter
         
+        # Configure agent
         self.agent = CodeAgent(
-            model= model,
+            model=self.model,
             tools=[
-                DownloadTaskAttachmentTool(rate_limiter=rate_limiter),
-#                DuckDuckGoSearchTool(),
-#                WikipediaSearchTool(),
-                SpeechToTextTool(),
-                ExcelReaderTool(),
-                VisitWebpageTool(),
-                PythonCodeReaderTool(),
-                DownloadImageTool(),
-                ChessEngineTool(),
-#                GoogleSearchTool,
-#                ImageAnalysisTool,
+                self.download_tool,
+                self.chess_engine,
+                self.SpeechToTextTool(),
+                self.ExcelReaderTool(),
+                self.VisitWebpageTool(),
+                self.PythonCodeReaderTool()
             ],
             verbosity_level=2,
             prompt_templates=prompt_templates,
             add_base_tools=True,
             max_steps=15
         )
+
         print("MagAgent initialized.")
 
+
+    class VisitWebpageTool(Tool):
+        name = "visit_webpage"
+        description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
+        inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
+        output_type = "string"
+    
+        def forward(self, url: str) -> str:
+            try:
+                response = requests.get(url, timeout=20)
+                response.raise_for_status()
+                markdown_content = markdownify(response.text).strip()
+                markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
+                from smolagents.utils import truncate_content
+                return truncate_content(markdown_content, 10000)
+            except requests.exceptions.Timeout:
+                return "The request timed out. Please try again later or check the URL."
+            except requests.exceptions.RequestException as e:
+                return f"Error fetching the webpage: {str(e)}"
+            except Exception as 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 and returns the local file path. Supports Excel (.xlsx), image (.png, .jpg), audio (.mp3), PDF (.pdf), and Python (.py) files."
+        inputs = {'task_id': {'type': 'string', 'description': 'The task id to download attachment from.'}}
+        output_type = "string"
+        DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+    
+        def __init__(self, rate_limiter: Optional[Limiter] = None, default_api_url: str = DEFAULT_API_URL, *args, **kwargs):
+            self.is_initialized = False
+            self.rate_limiter = rate_limiter
+            self.default_api_url = default_api_url
+    
+        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:
+                if self.rate_limiter:
+                    while not self.rate_limiter.consume(1):
+                        print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
+                        time.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'
+                elif 'application/pdf' in content_type:
+                    extension = '.pdf'
+                elif 'text/x-python' in content_type:
+                    extension = '.py'
+                else:
+                    return f"Error: Unsupported file type {content_type} for task {task_id}. Try using visit_webpage or web_search if the content is online."
+                
+                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.HTTPError as e:
+                if e.response.status_code == 429:
+                    return f"Error: Rate limit exceeded for task {task_id}. Try again later."
+                return f"Error downloading file for task {task_id}: {str(e)}"
+            except requests.exceptions.RequestException as e:
+                return f"Error downloading file for task {task_id}: {str(e)}"
+    
+    class SpeechToTextTool(Tool):
+        name = "speech_to_text"
+        description = (
+            "Converts an audio file to text using OpenAI Whisper."
+        )
+        inputs = {
+            "audio_path": {"type": "string", "description": "Path to audio file (.mp3, .wav)"},
+        }
+        output_type = "string"
+    
+        def __init__(self):
+            super().__init__()
+            self.model = whisper.load_model("base")
+    
+        def forward(self, audio_path: str) -> str:
+            if not os.path.exists(audio_path):
+                return f"Error: File not found at {audio_path}"
+            result = self.model.transcribe(audio_path)
+            return result.get("text", "")
+    
+    class ExcelReaderTool(Tool):
+        name = "excel_reader"
+    
+        description = """
+        This tool reads and processes Excel files (.xlsx, .xls).
+        It can extract data, calculate statistics, and perform data analysis on spreadsheets.
+        """
+        inputs = {
+            "excel_path": {
+                "type": "string"
+    ,
+                "description": "The path to the Excel file to read",
+            },
+            "sheet_name": {
+                "type": "string",
+    
+                "description": "The name of the sheet to read (optional, defaults to first sheet)",
+                "nullable": True
+            }
+        }
+        output_type = "string"
+        
+        def forward(self, excel_path: str, sheet_name: str = None) -> str:
+            """
+            Reads and processes the given Excel file.
+            """
+            try:
+                # Check if the file exists
+                if not os.path.exists(excel_path):
+                    return f"Error: Excel file not found at {excel_path}"
+                    
+                import pandas as pd
+                
+                # Read the Excel file
+                if sheet_name:
+                    df = pd.read_excel(excel_path, sheet_name=sheet_name)
+                else:
+                    df = pd.read_excel(excel_path)
+                    
+                # Get basic info about the data
+                info = {
+                    "shape": df.shape,
+                    "columns": list(df.columns),
+                    "dtypes": df.dtypes.to_dict(),
+                    "head": df.head(5).to_dict()
+                }
+                
+                # Return formatted info
+                result = f"Excel file: {excel_path}\n"
+                result += f"Shape: {info['shape'][0]} rows × {info['shape'][1]} columns\n\n"
+                result += "Columns:\n"
+                for col in info['columns']:
+                    result += f"- {col} ({info['dtypes'].get(col)})\n"
+                
+                result += "\nPreview (first 5 rows):\n"
+                result += df.head(5).to_string()
+                
+                return result
+                
+            except Exception as e:
+                return f"Error reading Excel file: {str(e)}"
+    
+    
+    
+    
+        class DownloadImageTool(Tool):
+            name = "download_chess_image"
+            description = "Downloads chess position image from task ID"
+            inputs = {'task_id': {'type': 'string'}}
+            output_type = "string"
+    
+            def forward(self, task_id: str) -> str:
+                try:
+                    response = requests.get(
+                        f"https://agents-course-unit4-scoring.hf.space/files/{task_id}",
+                        stream=True
+                    )
+                    response.raise_for_status()
+                    
+                    img_path = f"chess_{task_id}.png"
+                    with open(img_path, "wb") as f:
+                        for chunk in response.iter_content(8192):
+                            f.write(chunk)
+                    return img_path
+                except Exception as e:
+                    raise RuntimeError(f"Image download failed: {str(e)}")
+    
+    
+        
+        class ChessEngineTool(Tool):
+            import chess
+            import chess.engine
+            name = "stockfish_analysis"
+            description = "Analyzes chess position using Stockfish"
+            inputs = {'fen': {'type': 'string'}}
+            output_type = "string"
+    
+            def forward(self, fen: str) -> str:
+                try:
+                    board = chess.Board(fen)
+                    engine = chess.engine.SimpleEngine.popen_uci("stockfish")
+                    result = engine.play(board, chess.engine.Limit(time=2.0))
+                    engine.quit()
+                    return board.san(result.move)
+                except Exception as e:
+                    return f"Engine error: {str(e)}"
+    
+        async def analyze_position(self, task_id: str):
+            try:
+                # Step 1: Download image
+                img_path = await self.tools[0](task_id)
+                
+                # Step 2: Get multimodal analysis
+                response = await self.model.acreate(
+                    messages=[{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": """Analyze this chess position. 
+                            It's black's turn. Provide the winning move in algebraic notation.
+                            Respond ONLY with the move, nothing else."""},
+                            {"type": "image_url", "image_url": {"url": f"file://{img_path}"}}
+                        ]
+                    }],
+                    temperature=0.1
+                )
+                
+                return response.choices[0].message.content
+                
+            except Exception as e:
+                return f"Analysis failed: {str(e)}"
+    
+    
+    
+            
+    class PythonCodeReaderTool(Tool):
+        name = "read_python_code"
+        description = "Reads a Python (.py) file and returns its content as a string."
+        inputs = {
+            "file_path": {"type": "string", "description": "The path to the Python file to read"}
+        }
+        output_type = "string"
+    
+        def forward(self, file_path: str) -> str:
+            try:
+                if not os.path.exists(file_path):
+                    return f"Error: Python file not found at {file_path}"
+                with open(file_path, "r", encoding="utf-8") as file:
+                    content = file.read()
+                return content
+            except Exception as e:
+                return f"Error reading Python file: {str(e)}"
+
+    
     async def __call__(self, question: str, task_id: str) -> str:
         """Process a question asynchronously using the MagAgent."""
         print(f"MagAgent received question (first 50 chars): {question[:50]}... Task ID: {task_id}")
         try:
-            if self.rate_limiter:
-                while not self.rate_limiter.consume(1):
-                    print(f"Rate limit reached for task {task_id}. Waiting...")
-                    await asyncio.sleep(60 / 15)  # Assuming 15 RPM
-            # Include task_id in the task prompt to guide the agent
-            task = (
+            # Unified processing flow
+            img_path = await self.download_tool(task_id)
+            
+            response = await self.model.acreate(
+                messages=[{
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": f"{question}\nProvide answer in algebraic notation."},
+                        {"type": "image_url", "image_url": {"url": f"file://{img_path}"}}
+                    ]
+                }],
+                temperature=0.1
+            )
+            
+            return response.choices[0].message.content
+            
+            
+            
+#            if self.rate_limiter:
+#                while not self.rate_limiter.consume(1):
+#                    print(f"Rate limit reached for task {task_id}. Waiting...")
+#                    await asyncio.sleep(60 / 15)  # Assuming 15 RPM
+#            # Include task_id in the task prompt to guide the agent
+#            task = (
 #                f"Answer the following question accurately and concisely: \n"
-                f"{question} \n"
-                f"If the question references an attachment, use tool to download it with task_id: {task_id}\n"
+#                f"{question} \n"
+#                f"If the question references an attachment, use tool to download it with task_id: {task_id}\n"
 #                f"Return the answer as a string."
-            )
-            print(f"Calling agent.run for task {task_id}...")
-            response = await asyncio.to_thread(
-                self.agent.run,
-                task=task
-            )
-            print(f"Agent.run completed for task {task_id}.")
-            response = str(response)
-            if not response:
-                print(f"No answer found for task {task_id}.")
-                response = "No answer found."
-            print(f"MagAgent response: {response[:50]}...")
-            return response
+#            )
+#            print(f"Calling agent.run for task {task_id}...")
+#            response = await asyncio.to_thread(
+#                self.agent.run,
+#                task=task
+#            )
+#            print(f"Agent.run completed for task {task_id}.")
+#            response = str(response)
+#            if not response:
+#                print(f"No answer found for task {task_id}.")
+#                response = "No answer found."
+#            print(f"MagAgent response: {response[:50]}...")
+#            return response
         except Exception as e:
             error_msg = f"Error processing question for task {task_id}: {str(e)}. Check API key or network connectivity."
             print(error_msg)