From c89b635bda96105acb1e717272976ab9053250e7 Mon Sep 17 00:00:00 2001
From: Jovan Wong <77403495+jovanwongzixi@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:56:30 +0800
Subject: [PATCH] Add cli log flag

Added --log flag in cli arguments to set log level
---
 optillm.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/optillm.py b/optillm.py
index de767acd..891ced91 100644
--- a/optillm.py
+++ b/optillm.py
@@ -32,6 +32,14 @@
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
+logging_levels = {
+    "notset": logging.NOTSET,
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
 
 # Initialize Flask app
 app = Flask(__name__)
@@ -78,6 +86,7 @@
     'optillm_api_key': '',
     'return_full_response': False,
     'port': 8000,
+    'log': 'info',
 }
 
 # List of known approaches
@@ -370,6 +379,7 @@ def parse_args():
         ("--n", "OPTILLM_N", int, 1, "Number of final responses to be returned"),
         ("--return-full-response", "OPTILLM_RETURN_FULL_RESPONSE", bool, False, "Return the full response including the CoT with <thinking> tags"),
         ("--port", "OPTILLM_PORT", int, 8000, "Specify the port to run the proxy"),
+        ("--log", "OPTILLM_LOG", str, "info", "Specify the logging level", list(logging_levels.keys()))
     ]
 
     for arg, env, type_, default, help_text, *extra in args_env:
@@ -415,6 +425,12 @@ def main():
     server_config.update(vars(args))
 
     port = server_config['port']
+
+    # Set logging level from user request
+    logging_level = server_config['log']
+    if logging_level in logging_levels.keys():
+        logger.setLevel(logging_levels[logging_level])
+    
     logger.info(f"Starting server with approach: {server_config['approach']}")
     server_config_clean = server_config.copy()
     if server_config_clean['optillm_api_key']: