Create tools.py

tools.py

@@ -0,0 +1,363 @@
+# tools.py
+import os
+import json
+import logging
+import re
+import requests
+import shutil
+import urllib.parse
+import pandas as pd # For ExcelParsingTool
+from board_to_fen.predict import get_fen_from_image_path # For ChessBoardFENTool
+from google import genai
+from google.genai import types
+# from litellm import completion # Removed - no longer used for ConvertChessMoveTool
+from smolagents import Tool
+from settings import Settings
+from models import GoogleModelID # Import GoogleModelID
+
+logger = logging.getLogger(__name__)
+
+class BaseCustomTool(Tool):
+    """Base class for custom tools to easily pass settings."""
+    def __init__(self, settings: Settings):
+        super().__init__()
+        self.settings = settings
+        
+class GetTaskFileTool(BaseCustomTool):
+    name = "get_task_file_tool"
+    description = """If a file_name is provided in the task, use this tool to download the file associated with a given task_id. Returns the absolute file path to the downloaded file. This path can then be used by other tools like AudioUnderstandingTool or ExcelParsingTool. Example: get_task_file_tool(task_id="1234", file_name="example.mp3")"""
+    inputs = {
+        "task_id": {"type": "string", "description": "Task ID (required)"},
+        "file_name": {"type": "string", "description": "File name (required)"},
+    }
+    output_type = "string"
+
+    def __init__(self, settings: Settings):
+        super().__init__(settings)
+        self.directory_name = "downloads"
+        self.create_dir()
+        
+    def forward(self, task_id: str, file_name: str) -> str:
+        try:
+            # Use the scoring API base URL for file downloads
+            response = requests.get(f"{self.settings.scoring_api_base_url}/files/{task_id}", timeout=15)
+            response.raise_for_status()
+            
+            # Ensure the downloads directory exists
+            os.makedirs(self.directory_name, exist_ok=True)
+            
+            file_path = os.path.join(self.directory_name, file_name)
+            with open(file_path, 'wb') as file:
+                file.write(response.content)
+            
+            absolute_file_path = os.path.abspath(file_path)
+            logger.info(f"Successfully downloaded file '{file_name}' for task_id {task_id} to {absolute_file_path}")
+            return absolute_file_path
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Error downloading file for task_id {task_id} from API: {e}")
+            # Fallback to local 'files' directory if API download fails
+            local_file_path = os.path.join("files", file_name)
+            if os.path.exists(local_file_path):
+                destination_path = os.path.join(self.directory_name, file_name)
+                os.makedirs(self.directory_name, exist_ok=True)
+                shutil.copy2(local_file_path, destination_path)
+                absolute_local_file_path = os.path.abspath(destination_path)
+                logger.info(f"Copied local fallback file '{file_name}' to {absolute_local_file_path}")
+                return absolute_local_file_path
+            else:
+                logger.error(f"Local fallback file '{local_file_path}' not found.")
+                return f"Error: Could not download or find file '{file_name}' for task_id {task_id}. {e}"
+        except Exception as e:
+            logger.error(f"An unexpected error occurred in GetTaskFileTool: {e}")
+            return f"Error: An unexpected error occurred while getting file '{file_name}'. {e}"
+        
+    def create_dir(self):
+        """Creates the download 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 = "video_understanding_tool"
+    description = "Analyzes a YouTube video given its URL and a specific prompt/question about its content. Returns a text description or answer from the video. Use this for tasks involving video content. Example: video_understanding_tool(youtube_url=\"https://www.youtube.com/watch?v=VIDEO_ID\", prompt=\"What is the main topic of this video?\")"
+    inputs = {
+        "youtube_url": {"type": "string", "description": "The URL of the YouTube video"},
+        "prompt": {"type": "string", "description": "A question or request regarding the video content"},
+    }
+    output_type = "string"
+
+    def __init__(self, settings: Settings, model: GoogleModelID):
+        super().__init__(settings)
+        self.model = model
+        # Initialize Google GenAI client with API key
+        genai.configure(api_key=self.settings.gemini_api_key.get_secret_value())
+        logger.info(f"VideoUnderstandingTool initialized with model: {self.model}")
+        
+    def forward(self, youtube_url: str, prompt: str) -> str:
+        try:
+            # Use the genai.GenerativeModel for multimodal content
+            model_instance = genai.GenerativeModel(self.model)
+            
+            # Create a FileData part from the YouTube URL
+            video_file_data = types.Part(
+                file_data=types.FileData(
+                    file_uri=youtube_url,
+                    mime_type="video/mp4" # Assuming common video type, adjust if needed
+                )
+            )
+            
+            # Generate content with both video and text prompt
+            response = model_instance.generate_content(
+                contents=[video_file_data, types.Part(text=prompt)]
+            )
+            
+            return response.text
+        except Exception as e:
+            logger.error(f"Error understanding video from URL '{youtube_url}': {e}")
+            return f"Error understanding video: {e}"
+
+class AudioUnderstandingTool(BaseCustomTool):
+    name = "audio_understanding_tool"
+    description = "Analyzes a local audio file given its file path and a specific prompt/question about its content. Returns a text description or answer from the audio. Use this for tasks involving audio files. You must first download the audio file using 'get_task_file_tool'. Example: audio_understanding_tool(file_path=\"/tmp/audio.mp3\", prompt=\"What are the key ingredients mentioned?\")"
+    inputs = {
+        "file_path": {"type": "string", "description": "The local file path of the audio file (e.g., from get_task_file_tool)."},
+        "prompt": {"type": "string", "description": "A question or request regarding the audio content."},
+    }
+    output_type = "string"
+
+    def __init__(self, settings: Settings, model: GoogleModelID):
+        super().__init__(settings)
+        self.model = model
+        # Initialize Google GenAI client with API key
+        genai.configure(api_key=self.settings.gemini_api_key.get_secret_value())
+        logger.info(f"AudioUnderstandingTool initialized with model: {self.model}")
+
+    def forward(self, file_path: str, prompt: str) -> str:
+        try:
+            # Upload the local audio file to Gemini Files API
+            mp3_file = genai.upload_file(path=file_path)
+            logger.info(f"Uploaded audio file: {mp3_file.uri}")
+
+            # Use the genai.GenerativeModel for multimodal content
+            model_instance = genai.GenerativeModel(self.model)
+            
+            # Generate content with both audio file and text prompt
+            response = model_instance.generate_content(
+                contents=[mp3_file, types.Part(text=prompt)]
+            )
+            
+            # Delete the uploaded file from Gemini Files API (optional, but good practice)
+            # genai.delete_file(mp3_file.name) # This might require a separate API call or context manager
+
+            return response.text
+        except Exception as e:
+            logger.error(f"Error understanding audio from file '{file_path}': {e}")
+            return f"Error understanding audio: {e}"
+
+class ExcelParsingTool(BaseCustomTool):
+    name = "excel_parsing_tool"
+    description = "Parses an Excel (.xlsx) file given its local file path. It reads the first sheet and returns its content as a CSV formatted string. Use this for tasks involving Excel data. You must first download the Excel file using 'get_task_file_tool'. Example: excel_parsing_tool(file_path=\"/tmp/sales_data.xlsx\")"
+    inputs = {"file_path": {"type": "string", "description": "The local path to the Excel file (e.g., from get_task_file_tool)."}}
+    output_type = "string"
+
+    def __init__(self, settings: Settings):
+        super().__init__(settings)
+        logger.info("ExcelParsingTool initialized.")
+
+    def forward(self, file_path: str) -> str:
+        """
+        Reads an Excel file and returns its content (first sheet) as a CSV string.
+        """
+        try:
+            # Ensure the file exists before trying to read
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"Excel file not found at: {file_path}")
+
+            df = pd.read_excel(file_path)
+            csv_content = df.to_csv(index=False)
+            logger.info(f"Successfully parsed Excel file: {file_path}")
+            return csv_content
+        except Exception as e:
+            logger.error(f"Error parsing Excel file {file_path}: {e}")
+            return f"Error parsing Excel file: {e}"
+
+class ConvertChessMoveTool(BaseCustomTool):
+    name = "convert_chess_move_tool"
+    description = "Converts a chess move from coordinate notation (e.g., 'e2e4') to standard algebraic notation. Requires the current piece placement as plain text. Example: convert_chess_move_tool(piece_placement=\"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR\", move=\"e2e4\")"
+    inputs = {
+        "piece_placement": {"type": "string", "description": "The chess piece placement in plain text (e.g., a FEN board part)."},
+        "move": {"type": "string", "description": "The move in coordinate notation (e.g., 'e2e4')"},
+    }
+    output_type = "string"
+
+    def __init__(self, settings: Settings, model: GoogleModelID): # Changed model type to GoogleModelID
+        super().__init__(settings)
+        self.model = model
+        genai.configure(api_key=self.settings.gemini_api_key.get_secret_value()) # Configure genai for this tool
+        logger.info(f"ConvertChessMoveTool initialized with model: {self.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 board state for context: {piece_placement}. "
+            "Do not provide any additional thinking or commentary in the response, "
+            "return only the algebraic notation for the move."
+            )
+        messages = [{ "content": move_message, "role": "user"}]
+        try:
+            model_instance = genai.GenerativeModel(self.model) # Use genai.GenerativeModel
+            response = model_instance.generate_content(
+                contents=messages[0]['content'] # Pass content directly
+            )
+            return response.text
+        except Exception as e:
+            logger.error(f"Error converting chess move: {e}")
+            return f"Error converting chess move: {e}"
+
+class BestChessMoveTool(BaseCustomTool):
+    name = "best_chess_move_tool"
+    description = "Gets the best chess move in coordinate notation (e.g., 'e2e4') based on a FEN (Forsyth-Edwards Notation) representation of the chess position. Example: best_chess_move_tool(fen=\"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1\")"
+    inputs = {
+        "fen": {"type": "string", "description": "The FEN (Forsyth-Edwards Notation) representation of the chess position. Example: 'rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1'"},
+    }
+    output_type = "string"
+
+    def forward(self, fen: str) -> str:
+        try:
+            url = f"{self.settings.chess_eval_url}?fen={urllib.parse.quote(fen)}&depth=15" # Depth 15 for reasonable accuracy
+            response = requests.get(url, timeout=15)
+            response.raise_for_status() # Raise HTTPError for bad responses
+            
+            response_json = response.json()
+            if response_json.get('success') == True and 'bestmove' in response_json:
+                # Stockfish API often returns "bestmove e2e4 ponder e7e5"
+                # We need only the move itself, which is the second part
+                best_move = response_json['bestmove'].split()[1]
+                logger.info(f"Successfully retrieved best chess move: {best_move} for FEN: {fen}")
+                return best_move
+            else:
+                raise ValueError(f"Stockfish API returned unsuccessful response or missing 'bestmove': {response_json}")
+        except Exception as e:
+            logger.error(f"Error getting best chess move for FEN '{fen}': {e}")
+            return f"Error getting best chess move: {e}"
+
+class ChessBoardFENTool(Tool):
+    name = "chess_board_fen_tool"
+    description = "Generates the FEN (Forsyth-Edwards Notation) representation from a local image file of a chess board and the player whose turn it is. Returns the FEN string. You must first download the image file using 'get_task_file_tool'. Example: chess_board_fen_tool(image_path=\"/tmp/board.png\", player_turn=\"b\")"
+    inputs = {
+        "image_path": {"type": "string", "description": "The local file path of the chess board image (e.g., from get_task_file_tool)."},
+        "player_turn": {"type": "string", "description": "The player with the next turn in the match, must be 'w' (white) or 'b' (black)."}
+    }
+    output_type = "string"
+    
+    def _expand_fen_rank(self, rank_str):
+        """Expands a single rank string from FEN notation into a list of 8 characters."""
+        expanded_rank = []
+        for char in rank_str:
+            if char.isdigit():
+                expanded_rank.extend([' '] * int(char))
+            else:
+                expanded_rank.append(char)
+        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."""
+        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 empty_count > 0:
+                    compressed_rank += str(empty_count)
+                    empty_count = 0
+                compressed_rank += char
+        if empty_count > 0:
+            compressed_rank += str(empty_count)
+        return compressed_rank
+
+    def _invert_mirror_fen(self, fen_string: str) -> str:
+        """
+        Takes a FEN string, inverts the board vertically, mirrors it horizontally,
+        and returns the new FEN string representing this transformed view.
+        This is often needed to convert board_to_fen output to Stockfish compatible FEN.
+        """
+        try:
+            parts = fen_string.strip().split(' ')
+            if len(parts) != 6:
+                raise ValueError("FEN string must have 6 space-separated fields (board, turn, castling, ep, halfmove, fullmove).")
+            board_part = parts[0]
+            other_parts = parts[1:]
+
+            rank_strings = board_part.split('/')
+            if len(rank_strings) != 8:
+                raise ValueError("FEN board part must have 8 ranks separated by '/'.")
+
+            original_board = [self._expand_fen_rank(r) for r in rank_strings]
+            transformed_board = [[' ' for _ in range(8)] for _ in range(8)]
+
+            for r in range(8):
+                for c in range(8):
+                    transformed_board[7 - r][7 - c] = original_board[r][c]
+
+            new_rank_strings = [self._compress_fen_rank(row) for row in transformed_board]
+            new_board_part = "/".join(new_rank_strings)
+
+            return " ".join([new_board_part] + other_parts)
+
+        except Exception as e:
+            logger.error(f"Error processing FEN for inversion/mirroring: {e}. Input: '{fen_string}'")
+            return f"Error processing FEN: {e}"
+
+    def _add_fen_game_state(self, board_placement: str,
+                        side_to_move: str,
+                        castling: str = "-",
+                        en_passant: str = "-",
+                        halfmove_clock: int = 0,
+                        fullmove_number: int = 1) -> str:
+        """
+        Appends standard game state information to a FEN board placement string.
+        """
+        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}'"
+
+        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).")
+
+        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:
+        try:
+            board_placement = get_fen_from_image_path(image_path)
+            
+            # Add game state to the board placement
+            board_fen_with_state = self._add_fen_game_state(board_placement, player_turn)
+            
+            # Inversion makes board_to_fen output Stockfish compatible
+            board_fen_inverted = self._invert_mirror_fen(board_fen_with_state) 
+            
+            logger.info(f"Generated FEN from image '{image_path}': {board_fen_inverted}")
+            return board_fen_inverted
+        except Exception as e:
+            logger.error(f"Error generating FEN from image '{image_path}': {e}")
+            return f"Error generating FEN from image: {e}"
+