diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index 727dabc094..ab4ec234bf 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/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") diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd82ecdf2b..af51fe9ab7 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/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: