File size: 8,147 Bytes
476b409 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
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:
# 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
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)
|