Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/spend_tracking/cloudzero_endpoints.py: 70%
166 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import json
2from collections.abc import Mapping
3from typing import (
4 TYPE_CHECKING,
5 Final,
6 Protocol,
7 cast, # noqa: TID251 # the config repository's table protocol omits find_first
8)
10from fastapi import APIRouter, Depends, HTTPException
12from litellm._logging import verbose_proxy_logger
13from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
14from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
15from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
16from litellm.proxy.common_utils.encrypt_decrypt_utils import (
17 decrypt_value_helper,
18 encrypt_value_helper,
19)
20from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
21from litellm.repositories.config_repository import ConfigRepository
22from litellm.repositories.prisma_protocols import TableActions
23from litellm.types.proxy.cloudzero_endpoints import (
24 CloudZeroExportRequest,
25 CloudZeroExportResponse,
26 CloudZeroInitRequest,
27 CloudZeroInitResponse,
28 CloudZeroSettingsUpdate,
29 CloudZeroSettingsView,
30)
32if TYPE_CHECKING: 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true
33 from litellm.proxy.proxy_server import PrismaClient
35router: Final = APIRouter()
38# Initialize the sensitive data masker for API key masking
39_sensitive_masker: Final = SensitiveDataMasker()
42class _CloudZeroConfigRow(Protocol):
43 """The ``LiteLLM_Config`` row holding ``cloudzero_settings``, as this module reads it."""
45 @property
46 def param_value(self) -> str | Mapping[str, str] | None: ... 46 ↛ exitline 46 didn't return from function 'param_value' because
49def _config_table(prisma_client: "PrismaClient") -> TableActions[_CloudZeroConfigRow]:
50 repository_table: Final = ConfigRepository(prisma_client).table
51 return cast(TableActions[_CloudZeroConfigRow], repository_table) # cast-ok: repo protocol omits find_first
54async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: str):
55 """
56 Store CloudZero settings in the database with encrypted API key.
58 Args:
59 api_key: CloudZero API key to encrypt and store
60 connection_id: CloudZero connection ID
61 timezone: Timezone for date handling
62 """
63 from litellm.proxy.proxy_server import prisma_client
65 if prisma_client is None: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 raise HTTPException(
67 status_code=500,
68 detail={"error": CommonProxyErrors.db_not_connected_error.value},
69 )
71 # Encrypt the API key before storing
72 encrypted_api_key: Final = encrypt_value_helper(api_key)
74 cloudzero_settings: Final = {
75 "api_key": encrypted_api_key,
76 "connection_id": connection_id,
77 "timezone": timezone,
78 }
80 await ConfigRepository(prisma_client).table.upsert(
81 where={"param_name": "cloudzero_settings"},
82 data={
83 "create": {
84 "param_name": "cloudzero_settings",
85 "param_value": json.dumps(cloudzero_settings),
86 },
87 "update": {"param_value": json.dumps(cloudzero_settings)},
88 },
89 )
92async def _get_cloudzero_settings():
93 """
94 Retrieve CloudZero settings from the database with decrypted API key.
96 Returns:
97 dict: CloudZero settings with decrypted API key, or empty dict if not configured
98 """
99 from litellm.proxy.proxy_server import prisma_client
101 if prisma_client is None: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 raise HTTPException(
103 status_code=500,
104 detail={"error": CommonProxyErrors.db_not_connected_error.value},
105 )
107 cloudzero_config: Final = await _config_table(prisma_client).find_first(where={"param_name": "cloudzero_settings"})
108 if cloudzero_config is None or cloudzero_config.param_value is None:
109 return {}
111 # Handle both dict and JSON string cases
112 if isinstance(cloudzero_config.param_value, dict): 112 ↛ 114line 112 didn't jump to line 114 because the condition on line 112 was always true
113 settings = cloudzero_config.param_value
114 elif isinstance(cloudzero_config.param_value, str):
115 settings = json.loads(cloudzero_config.param_value)
116 else:
117 settings = dict(cloudzero_config.param_value)
119 # Decrypt the API key
120 encrypted_api_key: Final = settings.get("api_key")
121 if encrypted_api_key: 121 ↛ 130line 121 didn't jump to line 130 because the condition on line 121 was always true
122 decrypted_api_key = decrypt_value_helper(encrypted_api_key, key="cloudzero_api_key", exception_type="error")
123 if decrypted_api_key is None: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true
124 raise HTTPException(
125 status_code=500,
126 detail={"error": "Failed to decrypt CloudZero API key. Check your salt key configuration."},
127 )
128 settings["api_key"] = decrypted_api_key
130 return settings
133@router.get(
134 "/cloudzero/settings",
135 tags=["CloudZero"],
136 dependencies=[Depends(user_api_key_auth)],
137 response_model=CloudZeroSettingsView,
138)
139async def get_cloudzero_settings(
140 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
141):
142 """
143 View current CloudZero settings.
145 Returns the current CloudZero configuration with the API key masked for security.
146 Only the first 4 and last 4 characters of the API key are shown.
147 Returns null/empty values when settings are not configured (consistent with other settings endpoints).
149 Only admin users (Proxy Admin or Admin Viewer) can view CloudZero settings.
150 """
151 # Validation — Admin Viewer follows the read-parity rule.
152 if not _user_has_admin_view(user_api_key_dict): 152 ↛ 153line 152 didn't jump to line 153 because the condition on line 152 was never true
153 raise HTTPException(
154 status_code=403,
155 detail={"error": CommonProxyErrors.not_allowed_access.value},
156 )
158 try:
159 # Get CloudZero settings using the accessor method
160 settings: Final = await _get_cloudzero_settings()
162 # If settings are empty, return null/empty values (consistent with other endpoints)
163 if not settings:
164 return CloudZeroSettingsView(
165 api_key_masked=None,
166 connection_id=None,
167 timezone=None,
168 status=None,
169 )
171 # Use SensitiveDataMasker to mask the API key
172 masked_settings: Final = _sensitive_masker.mask_dict(settings)
174 return CloudZeroSettingsView(
175 api_key_masked=masked_settings.get("api_key"),
176 connection_id=settings.get("connection_id"),
177 timezone=settings.get("timezone"),
178 status="configured",
179 )
181 except HTTPException as e:
182 # Re-raise HTTPExceptions as-is
183 raise e
184 except Exception as e:
185 verbose_proxy_logger.error("Error retrieving CloudZero settings: %s", e)
186 raise HTTPException(
187 status_code=500,
188 detail={"error": f"Failed to retrieve CloudZero settings: {e}"},
189 )
192@router.put(
193 "/cloudzero/settings",
194 tags=["CloudZero"],
195 dependencies=[Depends(user_api_key_auth)],
196 response_model=CloudZeroInitResponse,
197)
198async def update_cloudzero_settings(
199 request: CloudZeroSettingsUpdate,
200 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
201):
202 """
203 Update existing CloudZero settings.
205 Allows updating individual CloudZero configuration fields without requiring all fields.
206 Only provided fields will be updated; others will remain unchanged.
208 Parameters:
209 - api_key: (Optional) New CloudZero API key for authentication
210 - connection_id: (Optional) New CloudZero connection ID for data submission
211 - timezone: (Optional) New timezone for date handling
213 Only admin users can update CloudZero settings.
214 """
215 # Validation
216 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 raise HTTPException(
218 status_code=403,
219 detail={"error": CommonProxyErrors.not_allowed_access.value},
220 )
222 # Check if at least one field is provided
223 if not any([request.api_key, request.connection_id, request.timezone]):
224 raise HTTPException(
225 status_code=400,
226 detail={"error": "At least one field must be provided for update"},
227 )
229 try:
230 # Get current settings
231 current_settings: Final = await _get_cloudzero_settings()
233 # Update only provided fields
234 updated_api_key: Final = request.api_key if request.api_key is not None else current_settings["api_key"]
235 updated_connection_id: Final = (
236 request.connection_id if request.connection_id is not None else current_settings["connection_id"]
237 )
238 updated_timezone: Final = request.timezone if request.timezone is not None else current_settings["timezone"]
240 # Store updated settings using the setter method with encryption
241 await _set_cloudzero_settings(
242 api_key=updated_api_key,
243 connection_id=updated_connection_id,
244 timezone=updated_timezone,
245 )
247 verbose_proxy_logger.info("CloudZero settings updated successfully")
249 return CloudZeroInitResponse(message="CloudZero settings updated successfully", status="success")
251 except HTTPException as e:
252 if e.status_code == 400:
253 # Settings not configured yet
254 raise HTTPException(
255 status_code=404,
256 detail={
257 "error": "CloudZero settings not found. Please initialize settings first using /cloudzero/init"
258 },
259 )
260 raise e
261 except Exception as e:
262 verbose_proxy_logger.error("Error updating CloudZero settings: %s", e)
263 raise HTTPException(
264 status_code=500,
265 detail={"error": f"Failed to update CloudZero settings: {e}"},
266 )
269# Global variable to track if CloudZero background job has been initialized
270_cloudzero_background_job_initialized: Final = False
273async def is_cloudzero_setup_in_db() -> bool:
274 """
275 Check if CloudZero is setup in the database.
277 CloudZero is considered setup in the database if:
278 - CloudZero settings exist in the database
279 - The settings have a non-None value
281 Returns:
282 bool: True if CloudZero is active, False otherwise
283 """
284 try:
285 from litellm.proxy.proxy_server import prisma_client
287 if prisma_client is None: 287 ↛ 288line 287 didn't jump to line 288 because the condition on line 287 was never true
288 return False
290 # Check for CloudZero settings in database
291 cloudzero_config: Final = await _config_table(prisma_client).find_first(
292 where={"param_name": "cloudzero_settings"}
293 )
295 # CloudZero is setup in the database if config exists and has non-None value
296 return cloudzero_config is not None and cloudzero_config.param_value is not None
298 except Exception as e:
299 verbose_proxy_logger.error("Error checking CloudZero status: %s", e)
300 return False
303def is_cloudzero_setup_in_config() -> bool:
304 """
305 Check if CloudZero is setup in config.yaml or environment variables.
307 CloudZero is considered setup in config if:
308 - "cloudzero" is in the callbacks list in config.yaml, OR
309 Returns:
310 bool: True if CloudZero is configured, False otherwise
311 """
312 import litellm
314 return "cloudzero" in litellm.callbacks
317async def is_cloudzero_setup() -> bool:
318 """
319 Check if CloudZero is setup in either config.yaml/env vars OR database.
321 CloudZero is considered setup if:
322 - CloudZero is configured in config.yaml callbacks, OR
323 - CloudZero environment variables are set, OR
324 - CloudZero settings exist in the database
326 Returns:
327 bool: True if CloudZero is configured anywhere, False otherwise
328 """
329 try:
330 # Check config.yaml/environment variables first
331 if is_cloudzero_setup_in_config(): 331 ↛ 332line 331 didn't jump to line 332 because the condition on line 331 was never true
332 return True
334 # Check database as fallback
335 if await is_cloudzero_setup_in_db(): 335 ↛ 336line 335 didn't jump to line 336 because the condition on line 335 was never true
336 return True
338 return False
340 except Exception as e:
341 verbose_proxy_logger.error("Error checking CloudZero setup: %s", e)
342 return False
345@router.post(
346 "/cloudzero/init",
347 tags=["CloudZero"],
348 dependencies=[Depends(user_api_key_auth)],
349 response_model=CloudZeroInitResponse,
350)
351async def init_cloudzero_settings(
352 request: CloudZeroInitRequest,
353 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
354):
355 """
356 Initialize CloudZero settings and store in the database.
358 This endpoint stores the CloudZero API key, connection ID, and timezone configuration
359 in the proxy database for use by the CloudZero logger.
361 Parameters:
362 - api_key: CloudZero API key for authentication
363 - connection_id: CloudZero connection ID for data submission
364 - timezone: Timezone for date handling (default: UTC)
366 Only admin users can configure CloudZero settings.
367 """
368 # Validation
369 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 369 ↛ 370line 369 didn't jump to line 370 because the condition on line 369 was never true
370 raise HTTPException(
371 status_code=403,
372 detail={"error": CommonProxyErrors.not_allowed_access.value},
373 )
375 try:
376 # Store settings using the setter method with encryption
377 await _set_cloudzero_settings(
378 api_key=request.api_key,
379 connection_id=request.connection_id,
380 timezone=request.timezone,
381 )
383 verbose_proxy_logger.info("CloudZero settings initialized successfully")
385 return CloudZeroInitResponse(message="CloudZero settings initialized successfully", status="success")
387 except Exception as e:
388 verbose_proxy_logger.error("Error initializing CloudZero settings: %s", e)
389 raise HTTPException(
390 status_code=500,
391 detail={"error": f"Failed to initialize CloudZero settings: {e}"},
392 )
395@router.post(
396 "/cloudzero/dry-run",
397 tags=["CloudZero"],
398 dependencies=[Depends(user_api_key_auth)],
399 response_model=CloudZeroExportResponse,
400)
401async def cloudzero_dry_run_export(
402 request: CloudZeroExportRequest,
403 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
404):
405 """
406 Perform a dry run export using the CloudZero logger.
408 This endpoint uses the CloudZero logger to perform a dry run export,
409 which returns the data that would be exported without actually sending it to CloudZero.
411 Parameters:
412 - limit: Optional limit on number of records to process (default: 10000)
414 Returns:
415 - usage_data: Sample of the raw usage data (first 50 records)
416 - cbf_data: CloudZero CBF formatted data ready for export
417 - summary: Statistics including total cost, tokens, and record counts
419 Only admin users can perform CloudZero exports.
420 """
421 # Validation
422 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 raise HTTPException(
424 status_code=403,
425 detail={"error": CommonProxyErrors.not_allowed_access.value},
426 )
428 try:
429 # Import and initialize CloudZero logger with credentials
430 from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
432 # Initialize logger with credentials directly
433 logger: Final = CloudZeroLogger()
434 dry_run_result: Final = await logger.dry_run_export_usage_data(limit=request.limit)
436 verbose_proxy_logger.info("CloudZero dry run export completed successfully")
438 return CloudZeroExportResponse(
439 message="CloudZero dry run export completed successfully.",
440 status="success",
441 dry_run_data=dry_run_result,
442 summary=dry_run_result.get("summary") if dry_run_result else None,
443 )
445 except Exception as e:
446 verbose_proxy_logger.error("Error performing CloudZero dry run export: %s", e)
447 raise HTTPException(
448 status_code=500,
449 detail={"error": f"Failed to perform CloudZero dry run export: {e}"},
450 )
453@router.post(
454 "/cloudzero/export",
455 tags=["CloudZero"],
456 dependencies=[Depends(user_api_key_auth)],
457 response_model=CloudZeroExportResponse,
458)
459async def cloudzero_export(
460 request: CloudZeroExportRequest,
461 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
462):
463 """
464 Perform an actual export using the CloudZero logger.
466 This endpoint uses the CloudZero logger to export usage data to CloudZero AnyCost API.
468 Parameters:
469 - limit: Optional limit on number of records to export
470 - operation: CloudZero operation type ("replace_hourly" or "sum", default: "replace_hourly")
472 Only admin users can perform CloudZero exports.
473 """
475 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 475 ↛ 476line 475 didn't jump to line 476 because the condition on line 475 was never true
476 raise HTTPException(
477 status_code=403,
478 detail={"error": CommonProxyErrors.not_allowed_access.value},
479 )
481 try:
482 # Get CloudZero settings using the accessor method with decryption
483 settings: Final = await _get_cloudzero_settings()
485 # Import and initialize CloudZero logger with credentials
486 from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
488 # Initialize logger with credentials directly
489 logger: Final = CloudZeroLogger(
490 api_key=settings.get("api_key"),
491 connection_id=settings.get("connection_id"),
492 timezone=settings.get("timezone"),
493 )
494 await logger.export_usage_data(
495 limit=request.limit,
496 operation=request.operation,
497 start_time_utc=request.start_time_utc,
498 end_time_utc=request.end_time_utc,
499 )
501 verbose_proxy_logger.info("CloudZero export completed successfully")
503 return CloudZeroExportResponse(
504 message="CloudZero export completed successfully",
505 status="success",
506 dry_run_data=None,
507 summary=None,
508 )
510 except Exception as e:
511 verbose_proxy_logger.error("Error performing CloudZero export: %s", e)
512 raise HTTPException(
513 status_code=500,
514 detail={"error": f"Failed to perform CloudZero export: {e}"},
515 )
518@router.delete(
519 "/cloudzero/delete",
520 tags=["CloudZero"],
521 dependencies=[Depends(user_api_key_auth)],
522 response_model=CloudZeroInitResponse,
523)
524async def delete_cloudzero_settings(
525 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
526):
527 """
528 Delete CloudZero settings from the database.
530 This endpoint removes the CloudZero configuration (API key, connection ID, timezone)
531 from the proxy database. Only the CloudZero settings entry will be deleted;
532 other configuration values in the database will remain unchanged.
534 Only admin users can delete CloudZero settings.
535 """
536 # Validation
537 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 537 ↛ 538line 537 didn't jump to line 538 because the condition on line 537 was never true
538 raise HTTPException(
539 status_code=403,
540 detail={"error": CommonProxyErrors.not_allowed_access.value},
541 )
543 try:
544 from litellm.proxy.proxy_server import prisma_client
546 if prisma_client is None: 546 ↛ 547line 546 didn't jump to line 547 because the condition on line 546 was never true
547 raise HTTPException(
548 status_code=500,
549 detail={"error": CommonProxyErrors.db_not_connected_error.value},
550 )
552 # Check if CloudZero settings exist
553 cloudzero_config: Final = await _config_table(prisma_client).find_first(
554 where={"param_name": "cloudzero_settings"}
555 )
557 if cloudzero_config is None:
558 raise HTTPException(
559 status_code=404,
560 detail={"error": "CloudZero settings not found"},
561 )
563 # Delete only the CloudZero settings entry
564 # This uses a specific where clause to target only the cloudzero_settings row
565 await ConfigRepository(prisma_client).table.delete(where={"param_name": "cloudzero_settings"})
567 verbose_proxy_logger.info("CloudZero settings deleted successfully")
569 return CloudZeroInitResponse(message="CloudZero settings deleted successfully", status="success")
571 except HTTPException as e:
572 raise e
573 except Exception as e:
574 verbose_proxy_logger.error("Error deleting CloudZero settings: %s", e)
575 raise HTTPException(
576 status_code=500,
577 detail={"error": f"Failed to delete CloudZero settings: {e}"},
578 )