diff --git a/server/app.py b/server/app.py index b939d2a..c22ea4b 100644 --- a/server/app.py +++ b/server/app.py @@ -20,29 +20,28 @@ generator = OllamaLLM() if args.provider == "ollama" else LlamacppLLM() @app.route('/generate', methods=['POST']) def start_generation(): if generator is None: - return jsonify({"error": "Generator not initialized"}), 400 + return jsonify({"error": "Generator not initialized"}), 401 data = request.get_json() history = data.get('messages', []) if generator.start(history): return jsonify({"message": "Generation started"}), 202 - return jsonify({"error": "Generation already in progress"}), 400 + return jsonify({"error": "Generation already in progress"}), 402 @app.route('/setup', methods=['POST']) def setup(): data = request.get_json() model = data.get('model', None) if model is None: - return jsonify({"error": "Model not provided"}), 400 + return jsonify({"error": "Model not provided"}), 403 generator.set_model(model) return jsonify({"message": "Model set"}), 200 @app.route('/get_complete_sentence', methods=['GET']) def get_complete_sentence(): if not generator: - return jsonify({"error": "Generator not initialized"}), 400 + return jsonify({"error": "Generator not initialized"}), 404 while True: status = generator.get_status() - print(status) if status["is_complete"]: return jsonify(status) return None @@ -50,7 +49,7 @@ def get_complete_sentence(): @app.route('/get_updated_sentence') def get_updated_sentence(): if not generator: - return jsonify({"error": "Generator not initialized"}), 400 + return jsonify({"error": "Generator not initialized"}), 405 return generator.get_status() if __name__ == '__main__': diff --git a/server/sources/generator.py b/server/sources/generator.py index 8c0197f..133d6e9 100644 --- a/server/sources/generator.py +++ b/server/sources/generator.py @@ -40,6 +40,7 @@ class GeneratorLLM(): with self.state.lock: if self.state.is_generating: return False + self.state.is_generating = True self.logger.info("Starting generation") threading.Thread(target=self.generate, args=(history,)).start() return True diff --git a/sources/llm_provider.py b/sources/llm_provider.py index b63d870..67e9ce6 100644 --- a/sources/llm_provider.py +++ b/sources/llm_provider.py @@ -79,6 +79,8 @@ class Provider: except AttributeError as e: raise NotImplementedError(f"{str(e)}\nIs {self.provider_name} implemented ?") except Exception as e: + if "RemoteDisconnected" in str(e): + return f"{self.server_ip} seem offline. RemoteDisconnected error." raise Exception(f"Provider {self.provider_name} failed: {str(e)}") from e return thought diff --git a/sources/utility.py b/sources/utility.py index 8639488..27eebb4 100644 --- a/sources/utility.py +++ b/sources/utility.py @@ -6,8 +6,7 @@ import threading import itertools import time - -def pretty_print(text, color = "info"): +def pretty_print(text, color="info"): """ Print text with color formatting. @@ -37,7 +36,8 @@ def pretty_print(text, color = "info"): print(text) pretty_print(f"Invalid color {color} in pretty_print", "warning") return - print(color_map[color], text, Fore.RESET) + print(f"{color_map[color]}{text}{Fore.RESET}") + print(' ' * 10) # prevent cut-off else: color_map = { "success": "green", @@ -51,6 +51,7 @@ def pretty_print(text, color = "info"): if color not in color_map: color = "default" print(colored(text, color_map[color])) + print(' ' * 10) # prevent cut-off def animate_thinking(text, color="status", duration=2): def _animate():