from langchain_community.chat_models import ChatOllama

class OllamaChatBot:
    def __init__(self, base_url, model, headers):
        self.base_url = base_url
        self.model = model
        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 get_request(self, prompt):
        messanges = []
        self.messanges.append(prompt)
        if len(self.messanges) > 5:
            messanges = messanges[:5]
        else:
            messanges = self.messanges
        return self.ollama.invoke(messanges).content