Improve logging logic and change license

Author: kenliao94

pyproject.toml

@@ -7,7 +7,7 @@ requires-python = ">=3.10"
 authors = [{ name = "Ken Liao" }]
 maintainers = [{ name = "Ken Liao", email = "[email protected]" }]
 keywords = ["rabbitmq", "mcp", "llm", "automation"]
-license = { text = "MIT" }
+license = { text = "Apache-2.0" }
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",

src/mcp_server_rabbitmq/logger.py

@@ -1,11 +1,43 @@
 import datetime
+from enum import IntEnum
 
-# Replace this with more powerful logging library
+class LOG_LEVEL(IntEnum):
+    DEBUG = 0
+    INFO = 1
+    WARNING = 2
+    ERROR = 3
+
+# TODO: see if there are better libraries available
 class Logger:
-    def __init__(self, log_file_path):
+    def __init__(self, log_file_path: str, log_level: LOG_LEVEL = LOG_LEVEL.WARNING):
         self.log_file_path = log_file_path
+        self.log_level = log_level
 
-    def log(self, msg):
+    def debug(self, msg):
+        if not self.__should_skip_logging(LOG_LEVEL.DEBUG):
+            self.__write_log(self.__post_process_log_entry(msg, LOG_LEVEL.DEBUG))
+
+    def info(self, msg):
+        if not self.__should_skip_logging(LOG_LEVEL.INFO):
+            self.__write_log(self.__post_process_log_entry(msg, LOG_LEVEL.INFO))
+
+    def warning(self, msg):
+        if not self.__should_skip_logging(LOG_LEVEL.WARNING):
+            self.__write_log(self.__post_process_log_entry(msg, LOG_LEVEL.WARNING))
+
+    def error(self, msg):
+        if not self.__should_skip_logging(LOG_LEVEL.ERROR):
+            self.__write_log(self.__post_process_log_entry(msg, LOG_LEVEL.ERROR))
+
+    def __post_process_log_entry(self, msg: str, log_level: LOG_LEVEL) -> str:
         timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        return f"[{log_level}] {timestamp} - {msg}\n"
+
+    def __write_log(self, msg):
         with open(self.log_file_path, "a") as f:
-            f.write(f"{timestamp} - {msg}\n")
+            f.write(msg)
+            # TODO: look for optimization opportunity here
+            f.flush()
+    
+    def __should_skip_logging(self, log_level: LOG_LEVEL) -> bool:
+        return int(log_level) < int(self.log_level)

src/mcp_server_rabbitmq/server.py

@@ -7,10 +7,10 @@
 import pika
 import ssl
 from .models import Enqueue
-from .logger import Logger
+from .logger import Logger, LOG_LEVEL
 
 
-async def serve(rabbitmq_host: str, port: int, username: str, password: str, use_tls: bool) -> None:
+async def serve(rabbitmq_host: str, port: int, username: str, password: str, use_tls: bool, logging_level: LOG_LEVEL = LOG_LEVEL.WARNING) -> None:
     # Setup server and logger
     server = Server("mcp-rabbitmq")
     # Make logger another option to pass
@@ -50,7 +50,7 @@ async def call_tool(
                 channel.basic_publish(exchange="", routing_key=queue, body=message)
                 return [TextContent(type="text", text=str("suceeded"))]
             except Exception as e:
-                logger.log(f"[ERROR] {e}")
+                logger.error(f"[ERROR] {e}")
                 return [TextContent(type="text", text=str("failed"))]
         raise ValueError(f"Tool not found: {name}")