from langchain_ollama import OllamaEmbeddings, OllamaLLM


class OllamaWrapper:
    def __init__(
        self,
        ollama_url="http://127.0.0.1:11434",
        embedding_model="nomic-embed-text",
        llm_model="llama3.2",
    ):
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        self._embeddings = None
        self._llm = None
        self._ollama_url = ollama_url
        self._chat = []

    @property
    def embeddings(self):
        if self._embeddings is None:
            self._embeddings = OllamaEmbeddings(
                model=self.embedding_model, base_url=self._ollama_url
            )
        return self._embeddings

    @property
    def llm(self):
        if self._llm is None:
            self._llm = OllamaLLM(model=self.llm_model, base_url=self._ollama_url)
        return self._llm

    def update_models(self, embedding_model=None, llm_model=None):
        """Update model names and reset clients"""
        if embedding_model:
            self.embedding_model = embedding_model
            self._embeddings = None
        if llm_model:
            self.llm_model = llm_model
            self._llm = None
