feat(ollama): Improved error management for Ollama. Also pull model from Ollama libary if needed.

Author: geoffroy-noel-ddh

describer/base.py

@@ -75,6 +75,13 @@ def log(self, message):
         timer = getattr(self, 'timer')
         if timer:
             timer.step(message)
+    
+    def log_fatal(self, message):
+        message = f'ERROR: {message}'
+        self.log(message)
+        
+        print(message)
+        exit(1)
 
     def get_model(self):
         if not self.model:

describer/ollama.py

@@ -45,18 +45,27 @@ def _get_ollama_model_code(self):
     def answer_question(self, image_path, question):
         if not self.model:
             self._init_model()
+        
+        import ollama
 
-        response = self.model.chat(
-            model=self._get_ollama_model_code(),
-            messages=[{
-                'role': 'user',
-                'content': question,
-                'images': [image_path]
-            }],
-            options={
-                'temperature': 0
-            }
-        )
+        model_code = self._get_ollama_model_code()
+
+        try:
+            response = self.model.chat(
+                model=model_code,
+                messages=[{
+                    'role': 'user',
+                    'content': question,
+                    'images': [image_path]
+                }],
+                options={
+                    'temperature': 0
+                }
+            )
+        except ollama._types.ResponseError as e:
+            if 'this model is missing data required for image input' in str(e):
+                self.log_fatal(f'This ollama model ("{model_code}") does not support image inputs. Please check model doc on ollama.com')
+            raise(e)
 
         # example of a response
         '''
@@ -81,13 +90,59 @@ def answer_question(self, image_path, question):
         return response['message']['content']
 
     def _init_model(self):
-        from ollama import Client
-        self.model = Client(host=BVQA_OLLAMA_HOST)
+        from ollama import Client, _types
+        self.client = Client(host=BVQA_OLLAMA_HOST)
+        self.model = self.client
+
+        self.stop_if_ollama_not_running()
+
+        self.pull_ollama_model()
+
+
         # load the model into memory so get_compute_info get stats.
         # and we separate the load from the actual inference in the logs.
-        res = self.model.generate(model=self._get_ollama_model_code(), prompt='Just say "yes". Nothing else.')
+        try:
+            res = self.model.generate(model=self._get_ollama_model_code(), prompt='Just say "yes". Nothing else.')
+        except _types.ResponseError as e:
+            print('H0'*40)
+            raise e
+
         return self.model
 
+    def stop_if_ollama_not_running(self):
+        self.get_ollama_models()
+
+    def is_ollama_model_pulled(self):
+        ret = False
+        model_code = self._get_ollama_model_code()
+        models = self.get_ollama_models()
+        matching_models = [
+            model
+            for model
+            in models
+            if model.model == model_code
+        ]
+        return bool(matching_models)
+
+    def pull_ollama_model(self):
+        if not self.is_ollama_model_pulled():
+            import ollama
+            model_code = self._get_ollama_model_code()
+            self.log(f'Pulling model "{model_code}" from Ollama repository...')
+            try:
+                res = self.client.pull(model_code)
+            except ollama._types.ResponseError as e:
+                self.log_fatal(f'failed to pull model "{model_code}" from Ollama repository. {e}')
+
+    def get_ollama_models(self):
+        ret = []
+        try:
+            res = self.client.list()
+            ret = res.models
+        except ConnectionError as e:
+            self.log_fatal(f'{e} ; BVQA_OLLAMA_HOST="{BVQA_OLLAMA_HOST}"')
+        return ret
+
     def get_compute_info(self):
         ret = {
             'type': 'ollama',
@@ -96,6 +151,13 @@ def get_compute_info(self):
         }
         model = self.get_model()
 
+        '''
+        # models=[Model(model='gemma3:4b', name='gemma3:4b', digest='a2af6cc3eb7fa8be8504abaf9b04e88f17a119ec3f04a3addf55f92841195f5a', expires_at=datetime.datetime(2025, 7, 18, 14, 3, 39, 405547, tzinfo=TzInfo(+01:00)), size=6169209844, size_vram=3700725812, details=ModelDetails(parent_model='', format='gguf', family='gemma3', families=['gemma3'], parameter_size='4.3B', quantization_level='Q4_K_M'))]
+        res = self.client.ps()
+        print(res)
+        exit(1)
+        '''
+
         if model:
             res = model.ps()
             for m in res.models:

requirements/ollama.txt

@@ -1 +1 @@
-ollama
+ollama<0.6