Fix rolling logger schema (#318)

whylabs · Nov 6, 2024 · ae9fc6d · ae9fc6d
1 parent ebe21c4
commit ae9fc6d
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/langkit/whylogs/rolling_logger.py b/langkit/whylogs/rolling_logger.py
@@ -1,7 +1,6 @@
 from typing import Any, Optional
 import whylogs as why
 from whylogs.core.schema import DatasetSchema
-from langkit import llm_metrics
 
 
 class RollingLogger:
@@ -12,15 +11,22 @@ def __init__(
         schema: Optional[DatasetSchema] = None,
         **kwargs: Any
     ):
-        llm_schema = llm_metrics.init()
+        if schema is None:
+            from langkit import llm_metrics
+
+            schema = llm_metrics.init()
         self.logger = why.logger(
             mode="rolling",
             interval=interval_minutes,
             when="M",
             base_name="langkit",
-            schema=llm_schema,
+            schema=schema,
         )
-        self.logger.append_writer(name="whylabs", **kwargs)
+
+        if "writer" in kwargs:
+            self.logger.append_writer(name=None, **kwargs)
+        else:
+            self.logger.append_writer(name="whylabs", **kwargs)
 
     def log(self, dict):
         self.logger.log(dict)