Merge pull request #22049 from BerriAI/litellm_spend_log_key_info

[Fix] Enrich Failure Spend Logs With Key/Team Metadata

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -12,7 +12,7 @@ from litellm.litellm_core_utils.core_helpers import (
 )
 from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.auth.auth_checks import log_db_metrics
+from litellm.proxy.auth.auth_checks import get_key_object, get_team_object, log_db_metrics
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.utils import ProxyUpdateSpend
 from litellm.types.utils import (
@@ -76,6 +76,10 @@ class _ProxyDBLogger(CustomLogger):
             traceback_str=traceback_str,
         )
 
+        _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(
+            metadata=_metadata,
+        )
+
         existing_metadata: dict = request_data.get("metadata", None) or {}
         existing_metadata.update(_metadata)
 
@@ -255,6 +259,72 @@ class _ProxyDBLogger(CustomLogger):
                 "Error in tracking cost callback - %s", str(e)
             )
 
+    @staticmethod
+    async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict:
+        """
+        Enriches failure spend log metadata by looking up the key object (and team object)
+        from cache/DB when key fields are missing.
+
+        This handles two scenarios:
+        1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other
+           fields are null. We look up the full key object to fill in alias, user_id,
+           team_id, etc.
+        2. Post-auth failures (provider errors, rate limits): key fields are populated
+           but team_alias is missing because LiteLLM_VerificationTokenView SQL view
+           doesn't include it. We look up the team object to fill in team_alias.
+        """
+        api_key_hash = metadata.get("user_api_key")
+        if not api_key_hash:
+            return metadata
+
+        from litellm.proxy.proxy_server import (
+            prisma_client,
+            proxy_logging_obj,
+            user_api_key_cache,
+        )
+
+        # Step 1: If key fields are missing, look up the full key object
+        if metadata.get("user_api_key_alias") is None:
+            try:
+                key_obj = await get_key_object(
+                    hashed_token=api_key_hash,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
+                if metadata.get("user_api_key_alias") is None:
+                    metadata["user_api_key_alias"] = key_obj.key_alias
+                if metadata.get("user_api_key_user_id") is None:
+                    metadata["user_api_key_user_id"] = key_obj.user_id
+                if metadata.get("user_api_key_team_id") is None:
+                    metadata["user_api_key_team_id"] = key_obj.team_id
+                if metadata.get("user_api_key_org_id") is None:
+                    metadata["user_api_key_org_id"] = key_obj.org_id
+            except Exception:
+                verbose_proxy_logger.debug(
+                    "Failed to enrich failure metadata with key info for api_key=%s",
+                    api_key_hash,
+                )
+
+        # Step 2: If team_id is known but team_alias is missing, look up the team object
+        team_id = metadata.get("user_api_key_team_id")
+        if team_id and metadata.get("user_api_key_team_alias") is None:
+            try:
+                team_obj = await get_team_object(
+                    team_id=team_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
+                if team_obj.team_alias is not None:
+                    metadata["user_api_key_team_alias"] = team_obj.team_alias
+            except Exception:
+                verbose_proxy_logger.debug(
+                    "Failed to enrich failure metadata with team_alias for team_id=%s",
+                    team_id,
+                )
+        return metadata
+
     @staticmethod
     def _should_track_errors_in_db():
         """

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -169,6 +169,223 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object():
         mock_proxy_logging.failed_tracking_alert.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_with_team_alias():
+    """
+    When team_id is set but team_alias is missing (and key_alias is present),
+    _enrich_failure_metadata_with_key_info should look up the team from cache
+    and populate user_api_key_team_alias.
+    """
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "my-team-alias"
+
+    with patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_team_obj,
+    ):
+        metadata = {
+            "user_api_key": "hashed_key",
+            "user_api_key_alias": "my-key-alias",  # already set
+            "user_api_key_team_id": "test_team_id",
+            "user_api_key_team_alias": None,
+        }
+        result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
+        assert result["user_api_key_team_alias"] == "my-team-alias"
+
+
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_with_full_key_lookup():
+    """
+    When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info
+    should look up the key object from cache/DB and populate alias, user_id, team_id,
+    then look up the team to get team_alias.
+    """
+    mock_key_obj = MagicMock()
+    mock_key_obj.key_alias = "fetched-key-alias"
+    mock_key_obj.user_id = "fetched-user-id"
+    mock_key_obj.team_id = "fetched-team-id"
+    mock_key_obj.org_id = "fetched-org-id"
+
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "fetched-team-alias"
+
+    with patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+        new_callable=AsyncMock,
+        return_value=mock_key_obj,
+    ), patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_team_obj,
+    ):
+        metadata = {
+            "user_api_key": "hashed_key",
+            "user_api_key_alias": None,  # all null - simulates auth error path
+            "user_api_key_user_id": None,
+            "user_api_key_team_id": None,
+            "user_api_key_team_alias": None,
+            "user_api_key_org_id": None,
+        }
+        result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
+        assert result["user_api_key_alias"] == "fetched-key-alias"
+        assert result["user_api_key_user_id"] == "fetched-user-id"
+        assert result["user_api_key_team_id"] == "fetched-team-id"
+        assert result["user_api_key_org_id"] == "fetched-org-id"
+        assert result["user_api_key_team_alias"] == "fetched-team-alias"
+
+
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_skips_when_team_alias_present():
+    """
+    When team_alias is already populated, _enrich_failure_metadata_with_key_info
+    should not perform a team cache lookup.
+    """
+    with patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+        new_callable=AsyncMock,
+    ) as mock_get_key, patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+        new_callable=AsyncMock,
+    ) as mock_get_team:
+        metadata = {
+            "user_api_key": "hashed_key",
+            "user_api_key_alias": "existing-alias",
+            "user_api_key_team_id": "test_team_id",
+            "user_api_key_team_alias": "already-set",
+        }
+        result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
+        assert result["user_api_key_team_alias"] == "already-set"
+        mock_get_key.assert_not_called()
+        mock_get_team.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_skips_when_no_api_key():
+    """
+    When api_key hash is absent, _enrich_failure_metadata_with_key_info should
+    not perform any lookups.
+    """
+    with patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+        new_callable=AsyncMock,
+    ) as mock_get_key:
+        metadata = {
+            "user_api_key": None,
+            "user_api_key_alias": None,
+            "user_api_key_team_id": None,
+            "user_api_key_team_alias": None,
+        }
+        result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
+        mock_get_key.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_async_post_call_failure_hook_enriches_auth_error_metadata():
+    """
+    Simulates a 401 ProxyException (e.g. can_key_call_model). In this case
+    UserAPIKeyAuth is created with only api_key set. The failure hook should
+    look up the key and team from cache/DB to populate all missing fields.
+    """
+    logger = _ProxyDBLogger()
+
+    # This is what auth_exception_handler creates for 401 errors
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="hashed_key",
+        # key_alias, user_id, team_id, team_alias are all None
+    )
+
+    request_data = {
+        "model": "claude-haiku-4-5",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "metadata": {},
+        "litellm_params": {},
+    }
+
+    mock_key_obj = MagicMock()
+    mock_key_obj.key_alias = "my-key-alias"
+    mock_key_obj.user_id = "my-user-id"
+    mock_key_obj.team_id = "my-team-id"
+    mock_key_obj.org_id = None
+
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "my-team-alias"
+
+    with patch(
+        "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
+        new_callable=AsyncMock,
+    ) as mock_update_database, patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+        new_callable=AsyncMock,
+        return_value=mock_key_obj,
+    ), patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_team_obj,
+    ):
+        await logger.async_post_call_failure_hook(
+            request_data=request_data,
+            original_exception=Exception("401 - model not allowed"),
+            user_api_key_dict=user_api_key_dict,
+        )
+
+        mock_update_database.assert_called_once()
+        call_args = mock_update_database.call_args[1]
+        metadata = call_args["kwargs"]["litellm_params"]["metadata"]
+        assert metadata["user_api_key_alias"] == "my-key-alias"
+        assert metadata["user_api_key_user_id"] == "my-user-id"
+        assert metadata["user_api_key_team_id"] == "my-team-id"
+        assert metadata["user_api_key_team_alias"] == "my-team-alias"
+
+
+@pytest.mark.asyncio
+async def test_async_post_call_failure_hook_enriches_missing_team_alias():
+    """
+    When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook
+    should look up the team from cache and populate user_api_key_team_alias in the
+    spend log metadata written to the DB.
+    """
+    logger = _ProxyDBLogger()
+
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test_api_key",
+        key_alias="test_alias",
+        user_id="test_user_id",
+        team_id="test_team_id",
+        team_alias=None,  # Missing - simulates regular key auth where SQL view omits team_alias
+    )
+
+    request_data = {
+        "model": "gpt-4",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "metadata": {},
+        "litellm_params": {},
+    }
+
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "enriched-team-alias"
+
+    with patch(
+        "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
+        new_callable=AsyncMock,
+    ) as mock_update_database, patch(
+        "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_team_obj,
+    ):
+        await logger.async_post_call_failure_hook(
+            request_data=request_data,
+            original_exception=Exception("Provider rate limit"),
+            user_api_key_dict=user_api_key_dict,
+        )
+
+        mock_update_database.assert_called_once()
+        call_args = mock_update_database.call_args[1]
+        metadata = call_args["kwargs"]["litellm_params"]["metadata"]
+        assert metadata["user_api_key_team_alias"] == "enriched-team-alias"
+        assert metadata["user_api_key_team_id"] == "test_team_id"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_value", [None, ""])
 async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value):