Spaces:
Running
Running
File size: 1,023 Bytes
224e4de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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
|