From 7e95498f7ac10703b0c9f07fa21341050fd97d8b Mon Sep 17 00:00:00 2001
From: martin legrand <martin.legrand@epitech.eu>
Date: Sat, 29 Mar 2025 13:14:37 +0100
Subject: [PATCH] feat : better logging

---
 server/app.py           | 11 ++---------
 sources/llm_provider.py |  6 +++---
 2 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/server/app.py b/server/app.py
index 242c576..2e7a1d9 100644
--- a/server/app.py
+++ b/server/app.py
@@ -13,10 +13,10 @@ args = parser.parse_args()
 
 app = Flask(__name__)
 
-generator = None
-
 assert args.provider in ["ollama", "llamacpp"], f"Provider {args.provider} does not exists. see --help for more information"
 
+generator = OllamaLLM if args.provider == "ollama" else LlamacppLLM 
+
 @app.route('/generate', methods=['POST'])
 def start_generation():
     if generator is None:
@@ -32,13 +32,6 @@ def setup():
     global generator
     data = request.get_json()
     model = data.get('model', None)
-    provider = args.provider
-    if provider == "ollama":
-        generator = OllamaLLM()
-    elif provider == "llamacpp":
-        generator = LlamacppLLM()
-    else:
-        raise ValueError(f"Provider {provider} does not exists. see --help for more information")
     if model is None:
         return jsonify({"error": "Model not provided"}), 400
     generator.set_model(model)
diff --git a/sources/llm_provider.py b/sources/llm_provider.py
index 51657ba..4dc3649 100644
--- a/sources/llm_provider.py
+++ b/sources/llm_provider.py
@@ -120,14 +120,14 @@ class Provider:
             is_complete = False
             while not is_complete:
                 response = requests.get(f"http://{self.server_ip}/get_updated_sentence")
-                print("raw response", response.json())
                 if "error" in response.json():
-                    continue
+                    pretty_print(response.json()["error"], color="failure")
+                    break
                 thought = response.json()["sentence"]
                 is_complete = bool(response.json()["is_complete"])
                 time.sleep(2)
         except KeyError as e:
-            raise Exception(f"{str(e)}\n\nError occured with server route. Are you using the correct address for the config.ini provider?") from e
+            raise Exception(f"{str(e)}\nError occured with server route. Are you using the correct address for the config.ini provider?") from e
         except Exception as e:
             raise e
         return thought