Spaces:
Sleeping
Sleeping
File size: 908 Bytes
d015f0c 0fadcb9 d015f0c 0fadcb9 d015f0c 0fadcb9 7b864ba 0fadcb9 d015f0c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import ollama
from ollama import Options
from llm.llm import ModelRoles, Model
class OllamaRepository:
def __init__(self, model:Model, system_msg):
self.model: Model = model
self.system_msg: str = system_msg
self.message_history: list[dict[str, str]] = [{"role": self.model.roles.system_role, "content": system_msg}]
def send_prompt(self, prompt:str, add_to_history:bool = True) -> dict[str, str]:
options: Options = Options(temperature=0)
self.message_history.append({"role": self.model.roles.user_role, "content":prompt})
response = ollama.chat(self.model.name, self.message_history, options=options)
answer = {"role": self.model.roles.ai_role, "content": response["message"]["content"]}
if add_to_history:
self.message_history.append(answer)
else:
self.message_history.pop()
return answer
|