diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index ca15962b72..403829deba 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, cast import litellm from litellm._logging import verbose_logger +from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger if TYPE_CHECKING: @@ -15,22 +16,30 @@ else: class CloudZeroLogger(CustomLogger): """ CloudZero Logger for exporting LiteLLM usage data to CloudZero AnyCost API. - + Environment Variables: CLOUDZERO_API_KEY: CloudZero API key for authentication CLOUDZERO_CONNECTION_ID: CloudZero connection ID for data submission CLOUDZERO_TIMEZONE: Timezone for date handling (default: UTC) """ - def __init__(self, api_key: Optional[str] = None, connection_id: Optional[str] = None, timezone: Optional[str] = None, **kwargs): + def __init__( + self, + api_key: Optional[str] = None, + connection_id: Optional[str] = None, + timezone: Optional[str] = None, + **kwargs, + ): """Initialize CloudZero logger with configuration from parameters or environment variables.""" super().__init__(**kwargs) - + # Get configuration from parameters first, fall back to environment variables self.api_key = api_key or os.getenv("CLOUDZERO_API_KEY") - self.connection_id = connection_id or os.getenv("CLOUDZERO_CONNECTION_ID") + self.connection_id = connection_id or os.getenv("CLOUDZERO_CONNECTION_ID") self.timezone = timezone or os.getenv("CLOUDZERO_TIMEZONE", "UTC") - verbose_logger.debug(f"CloudZero Logger initialized with connection ID: {self.connection_id}, timezone: {self.timezone}") + verbose_logger.debug( + f"CloudZero Logger initialized with connection ID: {self.connection_id}, timezone: {self.timezone}" + ) async def initialize_cloudzero_export_job(self): """ @@ -46,6 +55,7 @@ class CloudZeroLogger(CustomLogger): CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME, ) from litellm.proxy.proxy_server import proxy_logging_obj + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager # if using redis, ensure only one pod exports the data at a time @@ -62,7 +72,7 @@ class CloudZeroLogger(CustomLogger): else: # if not using redis, export the data directly await self._hourly_usage_data_export() - + async def _hourly_usage_data_export(self): """ Exports the hourly usage data to CloudZero. @@ -73,22 +83,25 @@ class CloudZeroLogger(CustomLogger): from datetime import timedelta, timezone from litellm.constants import CLOUDZERO_MAX_FETCHED_DATA_RECORDS + current_time_utc = datetime.now(timezone.utc) - one_hour_ago_utc = current_time_utc - timedelta(hours=1) + # Mitigates the possibility of missing spend if an hour is skipped due to a restart in an ephemeral environment + one_hour_ago_utc = current_time_utc - timedelta( + minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES * 2 + ) await self.export_usage_data( limit=CLOUDZERO_MAX_FETCHED_DATA_RECORDS, operation="replace_hourly", start_time_utc=one_hour_ago_utc, - end_time_utc=current_time_utc + end_time_utc=current_time_utc, ) - async def export_usage_data( - self, - limit: Optional[int] = None, + self, + limit: Optional[int] = None, operation: str = "replace_hourly", start_time_utc: Optional[datetime] = None, - end_time_utc: Optional[datetime] = None + end_time_utc: Optional[datetime] = None, ): """ Exports the usage data to CloudZero. @@ -96,7 +109,7 @@ class CloudZeroLogger(CustomLogger): - Reads data from the DB - Transforms the data to the CloudZero format - Sends the data to CloudZero - + Args: limit: Optional limit on number of records to export operation: CloudZero operation type ("replace_hourly" or "sum") @@ -104,9 +117,10 @@ class CloudZeroLogger(CustomLogger): 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") - + # Validate required configuration if not self.api_key or not self.connection_id: raise ValueError( @@ -117,61 +131,68 @@ class CloudZeroLogger(CustomLogger): database = LiteLLMDatabase() verbose_logger.debug("CloudZero Logger: Loading usage data from database") data = await database.get_usage_data( - limit=limit, - start_time_utc=start_time_utc, - end_time_utc=end_time_utc + limit=limit, start_time_utc=start_time_utc, end_time_utc=end_time_utc ) - + if data.is_empty(): verbose_logger.debug("CloudZero Logger: No usage data found to export") return verbose_logger.debug(f"CloudZero Logger: Processing {len(data)} records") - + # Transform data to CloudZero CBF format transformer = CBFTransformer() cbf_data = transformer.transform(data) - + if cbf_data.is_empty(): - verbose_logger.warning("CloudZero Logger: No valid data after transformation") + verbose_logger.warning( + "CloudZero Logger: No valid data after transformation" + ) return # Send data to CloudZero streamer = CloudZeroStreamer( api_key=self.api_key, connection_id=self.connection_id, - user_timezone=self.timezone + user_timezone=self.timezone, + ) + + verbose_logger.debug( + f"CloudZero Logger: Transmitting {len(cbf_data)} records to CloudZero" ) - - verbose_logger.debug(f"CloudZero Logger: Transmitting {len(cbf_data)} records to CloudZero") streamer.send_batched(cbf_data, operation=operation) - - verbose_logger.debug(f"CloudZero Logger: Successfully exported {len(cbf_data)} records to CloudZero") - + + verbose_logger.debug( + f"CloudZero Logger: Successfully exported {len(cbf_data)} records to CloudZero" + ) + except Exception as e: - verbose_logger.error(f"CloudZero Logger: Error exporting usage data: {str(e)}") + verbose_logger.error( + f"CloudZero Logger: Error exporting usage data: {str(e)}" + ) raise async def dry_run_export_usage_data(self, limit: Optional[int] = 10000): """ Returns the data that would be exported to CloudZero without actually sending it. - + Args: limit: Limit number of records to display (default: 10000) - + 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") - + # Initialize database connection and load data database = LiteLLMDatabase() verbose_logger.debug("CloudZero Logger: Loading usage data for dry run") data = await database.get_usage_data(limit=limit) - + if data.is_empty(): verbose_logger.warning("CloudZero Dry Run: No usage data found") return { @@ -182,44 +203,70 @@ class CloudZeroLogger(CustomLogger): "total_cost": 0, "total_tokens": 0, "unique_accounts": 0, - "unique_services": 0 - } + "unique_services": 0, + }, } - verbose_logger.debug(f"CloudZero Dry Run: Processing {len(data)} records...") - + verbose_logger.debug( + f"CloudZero Dry Run: Processing {len(data)} records..." + ) + # Convert usage data to dict format for response usage_data_sample = data.head(50).to_dicts() # Return first 50 rows # Transform data to CloudZero CBF format transformer = CBFTransformer() cbf_data = transformer.transform(data) - + if cbf_data.is_empty(): - verbose_logger.warning("CloudZero Dry Run: No valid data after transformation") + verbose_logger.warning( + "CloudZero Dry Run: No valid data after transformation" + ) return { "usage_data": usage_data_sample, "cbf_data": [], "summary": { "total_records": len(usage_data_sample), - "total_cost": sum(row.get('spend', 0) for row in usage_data_sample), - "total_tokens": sum(row.get('prompt_tokens', 0) + row.get('completion_tokens', 0) for row in usage_data_sample), + "total_cost": sum( + row.get("spend", 0) for row in usage_data_sample + ), + "total_tokens": sum( + row.get("prompt_tokens", 0) + + row.get("completion_tokens", 0) + for row in usage_data_sample + ), "unique_accounts": 0, - "unique_services": 0 - } + "unique_services": 0, + }, } # Convert CBF data to dict format for response cbf_data_dict = cbf_data.to_dicts() - + # Calculate summary statistics - total_cost = sum(record.get('cost/cost', 0) for record in cbf_data_dict) - unique_accounts = len(set(record.get('resource/account', '') for record in cbf_data_dict if record.get('resource/account'))) - unique_services = len(set(record.get('resource/service', '') for record in cbf_data_dict if record.get('resource/service'))) - total_tokens = sum(record.get('usage/amount', 0) for record in cbf_data_dict) - - verbose_logger.debug(f"CloudZero Logger: Dry run completed for {len(cbf_data)} records") - + total_cost = sum(record.get("cost/cost", 0) for record in cbf_data_dict) + unique_accounts = len( + set( + record.get("resource/account", "") + for record in cbf_data_dict + if record.get("resource/account") + ) + ) + unique_services = len( + set( + record.get("resource/service", "") + for record in cbf_data_dict + if record.get("resource/service") + ) + ) + total_tokens = sum( + record.get("usage/amount", 0) for record in cbf_data_dict + ) + + verbose_logger.debug( + f"CloudZero Logger: Dry run completed for {len(cbf_data)} records" + ) + return { "usage_data": usage_data_sample, "cbf_data": cbf_data_dict, @@ -228,10 +275,10 @@ class CloudZeroLogger(CustomLogger): "total_cost": total_cost, "total_tokens": total_tokens, "unique_accounts": unique_accounts, - "unique_services": unique_services - } + "unique_services": unique_services, + }, } - + except Exception as e: verbose_logger.error(f"CloudZero Logger: Error in dry run export: {str(e)}") verbose_logger.error(f"CloudZero Dry Run Error: {str(e)}") @@ -242,28 +289,38 @@ class CloudZeroLogger(CustomLogger): from rich.box import SIMPLE from rich.console import Console from rich.table import Table - + console = Console() - + if cbf_data.is_empty(): console.print("[yellow]No CBF data to display[/yellow]") return - console.print(f"\n[bold green]💰 CloudZero CBF Transformed Data ({len(cbf_data)} records)[/bold green]") + console.print( + f"\n[bold green]💰 CloudZero CBF Transformed Data ({len(cbf_data)} records)[/bold green]" + ) # Convert to dicts for easier processing records = cbf_data.to_dicts() # Create main CBF table - cbf_table = Table(show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1)) + cbf_table = Table( + show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1) + ) cbf_table.add_column("time/usage_start", style="blue", no_wrap=False) cbf_table.add_column("cost/cost", style="green", justify="right", no_wrap=False) - cbf_table.add_column("entity_type", style="magenta", justify="right", no_wrap=False) - cbf_table.add_column("entity_id", style="magenta", justify="right", no_wrap=False) + cbf_table.add_column( + "entity_type", style="magenta", justify="right", no_wrap=False + ) + cbf_table.add_column( + "entity_id", style="magenta", justify="right", no_wrap=False + ) cbf_table.add_column("team_id", style="cyan", no_wrap=False) cbf_table.add_column("team_alias", style="cyan", no_wrap=False) cbf_table.add_column("api_key_alias", style="yellow", no_wrap=False) - cbf_table.add_column("usage/amount", style="yellow", justify="right", no_wrap=False) + cbf_table.add_column( + "usage/amount", style="yellow", justify="right", no_wrap=False + ) cbf_table.add_column("resource/id", style="magenta", no_wrap=False) cbf_table.add_column("resource/service", style="cyan", no_wrap=False) cbf_table.add_column("resource/account", style="white", no_wrap=False) @@ -271,18 +328,18 @@ class CloudZeroLogger(CustomLogger): for record in records: # Use proper CBF field names - time_usage_start = str(record.get('time/usage_start', 'N/A')) - cost_cost = str(record.get('cost/cost', 0)) - usage_amount = str(record.get('usage/amount', 0)) - resource_id = str(record.get('resource/id', 'N/A')) - resource_service = str(record.get('resource/service', 'N/A')) - resource_account = str(record.get('resource/account', 'N/A')) - resource_region = str(record.get('resource/region', 'N/A')) - entity_type = str(record.get('entity_type', 'N/A')) - entity_id = str(record.get('entity_id', 'N/A')) - team_id = str(record.get('resource/tag:team_id', 'N/A')) - team_alias = str(record.get('resource/tag:team_alias', 'N/A')) - api_key_alias = str(record.get('resource/tag:api_key_alias', 'N/A')) + time_usage_start = str(record.get("time/usage_start", "N/A")) + cost_cost = str(record.get("cost/cost", 0)) + usage_amount = str(record.get("usage/amount", 0)) + resource_id = str(record.get("resource/id", "N/A")) + resource_service = str(record.get("resource/service", "N/A")) + resource_account = str(record.get("resource/account", "N/A")) + resource_region = str(record.get("resource/region", "N/A")) + entity_type = str(record.get("entity_type", "N/A")) + entity_id = str(record.get("entity_id", "N/A")) + team_id = str(record.get("resource/tag:team_id", "N/A")) + team_alias = str(record.get("resource/tag:team_alias", "N/A")) + api_key_alias = str(record.get("resource/tag:api_key_alias", "N/A")) cbf_table.add_row( time_usage_start, @@ -296,18 +353,30 @@ class CloudZeroLogger(CustomLogger): resource_id, resource_service, resource_account, - resource_region + resource_region, ) console.print(cbf_table) # Show summary statistics - total_cost = sum(record.get('cost/cost', 0) for record in records) - unique_accounts = len(set(record.get('resource/account', '') for record in records if record.get('resource/account'))) - unique_services = len(set(record.get('resource/service', '') for record in records if record.get('resource/service'))) + total_cost = sum(record.get("cost/cost", 0) for record in records) + unique_accounts = len( + set( + record.get("resource/account", "") + for record in records + if record.get("resource/account") + ) + ) + unique_services = len( + set( + record.get("resource/service", "") + for record in records + if record.get("resource/service") + ) + ) # Count total tokens from usage metrics - total_tokens = sum(record.get('usage/amount', 0) for record in records) + total_tokens = sum(record.get("usage/amount", 0) for record in records) console.print("\n[bold blue]📊 CBF Summary[/bold blue]") console.print(f" Records: {len(records):,}") @@ -316,8 +385,10 @@ class CloudZeroLogger(CustomLogger): console.print(f" Unique Accounts: {unique_accounts}") console.print(f" Unique Services: {unique_services}") - console.print("\n[dim]💡 This is the CloudZero CBF format ready for AnyCost ingestion[/dim]") - + console.print( + "\n[dim]💡 This is the CloudZero CBF format ready for AnyCost ingestion[/dim]" + ) + @staticmethod async def init_cloudzero_background_job(scheduler: AsyncIOScheduler): """ @@ -327,12 +398,11 @@ class CloudZeroLogger(CustomLogger): """ from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - - prometheus_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CloudZeroLogger - ) + prometheus_loggers: List[ + CustomLogger + ] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CloudZeroLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s cloudzero loggers", len(prometheus_loggers)) @@ -345,5 +415,5 @@ class CloudZeroLogger(CustomLogger): scheduler.add_job( cloudzero_logger.initialize_cloudzero_export_job, "interval", - minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES - ) \ No newline at end of file + minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES, + ) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 71b4125ed7..83ca01a5c0 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -26,6 +26,7 @@ import polars as pl class LiteLLMDatabase: """Handle LiteLLM PostgreSQL database connections and queries.""" + def _ensure_prisma_client(self): from litellm.proxy.proxy_server import prisma_client @@ -37,25 +38,25 @@ class LiteLLMDatabase: return prisma_client async def get_usage_data( - self, + self, limit: Optional[int] = None, start_time_utc: Optional[datetime] = None, - end_time_utc: Optional[datetime] = None + end_time_utc: Optional[datetime] = None, ) -> pl.DataFrame: """Retrieve usage data from LiteLLM daily user spend table.""" client = self._ensure_prisma_client() - + # Build WHERE clause for time filtering where_conditions = [] if start_time_utc: - where_conditions.append(f"dus.created_at >= '{start_time_utc.isoformat()}'") + where_conditions.append(f"dus.updated_at >= '{start_time_utc.isoformat()}'") if end_time_utc: - where_conditions.append(f"dus.created_at <= '{end_time_utc.isoformat()}'") - + where_conditions.append(f"dus.updated_at <= '{end_time_utc.isoformat()}'") + where_clause = "" if where_conditions: where_clause = "WHERE " + " AND ".join(where_conditions) - + # Query to get user spend data with team information query = f""" SELECT @@ -100,10 +101,10 @@ class LiteLLMDatabase: async def get_table_info(self) -> Dict[str, Any]: """Get information about the daily user spend table.""" client = self._ensure_prisma_client() - + try: # Get row count from user spend table - user_count = await self._get_table_row_count('LiteLLM_DailyUserSpend') + user_count = await self._get_table_row_count("LiteLLM_DailyUserSpend") # Get column structure from user spend table query = """ @@ -115,9 +116,9 @@ class LiteLLMDatabase: columns_response = await client.db.query_raw(query) return { - 'columns': columns_response, - 'row_count': user_count, - 'table_name': 'LiteLLM_DailyUserSpend' + "columns": columns_response, + "row_count": user_count, + "table_name": "LiteLLM_DailyUserSpend", } except Exception as e: raise Exception(f"Error getting table info: {str(e)}") @@ -125,13 +126,13 @@ class LiteLLMDatabase: async def _get_table_row_count(self, table_name: str) -> int: """Get row count from specified table.""" client = self._ensure_prisma_client() - + try: query = f'SELECT COUNT(*) as count FROM "{table_name}"' response = await client.db.query_raw(query) - + if response and len(response) > 0: - return response[0].get('count', 0) + return response[0].get("count", 0) return 0 except Exception: return 0 @@ -139,7 +140,7 @@ class LiteLLMDatabase: async def discover_all_tables(self) -> Dict[str, Any]: """Discover all tables in the LiteLLM database and their schemas.""" client = self._ensure_prisma_client() - + try: # Get all LiteLLM tables litellm_tables_query = """ @@ -150,7 +151,7 @@ class LiteLLMDatabase: ORDER BY table_name; """ tables_response = await client.db.query_raw(litellm_tables_query) - table_names = [row['table_name'] for row in tables_response] + table_names = [row["table_name"] for row in tables_response] # Get detailed schema for each table tables_info = {} @@ -181,7 +182,9 @@ class LiteLLMDatabase: WHERE i.indrelid = $1::regclass AND i.indisprimary; """ pk_response = await client.db.query_raw(pk_query, f'"{table_name}"') - primary_keys = [row['attname'] for row in pk_response] if pk_response else [] + primary_keys = ( + [row["attname"] for row in pk_response] if pk_response else [] + ) # Get foreign key information fk_query = """ @@ -226,18 +229,17 @@ class LiteLLMDatabase: row_count = 0 tables_info[table_name] = { - 'columns': columns_response, - 'primary_keys': primary_keys, - 'foreign_keys': foreign_keys, - 'indexes': indexes, - 'row_count': row_count + "columns": columns_response, + "primary_keys": primary_keys, + "foreign_keys": foreign_keys, + "indexes": indexes, + "row_count": row_count, } return { - 'tables': tables_info, - 'table_count': len(table_names), - 'table_names': table_names + "tables": tables_info, + "table_count": len(table_names), + "table_names": table_names, } except Exception as e: raise Exception(f"Error discovering tables: {str(e)}") - diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py new file mode 100644 index 0000000000..586ab43350 --- /dev/null +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -0,0 +1,85 @@ +import pytest +import polars as pl + +from unittest.mock import AsyncMock, MagicMock, patch +from datetime import datetime +from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger +from litellm.integrations.cloudzero.cz_stream_api import CloudZeroStreamer +from litellm.integrations.cloudzero.database import LiteLLMDatabase + + +class TestCloudZeroHourlyExport: + @pytest.mark.asyncio + async def test_hourly_export(self): + spend_mock_data = pl.LazyFrame( + { + "id": ["09327a4f-fa99-4613-86c5-23efb03640b1", "c7bcec65-0d76-4126-93b6-50fea1cdd2b"], + "user_id": ["069e8205-8f55-44fd-870b-0c036cab600c", "069e8205-8f55-44fd-870b-0c036cab600c"], + "date": ["2025-11-01", "2025-11-01"], + "api_key": [ + "c1465c9a821f420927b3d81972323fb516745bc93a4a54ceca0ce6ddf6100c39", + "c1465c9a821f420927b3d81972323fb516745bc93a4a54ceca0ce6ddf6100c39", + ], + "model": ["model_1", "model_2"], + "model_group": ["model_group_1", "model_group_2"], + "custom_llm_provider": ["provider_1", "provider_2"], + "prompt_tokens": [60, 60], + "completion_tokens": [71, 71], + "spend": [0.005, 0.005], + "api_requests": [1, 1], + "successful_requests": [1, 1], + "failed_requests": [0, 0], + "cache_creation_input_tokens": [0, 0], + "cache_read_input_tokens": [0, 0], + "created_at": [datetime(2025, 11, 1, 12), datetime(2025, 11, 1, 2)], + "updated_at": [datetime(2025, 11, 1, 12), datetime(2025, 11, 1, 12)], + } + ) + + team_mock_data = pl.LazyFrame( + { + "team_id": ["a3d6b0bb-098f-4260-81d6-fabae695b622"], + "team_alias": ["team_1"], + } + ) + verification_mock_data = pl.LazyFrame( + { + "team_id": ["a3d6b0bb-098f-4260-81d6-fabae695b622"], + "key_alias": ["key_1"], + "token": ["c1465c9a821f420927b3d81972323fb516745bc93a4a54ceca0ce6ddf6100c39"], + } + ) + + with ( + patch.object(LiteLLMDatabase, "_ensure_prisma_client") as mock_prisma_client_getter, + patch.object(CloudZeroStreamer, "send_batched") as send_batched_mock, + patch("litellm.integrations.cloudzero.cloudzero.datetime") as mock_datetime, + ): + fake_client = MagicMock() + fake_db = MagicMock() + + async def query_raw_mock(query: str): + sql_context = pl.SQLContext( + LiteLLM_DailyUserSpend=spend_mock_data, + LiteLLM_VerificationToken=verification_mock_data, + LiteLLM_TeamTable=team_mock_data, + ) + result = sql_context.execute(query).collect() + + return result + + fake_db.query_raw = AsyncMock(side_effect=query_raw_mock) + fake_client.db = fake_db + mock_prisma_client_getter.return_value = fake_client + + mock_datetime.now.return_value = datetime(2025, 11, 1, 12, 0, 1) + + def export_verifier(cbf_data, operation): + assert operation == "replace_hourly" + assert len(cbf_data) == 2 + + send_batched_mock.side_effect = export_verifier + + logger = CloudZeroLogger(api_key="test", connection_id="test") + + await logger._hourly_usage_data_export()