from ollama import Client
import os
import env_manager

env_manager.set_env()


class LLM:
    def __init__(
        self, system_message=None, num_ctx=20000, temperature=0, chat=True
    ) -> None:
        self.llm_model = "mistral-nemo:12b-instruct-2407-q5_K_M" #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

if __name__ == "__main__":
    llm = LLM()
    print(llm.generate("Why is the sky red?"))