Update agent.py

agent.py

@@ -4,13 +4,20 @@ import os
 import re
 import pandas as pd
 from typing import Optional
-from token_bucket import Limiter
+from token_bucket import Limiter, MemoryStorage
 import yaml
 from PIL import Image
 import requests
 from io import BytesIO
 from markdownify import markdownify
 import whisper
+import time
+from youtube_transcript_api import YouTubeTranscriptApi
+from SPARQLWrapper import SPARQLWrapper, JSON
+import chess
+import chess.engine
+import shutil
+import traceback
 
 # Simulated additional tools (implementation depends on external APIs or setup)
 #@tool
@@ -23,7 +30,6 @@ import whisper
 #    """
 #    cse_id = os.environ.get("GOOGLE_CSE_ID")
 #    if not api_key or not cse_id:
-
 #        raise ValueError("GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables.")
 #    url = "https://www.googleapis.com/customsearch/v1"
 #    params = {
@@ -50,14 +56,13 @@ import whisper
 #    """
 #    # Extract URL from question using regex
 #    url_pattern = r'https?://\S+'
-
+#
 #    match = re.search(url_pattern, question)
 #    if not match:
 #        return "No image URL found in the question."
 #    image_url = match.group(0)
 #
 #    headers = {
-
 #        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
 #    }
 #    try:
@@ -71,7 +76,6 @@ import whisper
 #        tools=[],
 #        model=model,
 #        max_steps=10,
-
 #        verbosity_level=2
 #    )
 #
@@ -82,6 +86,64 @@ import whisper
 #    
 #    return f"The image description: '{response}'"
 
+class ChessEngineTool(Tool):
+    name = "chess_engine"
+    description = "Analyzes a chess position (FEN) with Stockfish and returns the best move."
+    inputs = {
+        "fen": {"type": "string", "description": "FEN string of the position."},
+        "time_limit": {"type": "number", "description": "Time in seconds for engine analysis.", "nullable": True}
+    }
+    output_type = "string"
+
+    def forward(self, fen: str, time_limit: float = 0.1) -> str:
+        # figure out where the binary actually is
+        sf_bin = shutil.which("stockfish") or "/usr/games/stockfish"
+        if not sf_bin:
+            raise RuntimeError(
+                f"Cannot find stockfish on PATH or at /usr/games/stockfish. "
+                "Did you install it in apt.txt or via apt-get?"
+            )
+
+        board = chess.Board(fen)
+        engine = chess.engine.SimpleEngine.popen_uci(sf_bin)
+        result = engine.play(board, chess.engine.Limit(time=time_limit))
+        engine.quit()
+        return board.san(result.move)
+
+class DuckDuckGoSearchTool(Tool):
+    name = "web_search"
+    description = "Searches the web using DuckDuckGo and returns results as a string."
+    inputs = {'query': {'type': 'string', 'description': 'The search query.'}}
+    output_type = "string"
+
+    def __init__(self, rate_limiter: Optional[Limiter] = None, *args, **kwargs):
+        self.is_initialized = False
+        self.rate_limiter = rate_limiter
+
+    def forward(self, query: str) -> str:
+        max_retries = 3
+        retry_delay = 2  # Seconds
+        for attempt in range(max_retries):
+            try:
+                if self.rate_limiter:
+                    while not self.rate_limiter.consume(1):
+                        print(f"Rate limit reached for web search. Waiting...")
+                        time.sleep(retry_delay)
+                # Simplified DuckDuckGo search logic (replace with actual implementation)
+                response = requests.get(f"https://lite.duckduckgo.com/lite/?q={query}", timeout=10)
+                response.raise_for_status()
+                return response.text  # Process results as needed
+            except requests.exceptions.HTTPError as e:
+                if e.response.status_code == 202:  # Rate limit
+                    print(f"Rate limit hit for web search (attempt {attempt+1}/{max_retries}). Retrying in {retry_delay}s...")
+                    time.sleep(retry_delay)
+                    retry_delay *= 2  # Exponential backoff
+                    continue
+                return f"Error searching web: {str(e)}"
+            except Exception as e:
+                return f"Error searching web: {str(e)}"
+        return "Error: Web search failed due to rate limiting."
+
 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."
@@ -124,11 +186,11 @@ class VisitWebpageTool(Tool):
 
 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), and audio (.mp3) files."
+    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"
 
-    def __init__(self, rate_limiter: Optional[Limiter] = None, default_api_url: str = "https://agents-course-unit4-scoring.hf.space", *args, **kwargs):
+    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
@@ -154,8 +216,12 @@ class DownloadTaskAttachmentTool(Tool):
                 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}"
+                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)
@@ -164,6 +230,10 @@ class DownloadTaskAttachmentTool(Tool):
                     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)}"
 
@@ -242,20 +312,24 @@ class ExcelReaderTool(Tool):
         except Exception as e:
             return f"Error reading Excel file: {str(e)}"
 
-#@tool
-#class LocalFileAudioTool:
-#    """Tool for transcribing audio files"""
-#    
-#    @tool
-#    def transcribe(self, file_path: str) -> str:
-#        """Transcribe audio from file
-#        Args:
-#            file_path (str): Path to audio file
-#        Returns:
-#            str: Transcription text
-#        """
-#        return f"Transcribed audio from '{file_path}' (simulated)."
+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):
         """Initialize the MagAgent with search tools."""
@@ -270,17 +344,24 @@ class MagAgent:
         # Load prompt templates
         with open("prompts.yaml", 'r') as stream:
             prompt_templates = yaml.safe_load(stream)
+
+        # Initialize rate limiter for DuckDuckGoSearchTool
+        search_rate_limiter = Limiter(rate=30/60, capacity=30, storage=MemoryStorage()) if not rate_limiter else rate_limiter
         
         self.agent = CodeAgent(
             model= model,
             tools=[
-#                GoogleSearchTool,
-                DownloadTaskAttachmentTool(),
-                DuckDuckGoSearchTool(),
+                DownloadTaskAttachmentTool(rate_limiter=rate_limiter),
+                DuckDuckGoSearchTool(rate_limiter=search_rate_limiter),
                 WikipediaSearchTool(),
-#                ImageAnalysisTool,
                 SpeechToTextTool,
-                ExcelReaderTool()
+                ExcelReaderTool(),
+                VisitWebpageTool(),
+                PythonCodeReaderTool()
+                # Uncomment to add ChessEngineTool (requires python-chess and Stockfish)
+                 ChessEngineTool()
+#                GoogleSearchTool,
+#                ImageAnalysisTool,
 #                LocalFileAudioTool()
             ],
             verbosity_level=3,