From e07297fa8731676b0de83ae3443ceb2b037e216a Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:27:24 +0530 Subject: [PATCH] Address Greptile round 3 feedback for improved security and consistency - Encrypt integration_token alongside api_key in Vantage settings storage - Align dry-run summary with FocusExportEngine helper methods - Vectorize Tags JSON building using pl.struct + map_elements - Reuse registered VantageLogger in /export endpoint instead of creating fresh instances Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/transformer.py | 62 +++++++++---------- .../proxy/spend_tracking/vantage_endpoints.py | 57 +++++++++++++---- 2 files changed, 77 insertions(+), 42 deletions(-) diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 1e4a17796d..adac158651 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -10,27 +10,34 @@ import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA -def _build_tags_json(row: dict) -> str: - """Build a JSON string of metadata tags from a DB row. +_TAG_KEYS = ( + "team_id", + "team_alias", + "user_id", + "user_email", + "api_key_alias", + "model", + "model_group", + "custom_llm_provider", +) - Vantage uses this for Token Allocation — enriching billing data with - team, user, and API key metadata. + +def _build_tags_expr(available_keys: list[str]) -> pl.Expr: + """Build a Polars expression that produces a JSON Tags string per row. + + Uses ``pl.struct`` + ``map_elements`` so the heavy iteration stays inside + Polars rather than materialising every row to a Python dict first. """ - tags: dict[str, str] = {} - for key in ( - "team_id", - "team_alias", - "user_id", - "user_email", - "api_key_alias", - "model", - "model_group", - "custom_llm_provider", - ): - val = row.get(key) - if val is not None: - tags[key] = str(val) - return json.dumps(tags) if tags else "{}" + + def _struct_to_json(row: dict) -> str: + tags = {k: str(v) for k, v in row.items() if v is not None} + return json.dumps(tags) if tags else "{}" + + return ( + pl.struct(available_keys) + .map_elements(_struct_to_json, return_dtype=pl.String) + .alias("Tags") + ) class FocusTransformer: @@ -43,19 +50,12 @@ class FocusTransformer: if frame.is_empty(): return pl.DataFrame(schema=self.schema) - # Build Tags JSON from metadata columns - tag_col = "Tags" - tag_keys = [ - "team_id", "team_alias", "user_id", "user_email", - "api_key_alias", "model", "model_group", "custom_llm_provider", - ] - available_keys = [k for k in tag_keys if k in frame.columns] + # Build Tags JSON from metadata columns using vectorized Polars expression + available_keys = [k for k in _TAG_KEYS if k in frame.columns] if available_keys: - tags_series = frame.select(available_keys).to_dicts() - tags_json = [_build_tags_json(row) for row in tags_series] - frame = frame.with_columns(pl.Series(tag_col, tags_json)) + frame = frame.with_columns(_build_tags_expr(available_keys)) else: - frame = frame.with_columns(pl.lit("{}").alias(tag_col)) + frame = frame.with_columns(pl.lit("{}").alias("Tags")) # derive period start/end from usage date frame = frame.with_columns( @@ -124,5 +124,5 @@ class FocusTransformer: pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), none_str.alias("SubAccountType"), - pl.col(tag_col).cast(pl.String).alias("Tags"), + pl.col("Tags").cast(pl.String).alias("Tags"), ) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d0e01fe325..a553d1061e 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -1,5 +1,6 @@ import json +import litellm from fastapi import APIRouter, Depends, HTTPException from litellm._logging import verbose_proxy_logger @@ -26,6 +27,18 @@ _sensitive_masker = SensitiveDataMasker() VANTAGE_SETTINGS_PARAM_NAME = "vantage_settings" +def _get_registered_vantage_logger(): + """Return the VantageLogger already registered in litellm.callbacks, if any.""" + from litellm.integrations.vantage.vantage_logger import VantageLogger + + vantage_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=VantageLogger + ) + if vantage_loggers: + return vantage_loggers[0] + return None + + async def _set_vantage_settings( api_key: str, integration_token: str, base_url: str ): @@ -39,10 +52,11 @@ async def _set_vantage_settings( ) encrypted_api_key = encrypt_value_helper(api_key) + encrypted_integration_token = encrypt_value_helper(integration_token) vantage_settings = { "api_key": encrypted_api_key, - "integration_token": integration_token, + "integration_token": encrypted_integration_token, "base_url": base_url, } @@ -95,6 +109,22 @@ async def _get_vantage_settings(): ) settings["api_key"] = decrypted_api_key + encrypted_integration_token = settings.get("integration_token") + if encrypted_integration_token: + decrypted_integration_token = decrypt_value_helper( + encrypted_integration_token, + key="vantage_integration_token", + exception_type="error", + ) + if decrypted_integration_token is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to decrypt Vantage integration token. Check your salt key configuration." + }, + ) + settings["integration_token"] = decrypted_integration_token + return settings @@ -346,6 +376,7 @@ async def vantage_dry_run_export( # Dry-run uses the FOCUS database + transformer directly, # bypassing the destination so no Vantage credentials are required. from litellm.integrations.focus.database import FocusLiteLLMDatabase + from litellm.integrations.focus.export_engine import FocusExportEngine from litellm.integrations.focus.transformer import FocusTransformer database = FocusLiteLLMDatabase() @@ -357,11 +388,12 @@ async def vantage_dry_run_export( usage_sample = data.head(min(50, len(data))).to_dicts() if not data.is_empty() else [] normalized_sample = normalized.head(min(50, len(normalized))).to_dicts() if not normalized.is_empty() else [] + # Use the same column names as FocusExportEngine.dry_run_export_usage_data summary = { "total_records": len(normalized), - "total_spend": float(normalized.select("BilledCost").sum().item()) if not normalized.is_empty() and "BilledCost" in normalized.columns else 0.0, - "unique_teams": normalized["SubAccountId"].n_unique() if not normalized.is_empty() and "SubAccountId" in normalized.columns else 0, - "unique_models": normalized["ResourceType"].n_unique() if not normalized.is_empty() and "ResourceType" in normalized.columns else 0, + "total_spend": FocusExportEngine._sum_column(normalized, "BilledCost"), + "unique_teams": FocusExportEngine._count_unique(normalized, "SubAccountId"), + "unique_models": FocusExportEngine._count_unique(normalized, "ResourceType"), } dry_run_result = { @@ -420,15 +452,18 @@ async def vantage_export( ) try: - settings = await _get_vantage_settings() - from litellm.integrations.vantage.vantage_logger import VantageLogger - logger = VantageLogger( - api_key=settings.get("api_key"), - integration_token=settings.get("integration_token"), - base_url=settings.get("base_url"), - ) + # Prefer the already-registered logger to avoid recreating HTTP clients + # on every export call. + logger = _get_registered_vantage_logger() + if logger is None: + settings = await _get_vantage_settings() + logger = VantageLogger( + api_key=settings.get("api_key"), + integration_token=settings.get("integration_token"), + base_url=settings.get("base_url"), + ) await logger.export_usage_data( limit=request.limit, start_time_utc=request.start_time_utc,