Update agent.py

agent.py

@@ -245,6 +245,139 @@ class ExcelReaderTool(Tool):
         except Exception as e:
             return f"Error reading Excel file: {str(e)}"
 
+
+
+import pytesseract
+
+@tool
+def PNG2FENTool(png_file) -> 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 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."
@@ -292,6 +425,8 @@ class MagAgent:
                 ExcelReaderTool(),
                 VisitWebpageTool(),
                 PythonCodeReaderTool(),
+                PNG2FENTool(),
+                ChessEngineTool(),
 #                GoogleSearchTool,
 #                ImageAnalysisTool,
             ],