from ollama import Client
import os
import env_manager
env_manager.set_env()

class LLM:
    def __init__(self, system_message=None, num_ctx=2000, temperature=0, chat=True) -> None:
        self.llm_model = os.getenv("LLM_MODEL")
        self.system_message = system_message
        self.options = {"temperature": temperature, "num_ctx": num_ctx}
        self.messages = [{'role': 'system', 'content': self.system_message}]
        self.chat = chat
        self.ollama = Client(host=f'{os.getenv("LLM_URL")}:{os.getenv("LLM_PORT")}')

    def generate(self, prompt: str) -> str:
        self.messages.append({"role": "user", "content": prompt})

        result = self.ollama.chat(model=self.llm_model, messages=self.messages, options=self.options, )

        answer = result['message']['content']
        self.messages.append({"role": "assistant", "content": answer})
        if not self.chat:
            self.messages = [{'role': 'system', 'content': self.system_message}]

        return answer