diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4b3e4cfd6fb..1f78a462d54 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -277,6 +277,7 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.spend_log_insert import RejectedSpendLogRow from litellm.repositories.table_repositories import ( EndUserRepository, HealthCheckRepository, @@ -9174,7 +9175,10 @@ async def _keep_charges_of_dropped_logs(repo: SpendLogsRepository, rows: Sequenc try: await repo.insert_with_deliveries(charged) except Exception as e: # noqa: BLE001 # the request log is already lost; keep the flush going for other rows - if not (PrismaDBExceptionHandler.is_rejected_row_error(e) or isinstance(e, (UnicodeEncodeError, ValueError))): + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): raise spend_log_error( "Spend tracking - could not keep the charges of %d dropped spend log rows. error=%s", @@ -9215,7 +9219,10 @@ async def _create_spend_logs_with_poison_isolation( await repo.insert_with_deliveries(rows) return failure_budget except Exception as e: - if not (PrismaDBExceptionHandler.is_rejected_row_error(e) or isinstance(e, (UnicodeEncodeError, ValueError))): + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): raise if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): raise diff --git a/litellm/repositories/spend_log_insert.py b/litellm/repositories/spend_log_insert.py index a72f299ebb0..a958edb9f8a 100644 --- a/litellm/repositories/spend_log_insert.py +++ b/litellm/repositories/spend_log_insert.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import Final -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import LiteralString from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -50,12 +50,19 @@ _COLUMNS: Final[tuple[LiteralString, ...]] = ( _JSON_COLUMNS: Final = frozenset({"metadata", "request_tags", "messages", "response", "proxy_server_request"}) +class RejectedSpendLogRow(ValueError): + """A row the raw insert cannot encode, so it is isolated like a row Postgres refuses.""" + + def _json_value(column: str, value: object) -> object: if isinstance(value, datetime): return value.isoformat() - if column in _JSON_COLUMNS and isinstance(value, str): + if column not in _JSON_COLUMNS or not isinstance(value, str): + return value + try: return _JSON.validate_json(value) - return value + except ValidationError as error: + raise RejectedSpendLogRow(f"Spend log column {column} is not valid JSON") from error def spend_log_insert_statement(rows: Sequence[Mapping[str, object]]) -> tuple[LiteralString, str]: @@ -63,7 +70,7 @@ def spend_log_insert_statement(rows: Sequence[Mapping[str, object]]) -> tuple[Li raise ValueError("Spend log insert requires at least one row") fields: Final = frozenset(rows[0]) | {"created_at", "updated_at"} if not fields.issubset(_COLUMNS) or any(frozenset(row) | {"created_at", "updated_at"} != fields for row in rows): - raise ValueError("Spend log insert requires matching known columns") + raise RejectedSpendLogRow("Spend log insert requires matching known columns") columns: Final = ", ".join('"' + column + '"' for column in _COLUMNS if column in fields) now: Final = datetime.now(timezone.utc).isoformat() values: Final = tuple(