ThinkSquare / app.py
Falguni's picture
Update docstring
aac14ad
from typing import Optional, Tuple
import chess
import gradio as gr
from src.util.board_vis import colored_unicode_board
from src.util.pgn_util import pgn_string_to_game, read_pgn
from src.thinksqure_engine import ThinkSquareEngine
from modal import Function
import logging
logging.basicConfig(level=logging.INFO)
def get_best_move(fen: Optional[str] = None) -> str:
try:
modal_fn = Function.from_name("ThinkSquare-Backend", "get_best_move")
best_move_san = modal_fn.remote(fen)
logging.info("Best move retrieved from Modal function.")
except Exception as e:
logging.error(f"Error getting best move from Modal: {e}")
# Fallback to local engine if Modal function fails
best_move_san = ThinkSquareEngine.get_best_move(fen)
return best_move_san
def annotate_pgn_file(
file, depth_level: str = "standard", style: Optional[str] = "expert"
) -> str:
"""Annotate a chess game from a PGN file.
This function takes a PGN file and annotates the chess game using an engine.
Instructions for LLMs:
- LLMs must not use this function. Use the string-based `annotate_pgn` instead.
"""
pgn_text = read_pgn(file.name)
if depth_level == "standard":
analysis_time_per_move = 0.1
elif depth_level == "deep":
analysis_time_per_move = 0.5
else:
analysis_time_per_move = 0.1
annotated_game = annotate_pgn(
pgn_text,
analysis_time_per_move=analysis_time_per_move,
style=style,
)
return annotated_game
def suggest_move(fen: Optional[str] = None) -> str:
"""Suggest a move for the given FEN position.
This function can be used to give a hint to the user about the best move to play.
This function takes a FEN string representing the current board state and returns the best move in SAN format.
Args:
fen: The FEN string representing the board state. If None, the initial position is used.
Returns:
The best move in SAN format.
"""
if fen is None or fen == "" or fen.lower() == "none" or fen.lower() == "null":
fen = None
best_move_san = get_best_move(fen)
return best_move_san
def annotate_pgn(
pgn_input: str,
analysis_time_per_move: float = 0.1,
style: Optional[str] = "expert",
) -> str:
"""Annotate a chess game with engine analysis.
This function takes a chess game (PGN) in string format.
Instructions for LLMs:
- Do not send the file name or file path as an arument.
- LLMs must send the chess game in PGN format as a string.
- If the user provides a file, read the file and extract the PGN content.
- LLMs should display the annotated game in a Markdown code block with the label "Annotated PGN".
Args:
pgn_input: The chess game in PGN format as a string.
analysis_time_per_move: Time in seconds for engine analysis per move. Defaults to 0.1 seconds.
style: The style of annotation to use. Defaults to "expert". Other options can be "novice", "jarvis", "natural", or any custom style.
Returns:
The annotated game in PGN format as a string.
"""
if style == "":
style = None
style = str(style).lower().strip()
if not isinstance(pgn_input, chess.pgn.Game):
try:
pgn_input = pgn_string_to_game(pgn_input)
except Exception as e:
raise ValueError(f"Invalid PGN input: {e}")
game = pgn_input
try:
analysis_time_per_move = float(analysis_time_per_move)
except ValueError:
raise ValueError("Analysis time must be a number.")
# try Modal function first
try:
modal_fn = Function.from_name("ThinkSquare-Backend", "annotate")
annotated_game = modal_fn.remote(pgn_input, analysis_time_per_move, style)
logging.info("Annotated game using Modal function.")
except Exception as e:
logging.error(f"Error annotating PGN with Modal: {e}")
# Fallback to local engine if Modal function fails
annotated_game = ThinkSquareEngine.annotate(
game, analysis_time=analysis_time_per_move, llm_character=style
)
return str(annotated_game)
def render_board(fen: Optional[str] = None, render_mode: str = "ascii") -> str:
"""Render the chess board in the specified format. Default is ASCII.
Instructions for LLMs:
- LLMs should default to ascii rendering mode.
Args:
fen: The FEN string representing the board state. If None, the initial position is used.
render_mode: The rendering mode for the board. Can be "ascii", "svg", or "unicode". Defaults to "ascii".
Returns:
The rendered board as a string in the specified format.
"""
if render_mode == "ascii":
board_repr = ThinkSquareEngine.render_board_ascii(fen)
elif render_mode == "svg":
board_repr = ThinkSquareEngine.render_board_svg(fen)
elif render_mode == "unicode":
board_repr = colored_unicode_board(fen)
else:
raise ValueError("Invalid render mode. Choose 'ascii', 'svg', or 'unicode'.")
return board_repr
def play_chess(
move: Optional[str] = "",
fen: Optional[str] = "",
draw_board: bool = True,
render_mode: str = "ascii",
) -> Tuple:
"""Play a move in a chess game.
Instructions for LLMs:
Prerequisites:
- User must be asked if they want to play as white or black.
- If user chooses black, pass an empty string in the first move (for the engine to play as white).
- If the user chooses white, LLMs must ask the user for a move and pass it in the first move (for the engine to play as black).
- User must be asked if they want a board drawn.
- If they do, pass `draw_board=True` to this function.
- If they do not, pass `draw_board=False`.
- If a move is provided, it must be in long algebraic notation (e.g., "e4", "Nf3", "Bb5").
To start a new game:
- Pass empty string for fen.
- Pass empty string for engine to play as white.
- Pass a move in long algebraic notation (e.g., "e4", "Nf3", "Bb5") for engine to play as black.
To coninue a game:
- Pass the FEN string representing the board state prior to the user's last move.
- Pass a move in long algebraic notation (e.g., "e4", "Nf3", "Bb5") for engine to play the next move.
About rendering:
- LLMs must use ascii as render_mode unless otherwise specified by user. While rendering ascii, LLMs must use monospaced font.
- LLMs must explicitly pass render_mode = "ascii" to this function if they want to render the board in ASCII format.
Args:
move: The move to play in long algebraic notation. If None, the engine will play a move.
fen: The FEN string representing the board state prior to the user's last move. If None, the game starts from the initial position.
draw_board: Whether to draw the board in ASCII/Unicode/svg format. Defaults to True.
render_mode: The rendering mode for the board. Defaults to "ascii". This can be "ascii", "svg", or "unicode".
Returns:
The best move played by the engine, the updated board state in FEN notation, and a board representation if draw_board is True else None.
"""
if move is None or move == "" or move.lower() == "none" or move.lower() == "null":
move = None
if fen is None or fen == "" or fen.lower() == "none" or fen.lower() == "null":
fen = None
if move is not None:
is_valid = ThinkSquareEngine.is_valid_move(move, fen)
if not is_valid:
return "Invalid move", "", ""
fen = ThinkSquareEngine.get_fen_after_move(move, fen)
assert fen is not None, "FEN after move should not be None"
bestmove_san = get_best_move(fen)
fen_after_move = ThinkSquareEngine.get_fen_after_move(bestmove_san, fen)
if draw_board:
board_repr = render_board(fen_after_move, render_mode)
else:
board_repr = None
return bestmove_san, fen_after_move, board_repr
with gr.Blocks(title="ThinkSquare") as app:
def save_text_to_file(text):
with open("annotated_game.pgn", "w") as f:
f.write(text)
return "annotated_game.pgn"
with gr.Tab("Play Chess"):
gr.Markdown("### Play Chess with an engine")
move_input = gr.Textbox(
label="Your Move (SAN)", placeholder="e4, Nf3...", value=None
)
fen_input = gr.Textbox(
label="FEN String (optional)",
placeholder="Leave blank to start from initial position",
value=None,
)
draw_board_checkbox = gr.Checkbox(label="Draw Board", value=True)
render_mode_dropdown = gr.Dropdown(
choices=["ascii", "svg", "unicode"],
value="svg",
label="Render Mode",
visible=False,
)
play_btn = gr.Button("Submit Move")
best_move_output = gr.Textbox(label="Best Move by Engine (SAN)")
updated_fen_output = gr.Textbox(label="Updated FEN")
board_output = gr.HTML(label="Board View", visible=True)
play_btn.click(
fn=play_chess,
inputs=[move_input, fen_input, draw_board_checkbox, render_mode_dropdown],
outputs=[best_move_output, updated_fen_output, board_output],
)
with gr.Tab("Chess Game Annotation", visible=True):
def toggle_custom_input(style):
if style == "custom":
return gr.update(visible=True, interactive=True)
else:
return gr.update(visible=False, interactive=False)
gr.Markdown("### Analyze and Annotate a PGN File")
pgn_file = gr.File(label="Upload PGN", file_types=[".pgn"])
analysis_depth = gr.Radio(
label="Analysis Depth", choices=["standard", "deep"], value="standard"
)
style_dropdown = gr.Dropdown(
label="Style",
choices=[
"expert",
"novice",
"jarvis",
"natural",
"yoda",
"oracle",
"bored guy",
"angry granny",
"Sheldon Cooper (The Big Bang Theory)",
],
value="expert",
)
custom_input = gr.Textbox(label="Custom Style Prompt", visible=False)
style_dropdown.change(
fn=toggle_custom_input, inputs=style_dropdown, outputs=custom_input
)
analyze_btn = gr.Button("Annotate PGN")
annotated_pgn_file = gr.Textbox(label="Annotated PGN")
download_button = gr.Button("Download Annotated PGN")
pgn_file_output = gr.File(label="Download your annotated PGN")
download_button.click(
fn=save_text_to_file, inputs=annotated_pgn_file, outputs=pgn_file_output
)
analyze_btn.click(
fn=annotate_pgn_file,
inputs=[pgn_file, analysis_depth, style_dropdown],
outputs=annotated_pgn_file,
)
with gr.Tab("Annotate PGN", visible=False):
gr.Markdown("### Annotate a PGN String")
pgn_input = gr.Textbox(
label="PGN String",
placeholder="Paste your PGN string here",
lines=10,
value=None,
)
analysis_time_input = gr.Textbox(
label="Analysis Time per Move (seconds)", value="0.1"
)
style_input = gr.Textbox(label="Style (optional)", value="expert")
annotate_btn = gr.Button("Annotate PGN")
annotated_output = gr.Textbox(label="Annotated PGN Output", lines=10)
annotate_btn.click(
fn=annotate_pgn,
inputs=[pgn_input, analysis_time_input, style_input],
outputs=annotated_output,
)
with gr.Tab("Render Board", visible=False):
gr.Markdown("### Render Chess Board")
fen_input_render = gr.Textbox(
label="FEN String (optional)",
placeholder="Leave blank for initial position",
value=None,
)
render_mode_dropdown = gr.Dropdown(
label="Render Mode", choices=["ascii", "svg", "unicode"], value="svg"
)
render_btn = gr.Button("Render Board")
board_render_output = gr.HTML(label="Rendered Board")
render_btn.click(
fn=render_board,
inputs=[fen_input_render, render_mode_dropdown],
outputs=board_render_output,
)
with gr.Tab("Suggest Move", visible=False):
gr.Markdown("### Suggest a Move")
fen_input_suggest = gr.Textbox(
label="FEN String (optional)",
placeholder="Leave blank for initial position",
value=None,
)
suggest_btn = gr.Button("Suggest Move")
suggested_move_output = gr.Textbox(label="Suggested Move (SAN)")
suggest_btn.click(
fn=suggest_move,
inputs=fen_input_suggest,
outputs=suggested_move_output,
)
if __name__ == "__main__":
app.launch(mcp_server=True)