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)