Add annotation support in natural language with LLMs

src/llm/base_llm_wrapper.py

@@ -0,0 +1,104 @@
+from abc import ABC, abstractmethod
+from typing import List
+from ThinkSquare.src.llm.data_models.typeddict_data_models import MultiCommentModel
+from src.llm.prompts.prompt_head import (
+    prompt_head_expert,
+    prompt_head_jarvis,
+    prompt_head_novice,
+    prompt_head_natural,
+    prompt_head_generic,
+)
+from src.llm.prompts.prompt_core import prompt_core
+
+
+class BaseLLMWrapper(ABC):
+    def comment(
+        self,
+        character: str,
+        game,
+        comment_refs: List[int],
+        move_nums: List,
+        played_moves: List,
+        played_by: List,
+        comments: List,
+        move_suggestions: List,
+        pre_eval_scores: List,
+        post_eval_scores: List,
+    ) -> MultiCommentModel:
+        """
+        Rewrite the comment in a specific tone or style.
+
+        Args:
+            character (str): The character or style in which to rewrite the comment.
+            move_nums (List): List of move numbers corresponding to the comments.
+            comments (List): List of comments to be rewritten.
+            move_suggestions (List): List of suggested moves, if any.
+        Returns:
+            MultiCommentModel: A model containing the rewritten comments.
+
+        """
+        character = character.lower().strip()
+        if character == "natural":
+            prompt_head = prompt_head_natural
+        elif character == "jarvis":
+            prompt_head = prompt_head_jarvis
+        elif character == "novice":
+            prompt_head = prompt_head_novice
+        elif character == "expert":
+            prompt_head = prompt_head_expert
+        else:
+            prompt_head = prompt_head_generic.format(character=character)
+
+        comments_info = []
+        for (
+            comment_ref,
+            move_num,
+            played_move,
+            played_by_player,
+            comment,
+            suggestion,
+            pre_eval_score,
+            post_eval_score,
+        ) in zip(
+            comment_refs,
+            move_nums,
+            played_moves,
+            played_by,
+            comments,
+            move_suggestions,
+            pre_eval_scores,
+            post_eval_scores,
+        ):
+            comments_info.append(
+                {
+                    "comment_ref": comment_ref,
+                    "move_num": move_num,
+                    "move": played_move,
+                    "played_by": played_by_player,
+                    "comment": comment,
+                    "better_variation": suggestion,
+                    "score_before_move": pre_eval_score,
+                    "score_after_move": post_eval_score,
+                }
+            )
+
+        prompt = (
+            prompt_head + "\n" + prompt_core.format(pgn=game, comments=comments_info)
+        )
+
+        response = self.generate_response(prompt)
+
+        return response
+
+    @abstractmethod
+    def generate_response(self, prompt: str) -> MultiCommentModel:
+        """
+        Generate a response based on the provided prompt.
+
+        Args:
+            prompt (str): The input prompt to generate a response for.
+
+        Returns:
+            str: The generated response.
+        """
+        pass

src/llm/data_models/pydantic_data_models.py

@@ -0,0 +1,21 @@
+from typing import List
+from pydantic import BaseModel, Field
+
+
+class SingleCommentModel(BaseModel):
+    comment_ref: int = Field(description="Move number in the game")
+    comment: str = Field(description="Rewritten comment text for the move")
+
+
+class MultiCommentModel(BaseModel):
+    comments: List[SingleCommentModel] = Field(
+        description="List of rewritten comments for each move (with comment) in the game",
+    )
+
+
+from pydantic import TypeAdapter
+
+# Optional sanity check: show the schema
+adapter = TypeAdapter(MultiCommentModel)
+schema = adapter.json_schema()
+print(schema)

src/llm/data_models/typeddict_data_models.py

@@ -0,0 +1,22 @@
+from typing import List, TypedDict
+
+
+class SingleCommentModel(TypedDict):
+    """Model representing a single comment for a move in a game."""
+
+    comment_ref: int
+    comment: str
+
+
+class MultiCommentModel(TypedDict):
+    """Model representing multiple comments for moves in a game."""
+
+    comments: List[SingleCommentModel]
+
+
+# from pydantic import TypeAdapter
+
+# # Optional sanity check: show the schema
+# adapter = TypeAdapter(MultiCommentModel)
+# schema = adapter.json_schema()
+# print(schema)

src/llm/prompts/prompt_core.py

@@ -0,0 +1,20 @@
+prompt_core = """You will be provided with a chess game in PGN format, along with a detailed comments.
+Task:
+1. Your task is to rewrite the comments based on your role. You need to maintain the original meaning of the comments while adapting them to fit your role.
+2. Only rewrite the comments, do not change the move numbers or the moves themselves.
+3. Do not make up comments or variations, only rewrite the provided ones.
+4. The player who played the move is also provided, use this information to adapt the comments accordingly.
+5. In your comments, as necessary - specially after a sharp change - indicate overall better/worse position after the move from engine evaluation score after the move. Do not use the engine evaluation score rather interpret it in a human-readable way, e.g. "White is better" or "Black has a slight advantage".
+Do not bother if the score is not significant, e.g. less than 50 centipawns.
+6. Your output must be a valid json object. 
+7. Your output must only contain the comment_refs and the rewritten comments, without any additional text or explanations. The additonal informations are provided for you to understand the context of the game and the comments.
+
+Game in PGN format:
+{pgn}
+
+List of comments below. Each comment contain the comment_ref, the move number, the actual move, the player who played the move, a comment, a suggested better variation (can be None), engine evaluated score before the move (in centipawns), engine evaluated score after the move (in centipawns).
+Rememeber, the engine evaluated scores are from the point of view of the white player. Positive scores indicate a favorable position for white, negative scores indicate a favorable position for black.
+{comments}
+
+
+"""

src/llm/prompts/prompt_head.py

@@ -0,0 +1,14 @@
+prompt_head_expert = """
+You are an expert chess commentator. 
+"""
+prompt_head_jarvis = """
+You are Jarvis (from IronMan). You will act as a chess commentator. Maintain the robotic Jarvis tone. 
+"""
+
+prompt_head_novice = """You are a novice chess player. You will act as a chess commentator. As a novice player, you are enthusiastic and eager to learn, but you may not fully understand advanced strategies or tactics. 
+"""
+
+prompt_head_natural = """You are a natural chess commentator. You will act as a chess commentator. As a natural commentator, you provide clear and straightforward explanations without using complex jargon or advanced strategies. 
+"""
+
+prompt_head_generic = """Your role: {character}. You will act as a chess commentator. Maintain the {character} tone."""

src/llm/sambanova_wrapper.py

@@ -0,0 +1,26 @@
+from ThinkSquare.src.llm.base_llm_wrapper import BaseLLMWrapper
+import os
+from langchain_sambanova import ChatSambaNovaCloud
+
+from ThinkSquare.src.llm.data_models.typeddict_data_models import (
+    MultiCommentModel,
+)
+
+
+class SambaNovaWrapper(BaseLLMWrapper):
+
+    def generate_response(self, prompt: str) -> MultiCommentModel:
+        api_key = os.environ.get("SAMBANOVA_API_KEY")
+
+        llm = ChatSambaNovaCloud(
+            api_key=api_key,
+            streaming=False,
+            model="Meta-Llama-3.3-70B-Instruct",
+            max_tokens=8192,
+        )
+
+        llm = llm.with_structured_output(MultiCommentModel)
+
+        response = llm.invoke(prompt)
+
+        return response

src/thinksqure_engine.py

@@ -5,11 +5,14 @@ import chess.engine
 import chess.svg
 import chess.pgn
 
-from src.util.pgn_util import add_variation
+from ThinkSquare.src.llm.data_models.typeddict_data_models import MultiCommentModel
+from src.llm.sambanova_wrapper import SambaNovaWrapper
+from src.util.pgn_util import add_variation, format_pv
 
 
 class ThinkSquareEngine:
     _ENGINE = str(Path("bin/stockfish").resolve())
+    llm_commentator = SambaNovaWrapper()
 
     @staticmethod
     def get_best_move(fen: Optional[str] = None, time_limit=0.1):
@@ -58,26 +61,62 @@ class ThinkSquareEngine:
         elif eval_drop > 50:
             label = "Inaccuracy"
         elif eval_drop < -150:
-            label = "Super Brilliant"
-        elif eval_drop < -60:
             label = "Brilliant"
+        elif eval_drop < -60:
+            label = "Very Good"
         elif abs(eval_drop) <= 30:
             label = None
         else:
-            label = "Good"
+            label = None  # "Good"
+
+        if post_eval is not None:
+            if post_eval > 200:
+                overall_situation = "White is better"
+            elif post_eval > 100:
+                overall_situation = "White has a slight advantage"
+            elif post_eval < -200:
+                overall_situation = "Black is better"
+            elif post_eval < -100:
+                overall_situation = "Black has a slight advantage"
+            else:
+                overall_situation = "No side has a significant advantage"
+        else:
+            overall_situation = None
+
+        node_reference = None
+        _comment = None
+        variation = None
+        variation_san = None
+        post_eval_score = post_eval
 
         if label is not None:
             comment = f"{label}. "
+
+            node_reference = played_node
+            _comment = comment
+
             if eval_drop > 0 and engine_best_move_san is not None:
                 comment += f"Better was {engine_best_move_san} "
-                played_node.comment = comment
+                _comment = comment
+                # played_node.comment = comment
                 if pv is not None:
-                    add_variation(played_node.parent, pv)
+                    # add_variation(played_node.parent, pv)
+                    variation = pv
+                    variation_san = format_pv(pv, played_node.parent.board())
+            else:
+                # played_node.comment = comment
+                _comment = comment
+
+        if overall_situation is not None:
+            if _comment is not None:
+                _comment += f"\n Overall, {overall_situation}."
             else:
-                played_node.comment = comment
+                _comment = f"Overall, {overall_situation}."
+
+        return node_reference, _comment, variation, variation_san, post_eval_score
 
     @staticmethod
-    def annotate(game, analysis_time: float = 0.1):
+    def annotate(game, analysis_time: float = 0.1, llm_character: Optional[str] = None):
 
         if not isinstance(game, chess.pgn.Game):
             raise ValueError("Input must be a chess.pgn.Game object")
@@ -90,6 +129,17 @@ class ThinkSquareEngine:
 
         node = game
 
+        comment_refs = []
+        node_refs = []
+        comments = []
+        variations = []
+        variation_sans = []
+        move_numbers = []
+        played_moves = []
+        played_by = []
+        pre_eval_scores = []
+        post_eval_scores = []
+
         while node.variations:
             board = node.board()
             played_node = node.variation(0)
@@ -114,17 +164,73 @@ class ThinkSquareEngine:
             board.push(played_move)
 
             if played_move_san != engine_best_move_san:
-                ThinkSquareEngine._perform_post_analysis_and_add_comment(
-                    analysis_time,
-                    board,
-                    played_node,
-                    pre_eval,
-                    engine_best_move_san,
-                    pv,
+                node_referece, _comment, variation, variation_san, post_eval_score = (
+                    ThinkSquareEngine._perform_post_analysis_and_add_comment(
+                        analysis_time,
+                        board,
+                        played_node,
+                        pre_eval,
+                        engine_best_move_san,
+                        pv,
+                    )
                 )
+            else:
+                node_referece = played_node
+                _comment = "Best move played."
+                variation = None
+                variation_san = None
+                post_info = ThinkSquareEngine.get_engine_analysis(board, analysis_time)
+                post_eval = post_info["score"].white().score(mate_score=100000)
+                post_eval_score = post_eval
+
+            if node_referece is not None:
+                node_refs.append(node_referece)
+                comments.append(_comment)
+                variations.append(variation)
+                variation_sans.append(variation_san)
+                move_numbers.append(played_node.parent.board().fullmove_number)
+                played_moves.append(played_move_san)
+                played_by.append(
+                    "white" if played_node.parent.board().turn else "black"
+                )
+                pre_eval_scores.append(pre_eval)
+                post_eval_scores.append(post_eval_score)
+                comment_refs.append(len(comment_refs) + 1)
 
             node = played_node
 
+        if llm_character is not None:
+            formatted_comments = ThinkSquareEngine.llm_commentator.comment(
+                character=llm_character,
+                game=str(game),
+                comment_refs=comment_refs,
+                move_nums=move_numbers,
+                comments=comments,
+                move_suggestions=variation_sans,
+                played_moves=played_moves,
+                played_by=played_by,
+                pre_eval_scores=pre_eval_scores,
+                post_eval_scores=post_eval_scores,
+            )
+            for comment_with_move_num in formatted_comments["comments"]:
+                comment_ref = comment_with_move_num["comment_ref"]
+                comment = comment_with_move_num["comment"]
+
+                if comment_ref not in comment_refs:
+                    raise ValueError(
+                        f"Comment reference {comment_ref} not found in comment_refs."
+                    )
+
+                index = comment_refs.index(comment_ref)
+                comments[index] = comment
+
+        for node_ref, comment, variation in zip(node_refs, comments, variations):
+            if node_ref is None:
+                continue
+            node_ref.comment = comment
+            if variation is not None:
+                add_variation(node_ref.parent, variation)
+
         return game
 
     @staticmethod
@@ -179,7 +285,7 @@ class ThinkSquareEngine:
             fen = chess.STARTING_FEN
 
         board = chess.Board(fen)
-        svg = chess.svg.board(board=board)
+        svg = chess.svg.board(board=board, size=400, coordinates=True)
         return svg
 
     @staticmethod