Update agent.py

agent.py

@@ -53,7 +53,7 @@ class VisitWebpageTool(Tool):
 
     def forward(self, url: str) -> str:
         try:
-            response = requests.get(url, timeout=20)
+            response = requests.get(url, timeout=30)
             response.raise_for_status()
             markdown_content = markdownify(response.text).strip()
             markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
@@ -89,7 +89,7 @@ class DownloadTaskAttachmentTool(Tool):
                 while not self.rate_limiter.consume(1):
                     print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
                     time.sleep(4)  # Assuming 15 RPM
-            response = requests.get(file_url, stream=True, timeout=15)
+            response = requests.get(file_url, stream=True, timeout=25)
             response.raise_for_status()
             
             # Determine file extension based on Content-Type
@@ -204,167 +204,7 @@ class ExcelReaderTool(Tool):
             
         except Exception as e:
             return f"Error reading Excel file: {str(e)}"
-
-
-
-
-
-@tool
-def PNG2FENTool(png_file: str) -> str:
-    """Tool for converting a PNG file containing a chess board to a FEN position string.
-    Args:
-        png_file (str): The path to the PNG file.
-
-    Returns:
-        str: The FEN position string representing the chess board.
-    """
-#    Raises:
-#    - FileNotFoundError:
-#        If the PNG file does not exist.
-#    - ValueError:
-#        If the PNG file cannot be processed or does not contain a valid chess board.
-
-    try:
-        # Open and preprocess image with modern Pillow
-        img = Image.open(png_file)
-        img = ImageOps.exif_transpose(img).convert("L")
-        
-        # Use LANCZOS instead of ANTIALIAS
-        img = img.resize((img.width*2, img.height*2), Image.Resampling.LANCZOS)
-        
-        # Save temp file for OCR
-        temp_path = "chess_temp.png"
-        img.save(temp_path)
-        
-        # Perform OCR
-        import easyocr
-        reader = easyocr.Reader(['en'])
-        result = reader.readtext(png_file, detail=0)
-        fen_candidates = [text for text in result if validate_fen_format(text)]
-        
-        if not fen_candidates:
-            raise ValueError("No valid FEN found in image")
-            
-        return fen_candidates[0]
-        
-    except Exception as e:
-        raise ValueError(f"OCR processing failed: {str(e)}")
-    
-    
-#    try:
-#        # Open the PNG file using PIL
-#        image = Image.open(png_file)
-#
-#        # Use pytesseract to extract text from the image
-#        text = pytesseract.image_to_string(image)
-#
-#        # Process the extracted text to get the FEN position string
-#        fen_position = process_text_to_fen(text)
-#
-#        return fen_position
-#
-    except FileNotFoundError:
-        raise FileNotFoundError("PNG file not found.")
-#
-#    except Exception as e:
-#        raise ValueError("Error processing PNG file: " + str(e))
-
-def process_text_to_fen(text):
-    """
-    Processes the extracted text from the image to obtain the FEN position string.
-
-    Parameters:
-    - text: str
-        The extracted text from the image.
-
-    Returns:
-    - str:
-        The FEN position string representing the chess board.
-
-    Raises:
-    - ValueError:
-        If the extracted text does not contain a valid chess board.
-    """
-
-    # Process the text to remove any unnecessary characters or spaces
-    processed_text = text.strip().replace("\n", "").replace(" ", "")
-
-    # Check if the processed text matches the expected format of a FEN position string
-    if not validate_fen_format(processed_text):
-        raise ValueError("Invalid chess board.")
-
-    return processed_text
-
-def validate_fen_format(fen_string):
-    """
-    Validates if a given string matches the format of a FEN (Forsyth–Edwards Notation) position string.
-
-    Parameters:
-    - fen_string: str
-        The string to be validated.
-
-    Returns:
-    - bool:
-        True if the string matches the FEN format, False otherwise.
-    """
-
-    # FEN format: 8 sections separated by '/'
-    sections = fen_string.split("/")
-    if len(sections) != 8:
-        return False
-
-    # Check if each section contains valid characters
-    for section in sections:
-        if not validate_section(section):
-            return False
-
-    return True
-
-def validate_section(section):
-    """
-    Validates if a given section of a FEN (Forsyth–Edwards Notation) position string contains valid characters.
-
-    Parameters:
-    - section: str
-        The section to be validated.
-
-    Returns:
-    - bool:
-        True if the section contains valid characters, False otherwise.
-    """
-
-    # Valid characters: digits 1-8 or letters 'r', 'n', 'b', 'q', 'k', 'p', 'R', 'N', 'B', 'Q', 'K', 'P'
-    valid_chars = set("12345678rnbqkpRNBQKP")
-    return all(char in valid_chars for char in section)
-
-import chess
-import chess.engine
-
-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 PythonCodeReaderTool(Tool):
     name = "read_python_code"
     description = "Reads a Python (.py) file and returns its content as a string."
@@ -382,7 +222,20 @@ class PythonCodeReaderTool(Tool):
             return content
         except Exception as e:
             return f"Error reading Python file: {str(e)}"
-            
+
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+#from smolagents.tools import DuckDuckGoSearchException  # Replace with the actual exception if different
+
+class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
+    @retry(
+        stop=stop_after_attempt(3),  # Retry up to 3 times
+        wait=wait_exponential(multiplier=1, min=4, max=10),  # Wait 4s, 8s, then 10s
+        retry=retry_if_exception_type(Exception)  # Retry on any exception
+    )
+    def forward(self, query: str) -> str:
+        return super().forward(query)
+
+
 class MagAgent:
     def __init__(self, rate_limiter: Optional[Limiter] = None):
         """Initialize the MagAgent with search tools."""
@@ -407,26 +260,26 @@ class MagAgent:
             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
+        search_rate_limiter = Limiter(rate=10/60, capacity=10, storage=MemoryStorage()) if not rate_limiter else rate_limiter
         
         self.agent = CodeAgent(
             model= model,
             tools=[
                 DownloadTaskAttachmentTool(rate_limiter=rate_limiter),
-#                DuckDuckGoSearchTool(),
+                RetryDuckDuckGoSearchTool(),
 #                WikipediaSearchTool(),
                 SpeechToTextTool(),
                 ExcelReaderTool(),
                 VisitWebpageTool(),
                 PythonCodeReaderTool(),
-                PNG2FENTool,
-                ChessEngineTool(),
+#                PNG2FENTool,
+#                ChessEngineTool(),
 #                GoogleSearchTool,
 #                ImageAnalysisTool,
             ],
             verbosity_level=2,
             prompt_templates=prompt_templates,
-            add_base_tools=True,
+            add_base_tools=False,
             max_steps=15
         )
         print("MagAgent initialized.")