feat : better server provider

server/sources/generator.py

@@ -0,0 +1,55 @@
+
+from flask import jsonify
+import threading
+import logging
+from abc import abstractmethod
+
+class GenerationState:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.last_complete_sentence = ""
+        self.current_buffer = ""
+        self.is_generating = False
+    
+    def status(self) -> dict:
+        return {
+            "sentence": self.current_buffer,
+            "is_complete": not self.is_generating,
+            "last_complete_sentence": self.last_complete_sentence,
+            "is_generating": self.is_generating,
+        }
+
+class GeneratorLLM():
+    def __init__(self):
+        self.model = None
+        self.state = GenerationState()
+        self.logger = logging.getLogger(__name__)
+    
+    def set_model(self, model: str) -> None:
+        self.logger.info(f"Model set to {model}")
+        self.model = model
+    
+    def start(self, history: list) -> bool:
+        if self.model is None:
+            raise Exception("Model not set")
+        with self.state.lock:
+            if self.state.is_generating:
+                return False
+            self.logger.info("Starting generation")
+            threading.Thread(target=self.generate, args=(history,)).start()
+        return True
+    
+    def get_status(self) -> dict:
+        with self.state.lock:
+            return jsonify(self.state.status())
+
+    @abstractmethod
+    def generate(self, history: list) -> None:
+        """
+        Generate text using the model.
+        args:
+            history: list of strings
+        returns:
+            None
+        """
+        pass

server/sources/llamacpp.py

@@ -0,0 +1,22 @@
+
+from .generator import GeneratorLLM
+
+class LlamacppLLM(GeneratorLLM):
+    from llama_cpp import Llama
+
+    def __init__(self):
+        """
+        Handle generation using llama.cpp
+        """
+        super().__init__()
+        self.llm = Llama.from_pretrained(
+            repo_id=self.model,
+            filename="*q8_0.gguf",
+            verbose=True
+        )
+    
+    def generate(self, history):
+        self.logger.info(f"Using {self.model} for generation with Llama.cpp")
+        self.llm.create_chat_completion(
+              messages = history
+        )

server/sources/ollama.py

@@ -0,0 +1,46 @@
+
+from .generator import GeneratorLLM
+
+class OllamaLLM(GeneratorLLM):
+    import ollama
+
+    def __init__(self):
+        """
+        Handle generation using Ollama.
+        """
+        super().__init__()
+
+    def generate(self, history):
+        self.logger.info(f"Using {self.model} for generation with Ollama")
+        try:
+            with self.state.lock:
+                self.state.is_generating = True
+                self.state.last_complete_sentence = ""
+                self.state.current_buffer = ""
+
+            stream = ollama.chat(
+                model=self.model,
+                messages=history,
+                stream=True,
+            )
+
+            for chunk in stream:
+                content = chunk['message']['content']
+                print(content, end='', flush=True)
+
+                with self.state.lock:
+                    self.state.current_buffer += content
+
+        except ollama.ResponseError as e:
+            if e.status_code == 404:
+                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
+        finally:
+            with self.state.lock:
+                self.state.is_generating = False