Update tools/chess_tools.py

tools/chess_tools.py

@@ -1,88 +1,33 @@
-import shutil
-from smolagents import Tool
+from tools.tool_base import BaseTool
+from PIL import Image
 import requests
+from io import BytesIO
 
-class ChessEngineLocatorTool(Tool):
-    name = "locate_chess_engine"
-    description = "Locates the installed Stockfish chess engine binary on the system."
-    inputs = {}
-    output_type = "string"
-
-    def forward(self) -> str:
-        engine_path = shutil.which("stockfish")
-        return engine_path or "Stockfish engine not found."
-
+class ImageToChessBoardFENTool(BaseTool):
+    name = "image_to_chess_fen"
+    description = "Extracts the FEN representation of a chess board from an image."
 
-class ChessboardImageToFENTool(Tool):
-    name = "image_to_fen"
-    description = "Analyzes a chessboard image and returns its FEN representation."
     inputs = {
         "image_url": {
             "type": "string",
-            "description": "A public image URL or a base64 data URL containing a chessboard photo."
+            "description": "The URL of the chessboard image.",
         }
     }
     output_type = "string"
 
-    def __init__(self, endpoint_url: str, hf_token: str, **kwargs):
-        super().__init__(**kwargs)
-        self.endpoint_url = endpoint_url
-        self.hf_token = hf_token
-
     def forward(self, image_url: str) -> str:
-        headers = {
-            "Authorization": f"Bearer {self.hf_token}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "inputs": [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "input_text", "text": "Extract chess piece positions from this image."},
-                        {"type": "input_image", "image_url": image_url}
-                    ]
-                },
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "input_text",
-                            "text": (
-                                "Using the extracted pieces, list each one with its position in format <piece><square> (e.g., Kd4).\n"
-                                "Use uppercase for white, lowercase for black. Output only the lines with positions."
-                            )
-                        }
-                    ]
-                }
-            ]
-        }
-
-        response = requests.post(self.endpoint_url, headers=headers, json=payload)
-        response.raise_for_status()
-        raw_text = response.json().get("generated_text", "")
-
-        squares = {}
-        for line in raw_text.splitlines():
-            line = line.strip()
-            if len(line) == 3:
-                squares[line[1:3]] = line[0]
-
-        fen_rows = []
-        for rank in range(8, 0, -1):
-            row = ""
-            empty = 0
-            for file in "abcdefgh":
-                pos = f"{file}{rank}"
-                if pos in squares:
-                    if empty > 0:
-                        row += str(empty)
-                        empty = 0
-                    row += squares[pos]
-                else:
-                    empty += 1
-            if empty:
-                row += str(empty)
-            fen_rows.append(row)
-
-        return "/".join(fen_rows)
+        try:
+            response = requests.get(image_url)
+            response.raise_for_status()
+            image = Image.open(BytesIO(response.content)).convert("RGB")
+        except Exception as e:
+            return f"Failed to load image: {e}"
+
+        # This is a placeholder implementation.
+        # You should replace it with an actual model or logic that processes the image and returns the FEN.
+        fen = self._mock_fen_extraction(image)
+        return fen
+
+    def _mock_fen_extraction(self, image):
+        # Replace this logic with your actual FEN model prediction if available.
+        return "r1bqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"