repository/testing_repo.py

from typing import Tuple


from repository.repository_abc import Repository, Model, ModelRoles


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