From 859ff12e0c4e596336ff8f276d407c4aa315d4e9 Mon Sep 17 00:00:00 2001
From: Jamie Broomall <jamie@whylabs.ai>
Date: Wed, 6 Nov 2024 07:30:48 -0800
Subject: [PATCH] Fix rolling logger schema

---
 langkit/whylogs/rolling_logger.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/langkit/whylogs/rolling_logger.py b/langkit/whylogs/rolling_logger.py
index 97ab24c6..7808eb71 100644
--- 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)