|
|
import os |
|
|
import json |
|
|
import logging |
|
|
logger = logging.getLogger(__name__) |
|
|
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 = """Download the file_name associated with a given task_id. Get absolute file path""" |
|
|
inputs = { |
|
|
"task_id": {"type": "string", "description": "Task ID"}, |
|
|
"file_name": {"type": "string", "description": "File name"}, |
|
|
} |
|
|
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: |
|
|
|
|
|
shutil.copy2(f"files/{file_name}", f"{self.directory_name}/{file_name}") |
|
|
return f"{self.directory_name}/{file_name}" |
|
|
|
|
|
def create_dir(self): |
|
|
|
|
|
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 |
|
|
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 ChessGameFenTool(BaseCustomTool): |
|
|
name = "ChessGameFen" |
|
|
description = "Get a FEN representation given chess piece placement and a move." |
|
|
inputs = { |
|
|
"piece_placement": {"type": "string", "description": "The chess piece placement in plain text"}, |
|
|
"player_turn": {"type": "string", |
|
|
"description": "The player with the next turn in the match, black or white"}, |
|
|
} |
|
|
output_type = "string" |
|
|
|
|
|
def __init__(self, settings, model): |
|
|
super().__init__(settings) |
|
|
self.model = model |
|
|
|
|
|
def forward(self, piece_placement: str, player_turn: str) -> str: |
|
|
"""Use the tool.""" |
|
|
fen_message = f"""Assuming {player_turn} has the next turn, Use the following placement |
|
|
{piece_placement} and provide the board state as FEN. Do not provide any |
|
|
additional thinking or commentary in the response, the FEN only.""" |
|
|
messages = [{ "content": fen_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 ChessPiecePlacementTool(Tool): |
|
|
name = "ChessPiecePlacement" |
|
|
description = "Get chess piece placement information from an image of a board." |
|
|
inputs = { |
|
|
"image_path": {"type": "string", "description": "The local file of the chess board image"}, |
|
|
} |
|
|
output_type = "string" |
|
|
|
|
|
def forward(self, image_path: str) -> str: |
|
|
return get_fen_from_image_path(image_path) |
|
|
|
