Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/spend_tracking/vantage_endpoints.py: 67%
211 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
12import litellm
13from litellm._logging import verbose_proxy_logger
14from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
15from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
16from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
17from litellm.proxy.common_utils.encrypt_decrypt_utils import (
18 decrypt_value_helper,
19 encrypt_value_helper,
20)
21from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
22from litellm.repositories.config_repository import ConfigRepository
23from litellm.repositories.prisma_protocols import TableActions
24from litellm.types.proxy.vantage_endpoints import (
25 VantageDryRunRequest,
26 VantageExportRequest,
27 VantageExportResponse,
28 VantageInitRequest,
29 VantageInitResponse,
30 VantageSettingsUpdate,
31 VantageSettingsView,
32)
34if TYPE_CHECKING: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true
35 from litellm.proxy.proxy_server import PrismaClient
37router: Final = APIRouter()
39_sensitive_masker: Final = SensitiveDataMasker()
41VANTAGE_SETTINGS_PARAM_NAME: Final = "vantage_settings"
44class _VantageConfigRow(Protocol):
45 """The ``LiteLLM_Config`` row holding ``vantage_settings``, as this module reads it."""
47 @property
48 def param_value(self) -> str | Mapping[str, str] | None: ... 48 ↛ exitline 48 didn't return from function 'param_value' because
51def _config_table(prisma_client: "PrismaClient") -> TableActions[_VantageConfigRow]:
52 repository_table: Final = ConfigRepository(prisma_client).table
53 return cast(TableActions[_VantageConfigRow], repository_table) # cast-ok: repo protocol omits find_first
56def _get_registered_vantage_logger():
57 """Return the VantageLogger already registered in litellm.callbacks, if any."""
58 from litellm.integrations.vantage.vantage_logger import VantageLogger
60 vantage_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=VantageLogger)
61 if vantage_loggers: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 return vantage_loggers[0]
63 return None
66async def _set_vantage_settings(api_key: str, integration_token: str, base_url: str):
67 """Store Vantage settings in the database with encrypted API key."""
68 from litellm.proxy.proxy_server import prisma_client
70 if prisma_client is None: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 raise HTTPException(
72 status_code=500,
73 detail={"error": CommonProxyErrors.db_not_connected_error.value},
74 )
76 encrypted_api_key: Final = encrypt_value_helper(api_key)
77 encrypted_integration_token: Final = encrypt_value_helper(integration_token)
79 vantage_settings: Final = {
80 "api_key": encrypted_api_key,
81 "integration_token": encrypted_integration_token,
82 "base_url": base_url,
83 }
85 await ConfigRepository(prisma_client).table.upsert(
86 where={"param_name": VANTAGE_SETTINGS_PARAM_NAME},
87 data={
88 "create": {
89 "param_name": VANTAGE_SETTINGS_PARAM_NAME,
90 "param_value": json.dumps(vantage_settings),
91 },
92 "update": {"param_value": json.dumps(vantage_settings)},
93 },
94 )
97async def _get_vantage_settings():
98 """Retrieve Vantage settings from the database with decrypted API key."""
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 vantage_config: Final = await _config_table(prisma_client).find_first(
108 where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
109 )
110 if vantage_config is None or vantage_config.param_value is None:
111 return {}
113 if isinstance(vantage_config.param_value, dict): 113 ↛ 115line 113 didn't jump to line 115 because the condition on line 113 was always true
114 settings = vantage_config.param_value
115 elif isinstance(vantage_config.param_value, str):
116 settings = json.loads(vantage_config.param_value)
117 else:
118 settings = dict(vantage_config.param_value)
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="vantage_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 Vantage API key. Check your salt key configuration."},
127 )
128 settings["api_key"] = decrypted_api_key
130 encrypted_integration_token: Final = settings.get("integration_token")
131 if encrypted_integration_token: 131 ↛ 144line 131 didn't jump to line 144 because the condition on line 131 was always true
132 decrypted_integration_token: Final = decrypt_value_helper(
133 encrypted_integration_token,
134 key="vantage_integration_token",
135 exception_type="error",
136 )
137 if decrypted_integration_token is None: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 raise HTTPException(
139 status_code=500,
140 detail={"error": "Failed to decrypt Vantage integration token. Check your salt key configuration."},
141 )
142 settings["integration_token"] = decrypted_integration_token
144 return settings
147@router.get(
148 "/vantage/settings",
149 tags=["Vantage"],
150 dependencies=[Depends(user_api_key_auth)],
151 response_model=VantageSettingsView,
152)
153async def get_vantage_settings(
154 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
155):
156 """
157 View current Vantage settings.
159 Returns the current Vantage configuration with the API key masked for security.
160 Only admin users (Proxy Admin or Admin Viewer) can view Vantage settings.
161 """
162 # Admin Viewer follows the read-parity rule.
163 if not _user_has_admin_view(user_api_key_dict): 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 raise HTTPException(
165 status_code=403,
166 detail={"error": CommonProxyErrors.not_allowed_access.value},
167 )
169 try:
170 settings: Final = await _get_vantage_settings()
172 if not settings:
173 return VantageSettingsView(
174 api_key_masked=None,
175 integration_token_masked=None,
176 base_url=None,
177 status=None,
178 )
180 masked_settings: Final = _sensitive_masker.mask_dict(settings)
182 return VantageSettingsView(
183 api_key_masked=masked_settings.get("api_key"),
184 integration_token_masked=masked_settings.get("integration_token"),
185 base_url=settings.get("base_url"),
186 status="configured",
187 )
189 except HTTPException:
190 raise
191 except Exception as e:
192 verbose_proxy_logger.error("Error retrieving Vantage settings: %s", e)
193 raise HTTPException(
194 status_code=500,
195 detail={"error": f"Failed to retrieve Vantage settings: {e}"},
196 )
199@router.put(
200 "/vantage/settings",
201 tags=["Vantage"],
202 dependencies=[Depends(user_api_key_auth)],
203 response_model=VantageInitResponse,
204)
205async def update_vantage_settings(
206 request: VantageSettingsUpdate,
207 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
208):
209 """
210 Update existing Vantage settings.
212 Allows updating individual Vantage configuration fields without requiring all fields.
213 Only admin users can update Vantage settings.
214 """
215 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 215 ↛ 216line 215 didn't jump to line 216 because the condition on line 215 was never true
216 raise HTTPException(
217 status_code=403,
218 detail={"error": CommonProxyErrors.not_allowed_access.value},
219 )
221 if not any([request.api_key, request.integration_token, request.base_url]):
222 raise HTTPException(
223 status_code=400,
224 detail={"error": "At least one field must be provided for update"},
225 )
227 try:
228 current_settings: Final = await _get_vantage_settings()
230 if not current_settings: 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 raise HTTPException(
232 status_code=404,
233 detail={"error": "Vantage settings not found. Please initialize settings first using /vantage/init"},
234 )
236 updated_api_key: Final = request.api_key if request.api_key is not None else current_settings.get("api_key", "")
237 updated_token: Final = (
238 request.integration_token
239 if request.integration_token is not None
240 else current_settings.get("integration_token", "")
241 )
242 updated_base_url: Final = (
243 request.base_url
244 if request.base_url is not None
245 else current_settings.get("base_url", "https://api.vantage.sh")
246 )
248 await _set_vantage_settings(
249 api_key=updated_api_key,
250 integration_token=updated_token,
251 base_url=updated_base_url,
252 )
254 verbose_proxy_logger.info("Vantage settings updated successfully")
256 return VantageInitResponse(message="Vantage settings updated successfully", status="success")
258 except HTTPException:
259 raise
260 except Exception as e:
261 verbose_proxy_logger.error("Error updating Vantage settings: %s", e)
262 raise HTTPException(
263 status_code=500,
264 detail={"error": f"Failed to update Vantage settings: {e}"},
265 )
268async def is_vantage_setup_in_db() -> bool:
269 """Check if Vantage is setup in the database."""
270 try:
271 from litellm.proxy.proxy_server import prisma_client
273 if prisma_client is None: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 return False
276 vantage_config: Final = await _config_table(prisma_client).find_first(
277 where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
278 )
280 return vantage_config is not None and vantage_config.param_value is not None
282 except Exception as e:
283 verbose_proxy_logger.error("Error checking Vantage status: %s", e)
284 return False
287def is_vantage_setup_in_config() -> bool:
288 """Check if Vantage is setup in config.yaml, environment variables, or programmatically."""
289 from litellm.integrations.vantage.vantage_logger import VantageLogger
291 for cb in litellm.callbacks:
292 if cb == "vantage" or isinstance(cb, VantageLogger): 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true
293 return True
294 return False
297async def is_vantage_setup() -> bool:
298 """Check if Vantage is setup in either config or database."""
299 try:
300 if is_vantage_setup_in_config(): 300 ↛ 301line 300 didn't jump to line 301 because the condition on line 300 was never true
301 return True
302 if await is_vantage_setup_in_db(): 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true
303 return True
304 return False
305 except Exception as e:
306 verbose_proxy_logger.error("Error checking Vantage setup: %s", e)
307 return False
310@router.post(
311 "/vantage/init",
312 tags=["Vantage"],
313 dependencies=[Depends(user_api_key_auth)],
314 response_model=VantageInitResponse,
315)
316async def init_vantage_settings(
317 request: VantageInitRequest,
318 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
319):
320 """
321 Initialize Vantage settings and store in the database.
323 Parameters:
324 - api_key: Vantage API key for authentication
325 - integration_token: Vantage integration token for the cost-import endpoint
326 - base_url: Vantage API base URL (default: https://api.vantage.sh)
328 Only admin users can configure Vantage settings.
329 """
330 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 330 ↛ 331line 330 didn't jump to line 331 because the condition on line 330 was never true
331 raise HTTPException(
332 status_code=403,
333 detail={"error": CommonProxyErrors.not_allowed_access.value},
334 )
336 try:
337 await _set_vantage_settings(
338 api_key=request.api_key,
339 integration_token=request.integration_token,
340 base_url=request.base_url,
341 )
343 verbose_proxy_logger.info("Vantage settings initialized successfully")
345 return VantageInitResponse(message="Vantage settings initialized successfully", status="success")
347 except HTTPException:
348 raise
349 except Exception as e:
350 verbose_proxy_logger.error("Error initializing Vantage settings: %s", e)
351 raise HTTPException(
352 status_code=500,
353 detail={"error": f"Failed to initialize Vantage settings: {e}"},
354 )
357@router.post(
358 "/vantage/dry-run",
359 tags=["Vantage"],
360 dependencies=[Depends(user_api_key_auth)],
361 response_model=VantageExportResponse,
362)
363async def vantage_dry_run_export(
364 request: VantageDryRunRequest,
365 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
366):
367 """
368 Perform a dry run export using the Vantage logger.
370 Returns the data that would be exported without actually sending it to Vantage.
372 Parameters:
373 - limit: Limit on number of records to preview (default: 500)
375 Only admin users can perform Vantage exports.
376 """
377 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 377 ↛ 378line 377 didn't jump to line 378 because the condition on line 377 was never true
378 raise HTTPException(
379 status_code=403,
380 detail={"error": CommonProxyErrors.not_allowed_access.value},
381 )
383 try:
384 # Dry-run uses the FOCUS database + transformer directly,
385 # bypassing the destination so no Vantage credentials are required.
386 from litellm.integrations.focus.database import FocusLiteLLMDatabase
387 from litellm.integrations.focus.export_engine import FocusExportEngine
388 from litellm.integrations.focus.transformer import FocusTransformer
390 database: Final = FocusLiteLLMDatabase()
391 transformer: Final = FocusTransformer()
393 import polars as pl
395 data: Final = await database.get_usage_data(limit=request.limit)
396 normalized: Final = transformer.transform(data)
398 def _to_json_safe_dicts(frame: pl.DataFrame) -> list:
399 """Cast Decimal columns to Float64 so .to_dicts() produces
400 JSON-serializable float values instead of decimal.Decimal."""
401 decimal_cols = [col for col, dtype in zip(frame.columns, frame.dtypes) if isinstance(dtype, pl.Decimal)]
402 if decimal_cols:
403 frame = frame.with_columns([pl.col(c).cast(pl.Float64) for c in decimal_cols])
404 return frame.to_dicts()
406 usage_sample: Final = _to_json_safe_dicts(data.head(min(50, len(data)))) if not data.is_empty() else []
407 normalized_sample: Final = (
408 _to_json_safe_dicts(normalized.head(min(50, len(normalized)))) if not normalized.is_empty() else []
409 )
411 # Use the same pre-transform column names as
412 # FocusExportEngine.dry_run_export_usage_data for consistency.
413 total_spend: Final = FocusExportEngine._sum_column(data, "spend")
414 total_tokens: Final = FocusExportEngine._sum_column(data, "total_tokens")
415 summary: Final = {
416 "total_records": len(normalized),
417 "total_spend": float(total_spend) if total_spend is not None else 0,
418 "total_tokens": float(total_tokens) if total_tokens is not None else 0,
419 "unique_teams": FocusExportEngine._count_unique(data, "team_id"),
420 "unique_models": FocusExportEngine._count_unique(data, "model"),
421 }
423 dry_run_result: Final = {
424 "usage_data": usage_sample,
425 "normalized_data": normalized_sample,
426 "summary": summary,
427 }
429 verbose_proxy_logger.info("Vantage dry run export completed successfully")
431 return VantageExportResponse(
432 message="Vantage dry run export completed successfully.",
433 status="success",
434 dry_run_data=dry_run_result,
435 summary=summary,
436 )
438 except HTTPException:
439 raise
440 except Exception as e:
441 verbose_proxy_logger.error("Error performing Vantage dry run export: %s", e)
442 raise HTTPException(
443 status_code=500,
444 detail={"error": f"Failed to perform Vantage dry run export: {e}"},
445 )
448@router.post(
449 "/vantage/export",
450 tags=["Vantage"],
451 dependencies=[Depends(user_api_key_auth)],
452 response_model=VantageExportResponse,
453)
454async def vantage_export(
455 request: VantageExportRequest,
456 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
457):
458 """
459 Perform an actual export using the Vantage logger.
461 Exports usage data in FOCUS CSV format to the Vantage API.
463 Parameters:
464 - limit: Optional limit on number of records to export
465 - start_time_utc: Optional start time for data export
466 - end_time_utc: Optional end time for data export
468 Only admin users can perform Vantage exports.
469 """
470 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 470 ↛ 471line 470 didn't jump to line 471 because the condition on line 470 was never true
471 raise HTTPException(
472 status_code=403,
473 detail={"error": CommonProxyErrors.not_allowed_access.value},
474 )
476 try:
477 from litellm.integrations.vantage.vantage_logger import VantageLogger
479 # Prefer the already-registered logger to avoid recreating HTTP clients
480 # on every export call.
481 logger = _get_registered_vantage_logger()
482 if logger is None: 482 ↛ 496line 482 didn't jump to line 496 because the condition on line 482 was always true
483 settings: Final = await _get_vantage_settings()
484 if not settings: 484 ↛ 491line 484 didn't jump to line 491 because the condition on line 484 was always true
485 raise HTTPException(
486 status_code=404,
487 detail={
488 "error": "Vantage settings not found. Please initialize settings first using /vantage/init"
489 },
490 )
491 logger = VantageLogger(
492 api_key=settings.get("api_key"),
493 integration_token=settings.get("integration_token"),
494 base_url=settings.get("base_url"),
495 )
496 await logger.export_usage_data(
497 limit=request.limit,
498 start_time_utc=request.start_time_utc,
499 end_time_utc=request.end_time_utc,
500 )
502 verbose_proxy_logger.info("Vantage export completed successfully")
504 return VantageExportResponse(
505 message="Vantage export completed successfully",
506 status="success",
507 dry_run_data=None,
508 summary=None,
509 )
511 except HTTPException:
512 raise
513 except Exception as e:
514 verbose_proxy_logger.error("Error performing Vantage export: %s", e)
515 raise HTTPException(
516 status_code=500,
517 detail={"error": f"Failed to perform Vantage export: {e}"},
518 )
521@router.delete(
522 "/vantage/delete",
523 tags=["Vantage"],
524 dependencies=[Depends(user_api_key_auth)],
525 response_model=VantageInitResponse,
526)
527async def delete_vantage_settings(
528 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
529):
530 """
531 Delete Vantage settings from the database.
533 Only admin users can delete Vantage settings.
534 """
535 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 raise HTTPException(
537 status_code=403,
538 detail={"error": CommonProxyErrors.not_allowed_access.value},
539 )
541 try:
542 from litellm.proxy.proxy_server import prisma_client
544 if prisma_client is None: 544 ↛ 545line 544 didn't jump to line 545 because the condition on line 544 was never true
545 raise HTTPException(
546 status_code=500,
547 detail={"error": CommonProxyErrors.db_not_connected_error.value},
548 )
550 vantage_config: Final = await _config_table(prisma_client).find_first(
551 where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}
552 )
554 if vantage_config is None:
555 raise HTTPException(
556 status_code=404,
557 detail={"error": "Vantage settings not found"},
558 )
560 await ConfigRepository(prisma_client).table.delete(where={"param_name": VANTAGE_SETTINGS_PARAM_NAME})
562 # Deregister in-memory VantageLogger so the scheduler stops firing
563 from litellm.integrations.vantage.vantage_logger import VantageLogger
565 litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, VantageLogger)
567 verbose_proxy_logger.info("Vantage settings deleted successfully")
569 return VantageInitResponse(message="Vantage settings deleted successfully", status="success")
571 except HTTPException:
572 raise
573 except Exception as e:
574 verbose_proxy_logger.error("Error deleting Vantage settings: %s", e)
575 raise HTTPException(
576 status_code=500,
577 detail={"error": f"Failed to delete Vantage settings: {e}"},
578 )