From 0eb284b3d14182db2979da9e5e0e09baf97c5c06 Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 6 Oct 2026 21:43:28 +0000 Subject: [PATCH] fix(spend_logs): filter metadata at DB write so guardrail usage sees full rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 9 +--- litellm/proxy/utils.py | 13 +++-- .../spend/test_spend_logs_metadata_fields.py | 53 +++++++++++++++++++ 3 files changed, 64 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b4e85b7b82d..721caecb7cd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1518,16 +1518,9 @@ class DBSpendUpdateWriter: ) ) if prisma_client is not None and spend_logs_url is not None or prisma_client is not None: - from litellm.proxy.spend_tracking.spend_tracking_utils import ( - configured_spend_logs_metadata_fields, - spend_log_row_with_retained_metadata, - ) from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush - await enqueue_spend_logs( - prisma_client, - (spend_log_row_with_retained_metadata(payload, configured_spend_logs_metadata_fields()),), - ) + await enqueue_spend_logs(prisma_client, (payload,)) if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: request_spend_log_flush(prisma_client) else: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 29f2f46f001..4a30c7eabf7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7391,6 +7391,13 @@ class ProxyUpdateSpend: "Spend tracking - processing %d spend logs for DB write", len(logs_to_process), ) + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + configured_spend_logs_metadata_fields, + spend_log_row_with_retained_metadata, + ) + + retention: Final = configured_spend_logs_metadata_fields() + rows_to_write: Final = [spend_log_row_with_retained_metadata(row, retention) for row in logs_to_process] start_time: Final = time.time() try: for i in range(n_retry_times + 1): @@ -7400,7 +7407,7 @@ class ProxyUpdateSpend: if not base_url.endswith("/"): base_url += "/" verbose_proxy_logger.debug("base_url: %s", base_url) - json_data = json.dumps(logs_to_process) + json_data = json.dumps(rows_to_write) response = await db_writer_client.post( url=base_url + "spend/update", data=json_data, @@ -7411,8 +7418,8 @@ class ProxyUpdateSpend: # Items already removed from queue at start of function pass else: - for j in range(0, len(logs_to_process), BATCH_SIZE): - batch = logs_to_process[j : j + BATCH_SIZE] + for j in range(0, len(rows_to_write), BATCH_SIZE): + batch = rows_to_write[j : j + BATCH_SIZE] batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch] isolation_budget = MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH for statement_rows in spend_log_write_batches( diff --git a/tests/integration/spend/test_spend_logs_metadata_fields.py b/tests/integration/spend/test_spend_logs_metadata_fields.py index d3f75abda42..800887a354e 100644 --- a/tests/integration/spend/test_spend_logs_metadata_fields.py +++ b/tests/integration/spend/test_spend_logs_metadata_fields.py @@ -27,6 +27,7 @@ def _config( *, model_list: tuple[Mapping[str, JsonValue], ...] = (), litellm_settings: Mapping[str, JsonValue] | None = None, + guardrails: tuple[Mapping[str, JsonValue], ...] = (), ) -> Path: retention: Final = ( {} if spend_logs_metadata_fields is None else {"spend_logs_metadata_fields": dict(spend_logs_metadata_fields)} @@ -45,6 +46,7 @@ def _config( **retention, }, "litellm_settings": dict(litellm_settings or {}), + "guardrails": [dict(guardrail) for guardrail in guardrails], } ) ) @@ -126,6 +128,57 @@ def test_included_metadata_fields_are_the_only_ones_stored_besides_always_kept( assert set(object_value(row["metadata"])) == {"status", "cold_storage_object_key", "user_api_key_alias"} +def test_guardrail_usage_is_tracked_when_guardrail_information_is_not_stored(gateway: Gateway, tmp_path: Path) -> None: + guardrail_name: Final = f"metadata-fields-guardrail-{uuid4()}" + with ( + wire_server(lambda _: Reply(body=b'{"action":"NONE"}')) as policy, + wire_server(_respond) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=_config( + tmp_path, + {"exclude": ["guardrail_information"]}, + guardrails=( + { + "guardrail_name": guardrail_name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": policy.url, + "api_key": "synthetic-guardrail-key", + }, + }, + ), + ), + ) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response_id: Final = string_value(isolated.chat(model, key=scenario.key(models=[model]))["id"]) + assert len(policy.drain()) == 1 + rows: Final = eventually( + lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert "guardrail_information" not in object_value(rows[0]["metadata"]), rows[0] + indexed: Final = eventually( + lambda: read_rows( + 'SELECT guardrail_id FROM "LiteLLM_SpendLogGuardrailIndex" WHERE request_id=%s', (response_id,) + ), + lambda values: len(values) == 1, + seconds=70, + ) + metrics: Final = read_rows( + 'SELECT requests_evaluated, passed_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (string_value(indexed[0]["guardrail_id"]),), + ) + assert metrics == [{"requests_evaluated": 1, "passed_count": 1}], metrics + + EXCLUDED: Final = ("model_map_information", "user_api_key_alias")