fix test_init_custom_logger_compatible_class_as_callback

litellm/integrations/cloudzero/cloudzero.py

@@ -6,10 +6,6 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 
-from .cz_stream_api import CloudZeroStreamer
-from .database import LiteLLMDatabase
-from .transform import CBFTransformer
-
 if TYPE_CHECKING:
     from apscheduler.schedulers.asyncio import AsyncIOScheduler
 else:
@@ -105,6 +101,9 @@ class CloudZeroLogger(CustomLogger):
             limit: Optional limit on number of records to export
             operation: CloudZero operation type ("replace_hourly" or "sum")
         """
+        from litellm.integrations.cloudzero.cz_stream_api import CloudZeroStreamer
+        from litellm.integrations.cloudzero.database import LiteLLMDatabase
+        from litellm.integrations.cloudzero.transform import CBFTransformer
         try:
             verbose_logger.debug("CloudZero Logger: Starting usage data export")
             
@@ -163,6 +162,8 @@ class CloudZeroLogger(CustomLogger):
         Returns:
             dict: Contains usage_data, cbf_data, and summary statistics
         """
+        from litellm.integrations.cloudzero.database import LiteLLMDatabase
+        from litellm.integrations.cloudzero.transform import CBFTransformer
         try:
             verbose_logger.debug("CloudZero Logger: Starting dry run export")
             

litellm/litellm_core_utils/custom_logger_registry.py

@@ -38,6 +38,7 @@ try:
     from litellm_enterprise.integrations.prometheus import PrometheusLogger
 except Exception:
     PrometheusLogger = None
+from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
 from litellm.integrations.dotprompt import DotpromptManager
 from litellm.integrations.s3_v2 import S3Logger
 from litellm.integrations.sqs import SQSLogger
@@ -86,6 +87,7 @@ class CustomLoggerRegistry:
         "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler,
         "vector_store_pre_call_hook": VectorStorePreCallHook,
         "dotprompt": DotpromptManager,
+        "cloudzero": CloudZeroLogger,
     }
 
     try: