tools/chess_tools.py

import shutil
from smolagents import Tool
import requests

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 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."
        }
    }
    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)