dam_helper/repository/testing_repo.py

from repository import ModelRoles, Model, Repository


class TestingRepository(Repository):

    def __init__(self, prompts_answers: list[dict[str, str]], model_info:Model=None):
        self.prompt_answers = prompts_answers
        self.next_answer = iter(self.prompt_answers)
        self.message_history = []
        self.model_info = model_info or Model("fake_model",
                                              ModelRoles("system", "user", "assistant"))
    def init(self):
        pass

    def send_prompt(self, prompt: str, add_to_history: bool = True) -> dict[str, str]:
        response = next(self.next_answer)
        if add_to_history:
            self.get_message_history().append(response)
        return response

    def get_message_history(self) -> list[dict[str, str]]:
        return self.message_history

    def get_model_info(self) -> Model:
        return self.model_info

    def get_model_roles(self) -> ModelRoles:
        return self.model_info.roles