import os import json import logging logger = logging.getLogger(__name__) import re import requests import shutil from typing import Any import urllib.parse from board_to_fen.predict import get_fen_from_image_path from google import genai from google.genai import types from litellm import completion from smolagents import Tool from settings import Settings class BaseCustomTool(Tool): def __init__(self, settings): super().__init__() self.settings = settings class GetTaskFileTool(BaseCustomTool): name = "get_task_file_tool" description = """If a file_name is provided, download file associated with a given task_id. Get absolute file path""" inputs = { "task_id": {"type": "string", "description": "Task ID (required)"}, "file_name": {"type": "string", "description": "File name (required)"}, } output_type = "string" def __init__(self, settings): super().__init__(settings) self.directory_name = "downloads" self.create_dir() def forward(self, task_id: str, file_name: str) -> str: try: response = requests.get(f"{self.settings.evaluation_api_base_url}/files/{task_id}", timeout=15) response.raise_for_status() with open(f"{self.directory_name}/{file_name}", 'wb') as file: file.write(response.content) return os.path.abspath(f"{self.directory_name}/{file_name}") except Exception as e: # Fetch the local file instead, dealing with rate limits, etc. shutil.copy2(f"files/{file_name}", f"{self.directory_name}/{file_name}") return f"{self.directory_name}/{file_name}" def create_dir(self): # Create the directory if it doesn't exist if not os.path.exists(self.directory_name): os.makedirs(self.directory_name) logger.info(f"Directory '{self.directory_name}' created successfully.") else: logger.debug(f"Directory '{self.directory_name}' already exists.") class VideoUnderstandingTool(BaseCustomTool): name = "VideoUnderstanding" description = "Prompt a YouTube video with questions to understand its content." inputs = { "youtube_url": {"type": "string", "description": "The URL of the YouTube video"}, "prompt": {"type": "string", "description": "A question or request regarding the video"}, } output_type = "string" def __init__(self, settings, model): super().__init__(settings) self.model = model def forward(self, youtube_url: str, prompt: str) -> str: client = genai.Client(api_key=self.settings.gemini_api_key.get_secret_value()) try: video_description = client.models.generate_content( model=self.model, contents=types.Content( parts=[ types.Part( file_data=types.FileData(file_uri=youtube_url) ), types.Part(text=prompt) ] ) ) return video_description.text except Exception as e: logger.error(f"Error understanding video: {e}") return False class AudioUnderstandingTool(BaseCustomTool): name = "AudioUnderstanding" description = "Prompt a local audio file with questions to understand its content." inputs = { "file_path": {"type": "string", "description": "The local file of the audio"}, "prompt": {"type": "string", "description": "A question or request regarding the audio"}, } output_type = "string" def __init__(self, settings, model): super().__init__(settings) self.model = model def forward(self, file_path: str, prompt: str) -> str: client = genai.Client(api_key=self.settings.gemini_api_key.get_secret_value()) try: mp3_file = client.files.upload(file=f"{file_path}") audio_description = client.models.generate_content( model=self.model, contents=[prompt, mp3_file] ) return audio_description.text except Exception as e: logger.error(f"Error understanding audio: {e}") return False class ConvertChessMoveTool(BaseCustomTool): name = "ConvertChessMove" description = "Convert a chess move from coordinate notation to algebraic notation." inputs = { "piece_placement": {"type": "string", "description": "The chess piece placement in plain text"}, "move": {"type": "string", "description": "The move in coordinate notation (e.g., e2e4)"}, } output_type = "string" def __init__(self, settings, model): super().__init__(settings) self.model = model def forward(self, piece_placement: str, move: str) -> str: move_message = ( f"Convert this chess move from coordinate notation to algebraic " f"notation: {move}. Use the following {piece_placement}. Do not provide any additional " "thinking or commentary in the response, the algebraic notation only." ) messages = [{ "content": move_message, "role": "user"}] response = completion( model=self.model, temperature=0.0, messages=messages, api_key=self.settings.openrouter_api_key.get_secret_value() ) return response.choices[0].message.content class BestChessMoveTool(BaseCustomTool): name = "BestChessMove" description = "Get best chess move in coordinate notation based on a FEN representation." inputs = { "fen": {"type": "string", "description": "The FEN (Forsyth-Edwards Notation) \ representation of the chess position. Example \ rn1q1rk1/pp2b1pp/2p2n2/3p1pB1/3P4/1QP2N2/PP1N1PPP/R4RK1 b - - 1 11"}, } output_type = "string" def forward(self, fen: str) -> str: try: url = f"{self.settings.chess_eval_url}?fen={urllib.parse.quote(fen)}&depth=15" response = requests.get(url, timeout=15) if response.status_code == 200 and json.loads(response.text)['success'] == True: return json.loads(response.text)['bestmove'].split()[1] else: raise ValueError(f"Error getting chess evaluation: {response.status_code}") except Exception as e: logger.error(f"Error getting chess evaluation: {e}") class ChessBoardFENTool(Tool): name = "ChessBoardFEN" description = "Get the FEN representation from an image of a chess board and a player turn." inputs = { "image_path": {"type": "string", "description": "The local file of the chess board image"}, "player_turn": {"type": "string", "description": "The player with the next turn in the match, black or white"} } output_type = "string" def _expand_fen_rank(self, rank_str): """ Expands a single rank string from FEN notation (e.g., 'p2b4') into a list of 8 characters representing the squares. Uses ' ' for empty squares. """ expanded_rank = [] for char in rank_str: if char.isdigit(): # Add number of empty squares specified by the digit expanded_rank.extend([' '] * int(char)) else: # Add the piece character expanded_rank.append(char) # Validate rank length if len(expanded_rank) != 8: raise ValueError(f"Invalid FEN rank string (length != 8): {rank_str}") return expanded_rank def _compress_fen_rank(self, rank_list): """ Compresses a list of 8 characters (representing a rank) back into FEN rank notation (e.g., turns [' ', 'K', ...] into '1K6'). Assumes ' ' represents an empty square. """ if len(rank_list) != 8: raise ValueError(f"Invalid rank list (length != 8): {rank_list}") compressed_rank = "" empty_count = 0 for char in rank_list: if char == ' ': empty_count += 1 else: # If we encountered a piece after empty squares, add the count if empty_count > 0: compressed_rank += str(empty_count) empty_count = 0 # Add the piece compressed_rank += char # If the rank ends with empty squares, add the final count if empty_count > 0: compressed_rank += str(empty_count) return compressed_rank def _invert_mirror_fen(self, fen_string): """ Takes a FEN string, inverts the board vertically, mirrors it horizontally, and returns the new FEN string representing this transformed view. The other FEN fields (turn, castling, etc.) are preserved. """ try: # 1. Split FEN into parts parts = fen_string.strip().split(' ') if len(parts) != 6: raise ValueError("FEN string must have 6 space-separated fields.") board_part = parts[0] other_parts = parts[1:] # Side-to-move, castling, ep, halfmove, fullmove # 2. Parse the board part into an 8x8 representation rank_strings = board_part.split('/') if len(rank_strings) != 8: raise ValueError("FEN board part must have 8 ranks separated by '/'.") # original_board[0] corresponds to rank 8, original_board[7] to rank 1 original_board = [self._expand_fen_rank(r) for r in rank_strings] # 3. Create a new empty 8x8 board for the transformed state # Using ' ' as the placeholder for empty squares transformed_board = [[' ' for _ in range(8)] for _ in range(8)] # 4. Apply the inversion (vertical flip) and mirror (horizontal flip) for r in range(8): # Iterate through original rows (ranks 8 down to 1) for c in range(8): # Iterate through original columns (files a to h) # The piece at original [r][c] moves to transformed [7-r][7-c] transformed_board[7 - r][7 - c] = original_board[r][c] # 5. Generate the new FEN board string from the transformed board # Read ranks from top (index 0 = rank 8) to bottom (index 7 = rank 1) new_rank_strings = [self._compress_fen_rank(row) for row in transformed_board] new_board_part = "/".join(new_rank_strings) # 6. Reassemble the full FEN string return " ".join([new_board_part] + other_parts) except Exception as e: # Return error message if parsing or processing fails return f"Error processing FEN: {e}. Input: '{fen_string}'" def _add_fen_game_state(self, board_placement, side_to_move, castling="-", en_passant="-", halfmove_clock=0, fullmove_number=1): """ Appends standard game state information to a FEN board placement string. Args: board_placement (str): The board layout part of the FEN string (e.g., "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR"). side_to_move (str): The active color ('w' for White, 'b' for Black). Case-insensitive, will be converted to lowercase. castling (str, optional): Castling availability string (e.g., "KQkq", "-"). Defaults to "-". en_passant (str, optional): En passant target square string (e.g., "e3", "-"). Defaults to "-". halfmove_clock (int, optional): The number of halfmoves since the last capture or pawn advance. Defaults to 0. fullmove_number (int, optional): The number of the full move. Starts at 1 and increments after Black's move. Defaults to 1. Returns: str: The complete FEN string including the game state, or an error message string if inputs are invalid. """ # Validate side_to_move side_to_move_lower = str(side_to_move).lower() if side_to_move_lower not in ['w', 'b']: return f"Error: side_to_move must be 'w' or 'b', received '{side_to_move}'" # Validate clock values (should be non-negative integers, fullmove >= 1) try: halfmove_clock = int(halfmove_clock) fullmove_number = int(fullmove_number) if halfmove_clock < 0: raise ValueError("halfmove_clock cannot be negative.") if fullmove_number < 1: raise ValueError("fullmove_number must be 1 or greater.") except (ValueError, TypeError): return (f"Error: halfmove_clock ('{halfmove_clock}') and " f"fullmove_number ('{fullmove_number}') must be valid integers " f"(non-negative and positive respectively).") # Assemble the full FEN string using the validated/defaulted values # Note: castling and en_passant strings are used directly as passed or defaulted. # More complex validation could be added for them if needed. full_fen = (f"{board_placement} {side_to_move_lower} {castling} " f"{en_passant} {halfmove_clock} {fullmove_number}") return full_fen def forward(self, image_path: str, player_turn: str) -> str: board_placement = get_fen_from_image_path(image_path) # Inversion makes board_to_fen output Stockfish compatible board_fen = self._add_fen_game_state(board_placement, player_turn) board_fen_inverted = self._invert_mirror_fen(board_fen) return board_fen_inverted