Download Ollama model if not present (#1397)

embedchain/llm/ollama.py

@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Iterable
 from typing import Optional, Union
 
@@ -5,11 +6,14 @@ from langchain.callbacks.manager import CallbackManager
 from langchain.callbacks.stdout import StdOutCallbackHandler
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain_community.llms.ollama import Ollama
+from ollama import Client
 
 from embedchain.config import BaseLlmConfig
 from embedchain.helpers.json_serializable import register_deserializable
 from embedchain.llm.base import BaseLlm
 
+logger = logging.getLogger(__name__)
+
 
 @register_deserializable
 class OllamaLlm(BaseLlm):
@@ -18,6 +22,12 @@ class OllamaLlm(BaseLlm):
         if self.config.model is None:
             self.config.model = "llama2"
 
+        client = Client(host=config.base_url)
+        local_models = client.list()["models"]
+        if not any(model.get("name") == self.config.model for model in local_models):
+            logger.info(f"Pulling {self.config.model} from Ollama!")
+            client.pull(self.config.model)
+
     def get_llm_model_answer(self, prompt):
         return self._get_answer(prompt=prompt, config=self.config)