feat : better logging

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)

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