src/llm/base_llm_wrapper.py

from abc import ABC, abstractmethod
from typing import List
from 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