fix : server crash with ollama

server/app.py

@@ -3,8 +3,8 @@
 import argparse
 from flask import Flask, jsonify, request
 
-from sources.llamacpp import LlamacppLLM
-from sources.ollama import OllamaLLM
+from sources.llamacpp_handler import LlamacppLLM
+from sources.ollama_handler import OllamaLLM
 
 parser = argparse.ArgumentParser(description='AgenticSeek server script')
 parser.add_argument('--provider', type=str, help='LLM backend library to use. set to [ollama] or [llamacpp]', required=True)

server/sources/ollama_handler.py

@@ -1,4 +1,5 @@
 
+import time
 from .generator import GeneratorLLM
 import ollama
 
@@ -31,17 +32,28 @@ class OllamaLLM(GeneratorLLM):
 
                 with self.state.lock:
                     self.state.current_buffer += content
-
-        except ollama.ResponseError as e:
-            if e.status_code == 404:
+        except Exception as e:
+            if "404" in str(e):
                 self.logger.info(f"Downloading {self.model}...")
                 ollama.pull(self.model)
-            with self.state.lock:
-                self.state.is_generating = False
-            print(f"Error: {e}")
-        except Exception as e:
             if "refused" in str(e).lower():
                 raise Exception("Ollama connection failed. is the server running ?") from e
+            raise e
         finally:
+            self.logger.info("Generation complete")
             with self.state.lock:
-                self.state.is_generating = False
+                self.state.is_generating = False
+
+if __name__ == "__main__":
+    generator = OllamaLLM()
+    history = [
+        {
+            "role": "user",
+            "content": "Hello, how are you ?"
+        }
+    ]
+    generator.set_model("deepseek-r1:1.5b")
+    generator.start(history)
+    while True:
+        print(generator.get_status())
+        time.sleep(1)