tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 04:28:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Bug] Updated spend would not be sent to CloudZero (#16201)

Browse Source 
* Address a bug where cloudzero spend is not sent to cloudzero if a spend
update happens

* revert change unrelated to PR

* use polars for mocking instead of sqlite
This commit is contained in:
Daniel Sabanov
2025-11-10 19:43:15 -08:00
committed by GitHub
parent b6dbd4fa28
commit 0ecc38519e
3 changed files with 270 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/cloudzero/cloudzero.py
+156 -86
View File
@@ -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
)
minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES,
)
litellm/integrations/cloudzero/database.py
+29 -27
View File
@@ -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)}")
tests/test_litellm/integrations/cloudzero/test_cloudzero.py
+85
View File
@@ -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()