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

Canner · Jan 8, 2025 · 7fbf3ad · 7fbf3ad
1 parent 0892a41
commit 7fbf3ad
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 27 deletions.
diff --git a/wren-ai-service/eval/llm_trace_analysis/report_gen.py b/wren-ai-service/eval/llm_trace_analysis/report_gen.py
@@ -153,7 +153,7 @@ def save_report(report_by_release: dict):
 
 
 async def main():
-    RELEASES = ["0.15.3", "0.14.2"]
+    RELEASES = ["0.15.4", "0.15.3", "0.14.2"]
     TRACE_NAMES = [
         # indexing
         "Prepare Semantics",

diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py
@@ -29,7 +29,7 @@ async def lifespan(app: FastAPI):
     pipe_components = generate_components(settings.components)
     app.state.service_container = create_service_container(pipe_components, settings)
     app.state.service_metadata = create_service_metadata(pipe_components)
-    init_langfuse()
+    init_langfuse(settings)
 
     yield
 

diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py
@@ -41,7 +41,7 @@ def dry_run_pipeline(pipeline_cls: BasicPipeline, pipeline_name: str, **kwargs):
 
     pipe_components = generate_components(settings.components)
     pipeline = pipeline_cls(**pipe_components[pipeline_name])
-    init_langfuse()
+    init_langfuse(settings)
 
     async_validate(lambda: pipeline.run(**kwargs))
 

diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py
@@ -6,31 +6,37 @@
 from dotenv import load_dotenv
 from langfuse.decorators import langfuse_context
 
+from src.config import Settings
+
 logger = logging.getLogger("wren-ai-service")
 
 
 class CustomFormatter(logging.Formatter):
-    grey = "\x1b[38;20m"
-    yellow = "\x1b[33;20m"
-    red = "\x1b[31;20m"
-    bold_red = "\x1b[31;1m"
-    reset = "\x1b[0m"
-    format = (
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
-    )
-
-    FORMATS = {
-        logging.DEBUG: yellow + format + reset,
-        logging.INFO: grey + format + reset,
-        logging.WARNING: yellow + format + reset,
-        logging.ERROR: red + format + reset,
-        logging.CRITICAL: bold_red + format + reset,
-    }
-
-    def format(self, record):
-        log_fmt = self.FORMATS.get(record.levelno)
-        formatter = logging.Formatter(log_fmt)
-        return formatter.format(record)
+    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 ImportError:
+        pass
+    finally:
+        _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)
 
 
 def setup_custom_logger(name, level_str: str):
@@ -63,9 +69,7 @@ def remove_trailing_slash(endpoint: str) -> str:
     return endpoint.rstrip("/") if endpoint.endswith("/") else endpoint
 
 
-def init_langfuse():
-    from src.config import settings
-
+def init_langfuse(settings: Settings):
     langfuse_context.configure(
         enabled=settings.langfuse_enable,
         host=settings.langfuse_host,