from langchain_community.chat_models import ChatOllama

class OllamaChatBot:
    def __init__(self, base_url, model, headers):
        self.base_url = base_url
        self.model = model
        # check if header exists
        if self.is_empty(headers):
            self.headers = ""
        else: 
            self.headers = headers
        self.messanges = []

        if headers is None:
            self.ollama = ChatOllama(
            base_url=self.base_url,
            model=self.model,
        )
        else:
            self.ollama = ChatOllama(
                base_url=self.base_url,
                model=self.model,
                headers = self.headers
            )

    def is_empty(self, dictionary: dict) -> bool:
        return len(dictionary) == 1 and list(dictionary.keys())[0] == '' and list(dictionary.values())[0] == ''

        
    def get_request(self, prompt: str) -> str:
        messanges = []
        self.messanges.append(prompt)
        if len(self.messanges) > 5:
            messanges = messanges[:5]
        else:
            messanges = self.messanges
        try:
            return self.ollama.invoke(messanges).content
        except ValueError as e:
            return f"An unexpected Error occuried: {e}"