Add funcionalities

app.py

@@ -1,32 +1,85 @@
 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 cleanup_tmp_pgn, export_pgn, read_pgn
+from src.util.pgn_util import pgn_string_to_game, read_pgn
 from src.thinksqure_engine import ThinkSquareEngine
 
 
-def annotate_pgn(
+def annotate_pgn_file(
     file, analysis_time_per_move: float = 0.1, style: Optional[str] = "expert"
 ) -> str:
-    """Annotate a PGN file with engine analysis.
+    """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)
+
+    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:
-        file: The PGN file to be annotated.
-        analysis_time_per_move: Time in seconds for engine analysis per move.
+        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 = ThinkSquareEngine.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 PGN file.
+        The annotated game in PGN format as a string.
     """
-    cleanup_tmp_pgn()
 
     if style == "":
         style = None
 
     style = str(style).lower().strip()
 
-    game = read_pgn(file.name)
+    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)
@@ -40,13 +93,15 @@ def annotate_pgn(
     return str(annotated_game)
 
 
-def render_board(fen: Optional[str] = None, render_mode: str = "svg") -> str:
-    """Render the chess board in the specified format.
-    LLMs should use html rendering mode to render the board in SVG format.
+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 "svg".
+        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.
     """
@@ -66,37 +121,39 @@ def play_chess(
     move: Optional[str] = "",
     fen: Optional[str] = "",
     draw_board: bool = True,
-    render_mode: str = "svg",
+    render_mode: str = "ascii",
 ) -> Tuple:
     """Play a move in a chess game.
-    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, 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. LLMs should use html rendering mode to render the board in SVG format.
-            - If they do not, pass `draw_board=False`.
+    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, 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").
+            - 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 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.
+        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.
 
     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 format. Defaults to True.
-        render_mode: The rendering mode for the board. Defaults to "svg". This can be "ascii", "svg", or "unicode".
+        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". LLMs should default to "ascii".
 
     Returns:
-        The best move played by the engine, the updated board state in FEN notation, and  a string representation of the board if draw_board is True else None.
+        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":
@@ -126,6 +183,7 @@ def play_chess(
 
 
 with gr.Blocks(title="ThinkSquare") as app:
+
     def save_text_to_file(text):
         with open("annotated_game.pgn", "w") as f:
             f.write(text)
@@ -142,6 +200,12 @@ with gr.Blocks(title="ThinkSquare") as app:
             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")
 
@@ -151,11 +215,10 @@ with gr.Blocks(title="ThinkSquare") as app:
 
         play_btn.click(
             fn=play_chess,
-            inputs=[move_input, fen_input, draw_board_checkbox],
+            inputs=[move_input, fen_input, draw_board_checkbox, render_mode_dropdown],
             outputs=[best_move_output, updated_fen_output, board_output],
         )
-
-    with gr.Tab("PGN Annotation"):
+    with gr.Tab("Chess Game Annotation", visible=True):
 
         def toggle_custom_input(style):
             if style == "custom":
@@ -163,7 +226,7 @@ with gr.Blocks(title="ThinkSquare") as app:
             else:
                 return gr.update(visible=False, interactive=False)
 
-        gr.Markdown("### PGN Analyzer")
+        gr.Markdown("### Analyze and Annotate a PGN File")
         pgn_file = gr.File(label="Upload PGN", file_types=[".pgn"])
         analysis_time = gr.Textbox(
             label="Analysis Time per Move (seconds)", value="0.1"
@@ -189,11 +252,34 @@ with gr.Blocks(title="ThinkSquare") as app:
         )
 
         analyze_btn.click(
-            fn=annotate_pgn,
+            fn=annotate_pgn_file,
             inputs=[pgn_file, analysis_time, style_dropdown],
             outputs=annotated_pgn_file,
         )
 
+    with gr.Tab("Annotate PGN", visible=True):
+        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(
@@ -213,7 +299,22 @@ with gr.Blocks(title="ThinkSquare") as app:
             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)