chore(wren-ai-service): minor update (#1095)

Canner · Jan 9, 2025 · 89333e4 · 89333e4
1 parent ded2387
commit 89333e4
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 25 deletions.
diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py
@@ -19,7 +19,9 @@
 )
 from src.web.v1 import routers
 
-setup_custom_logger("wren-ai-service", level_str=settings.logging_level)
+setup_custom_logger(
+    "wren-ai-service", level_str=settings.logging_level, is_dev=settings.development
+)
 
 
 # https://fastapi.tiangolo.com/advanced/events/#lifespan

diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py
@@ -12,34 +12,38 @@
 
 
 class CustomFormatter(logging.Formatter):
-    try:
-        # Imports the Cloud Logging client library
-        import google.cloud.logging
-
-        # Instantiates a client
-        client = google.cloud.logging.Client()
-
-        # Retrieves a Cloud Logging handler based on the environment
-        # you're running in and integrates the handler with the
-        # Python logging module. By default this captures all logs
-        # at INFO level and higher
-        client.setup_logging()
-    except Exception:
-        pass
-    finally:
+    def __init__(self, is_dev: bool):
+        super().__init__()
+
+        try:
+            if not is_dev:
+                # Imports the Cloud Logging client library
+                import google.cloud.logging
+
+                # Instantiates a client
+                client = google.cloud.logging.Client()
+
+                # Retrieves a Cloud Logging handler based on the environment
+                # you're running in and integrates the handler with the
+                # Python logging module. By default this captures all logs
+                # at INFO level and higher
+                client.setup_logging()
+        except Exception:
+            pass
+
+    def format(self, record):
         _LOGGING_FORMAT = "{levelname:<.1}{asctime}.{msecs:03.0f} {process} {name}:{lineno}] {message}"
         _DATE_FMT = "%m%d %H:%M:%S"
 
-        def format(self, record):
-            formatter = logging.Formatter(
-                fmt=self._LOGGING_FORMAT,
-                datefmt=self._DATE_FMT,
-                style="{",
-            )
-            return formatter.format(record)
+        formatter = logging.Formatter(
+            fmt=_LOGGING_FORMAT,
+            datefmt=_DATE_FMT,
+            style="{",
+        )
+        return formatter.format(record)
 
 
-def setup_custom_logger(name, level_str: str):
+def setup_custom_logger(name, level_str: str, is_dev: bool):
     level_str = level_str.upper()
 
     if level_str not in logging._nameToLevel:
@@ -48,7 +52,7 @@ def setup_custom_logger(name, level_str: str):
     level = logging._nameToLevel[level_str]
 
     handler = logging.StreamHandler()
-    handler.setFormatter(CustomFormatter())
+    handler.setFormatter(CustomFormatter(is_dev))
 
     logger = logging.getLogger(name)
     logger.setLevel(level)