From 7ee492f749f2a573bacbded0c060b0ea719fd431 Mon Sep 17 00:00:00 2001 From: Rick <26716961+Bytechoreographer@users.noreply.github.com> Date: Wed, 24 Jun 2026 19:05:34 +0800 Subject: [PATCH] feat(proxy): read cold-storage prompts back in the logs detail view (#30364) * feat(proxy): read cold-storage prompts back in the logs detail view When a deployment offloads prompts and responses to cold storage instead of Postgres, the spend-log row holds only "{}" placeholders plus a metadata.cold_storage_object_key pointer, so the UI logs detail drawer showed nothing. The detail endpoint only read the placeholder columns and never fetched the object back. Resolve the payload per row based on actual content, not a config flag: if Postgres has content, return it; otherwise read the exact stored object key and fetch from the configured cold storage backend through ColdStorageHandler. Reading the persisted key is a single GET. The key embeds a microsecond timestamp that cannot be reconstructed from the millisecond-precision startTime column, and listing the day's prefix to match on request_id would be too expensive for this per-open path. Also teach the detail drawer's pretty-view parser to accept a bare messages array. The cold storage payload carries the prompt as a top-level messages list with no proxy_server_request, so without this the output rendered while the input stayed blank. ColdStorageHandler gains an optional injected logger so the resolver can be unit tested without monkeypatching. Postgres-stored prompts are unaffected: the fast path returns the existing columns and the request-body object still renders the same way. * Update litellm/proxy/spend_tracking/spend_management_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * test(proxy): cover ColdStorageHandler resolution paths and cold-storage fetch failure Add unit tests for ColdStorageHandler (injected logger, graceful None when no logger is configured, and resolution of a configured logger from the callback registry) and a regression test asserting a cold storage backend exception degrades to the Postgres values instead of surfacing a 500. --------- Co-authored-by: Bytechoreographer Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../spend_tracking/cold_storage_handler.py | 37 ++- .../spend_management_endpoints.py | 121 ++++++- .../test_spend_management_endpoints.py | 295 ++++++++++++++++++ .../PrettyMessagesView.test.tsx | 11 + .../LogDetailsDrawer/prettyMessagesUtils.ts | 24 +- 5 files changed, 447 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py index 57c41bafccd..3974d4df618 100644 --- a/litellm/proxy/spend_tracking/cold_storage_handler.py +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -16,8 +16,14 @@ class ColdStorageHandler: This class is responsible for handling Getting/Setting the proxy server request from cold storage. It allows fetching a dict of the proxy server request from s3 or GCS bucket. + + The cold storage logger can be injected for testing; when omitted it is + resolved from the configured ``litellm.cold_storage_custom_logger``. """ + def __init__(self, cold_storage_logger: Optional[CustomLogger] = None): + self._injected_cold_storage_logger = cold_storage_logger + async def get_proxy_server_request_from_cold_storage_with_object_key( self, object_key: str, @@ -31,33 +37,26 @@ class ColdStorageHandler: Returns: Optional[dict]: The proxy server request dict or None if not found """ - - # select the custom logger to use for cold storage - custom_logger_name: Optional[_custom_logger_compatible_callbacks_literal] = ( - self._select_custom_logger_for_cold_storage() + custom_logger = ( + self._injected_cold_storage_logger or self._resolve_cold_storage_logger() ) - - # if no custom logger name is configured, return None - if custom_logger_name is None: + if custom_logger is None: return None - # get the active/initialized custom logger - custom_logger: Optional[CustomLogger] = ( + return await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=object_key, + ) + + def _resolve_cold_storage_logger(self) -> Optional[CustomLogger]: + custom_logger_name = self._select_custom_logger_for_cold_storage() + if custom_logger_name is None: + return None + return ( litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( custom_logger_name ) ) - # if no custom logger is found, return None - if custom_logger is None: - return None - - proxy_server_request = await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key( - object_key=object_key, - ) - - return proxy_server_request - def _select_custom_logger_for_cold_storage( self, ) -> Optional[_custom_logger_compatible_callbacks_literal]: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0ba77dcd2f0..48f12d44370 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,7 +3,17 @@ import collections import json import os from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + NamedTuple, + Optional, + Union, +) import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status @@ -29,6 +39,7 @@ from litellm.repositories.verification_token_repository import ( if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler else: PrismaClient = Any @@ -2175,6 +2186,89 @@ async def ui_view_spend_logs( raise handle_exception_on_proxy(e) +class RequestResponsePayload(NamedTuple): + messages: Optional[Union[str, list, dict]] + response: Optional[Union[str, list, dict]] + proxy_server_request: Optional[Union[str, dict]] + + +_EMPTY_SPEND_LOG_VALUES = frozenset({"", "{}", "[]", "null"}) + + +def _spend_log_field_has_content(value: Optional[Union[str, list, dict]]) -> bool: + if value is None: + return False + if isinstance(value, str): + return value.strip() not in _EMPTY_SPEND_LOG_VALUES + if isinstance(value, (list, dict)): + return len(value) > 0 + return True + + +def _cold_storage_object_key_from_metadata( + metadata: Optional[Union[str, dict]], +) -> Optional[str]: + if isinstance(metadata, str): + try: + metadata = json.loads(metadata) + except (json.JSONDecodeError, TypeError): + return None + if not isinstance(metadata, dict): + return None + object_key = metadata.get("cold_storage_object_key") + return object_key if isinstance(object_key, str) and object_key else None + + +async def _resolve_request_response_payload( + row: Mapping[str, Any], + cold_storage_handler: "ColdStorageHandler", +) -> RequestResponsePayload: + """ + Decide where the prompt/response come from for a single spend-log row. + + PG holds the content when ``store_prompts_in_spend_logs`` is on; otherwise it + holds ``"{}"`` placeholders and the real payload lives in cold storage keyed + by ``metadata.cold_storage_object_key``. The choice is made on actual row + content, not config flags, so historical and mixed-storage rows both resolve + correctly. + """ + messages = row.get("messages") + response = row.get("response") + proxy_server_request = row.get("proxy_server_request") + + pg_payload = RequestResponsePayload(messages, response, proxy_server_request) + if ( + _spend_log_field_has_content(messages) + or _spend_log_field_has_content(response) + or _spend_log_field_has_content(proxy_server_request) + ): + return pg_payload + + object_key = _cold_storage_object_key_from_metadata(row.get("metadata")) + if object_key is None: + return pg_payload + + try: + payload = await cold_storage_handler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=object_key + ) + except Exception: + verbose_proxy_logger.warning( + "Failed to fetch cold storage payload for key %s; falling back to DB values", + object_key, + exc_info=True, + ) + return pg_payload + if payload is None: + return pg_payload + + return RequestResponsePayload( + messages=payload.get("messages"), + response=payload.get("response"), + proxy_server_request=payload.get("proxy_server_request"), + ) + + @router.get( "/spend/logs/ui/{request_id}", tags=["Budget & Spend Tracking"], @@ -2241,26 +2335,27 @@ async def ui_view_request_response_for_request_id( if payload is not None: return payload - # Fallback: fetch heavy columns directly from the database. - # The list endpoint (/spend/logs/ui) intentionally excludes messages, - # response, and proxy_server_request for performance. When no custom - # logger (S3, GCS, etc.) is configured, we still need to serve these - # fields from the DB for the detail/drawer view. + # Fallback: the list endpoint omits the heavy columns for performance, so + # serve them here. When prompts were offloaded to cold storage the DB holds + # only placeholders, so _resolve_request_response_payload fetches the real + # payload from the configured cold storage backend by object key. if prisma_client is not None: + from litellm.proxy.spend_tracking.cold_storage_handler import ( + ColdStorageHandler, + ) + sql_query = """ - SELECT messages, response, proxy_server_request + SELECT messages, response, proxy_server_request, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = $1 LIMIT 1 """ db_result = await prisma_client.db.query_raw(sql_query, request_id) if db_result and len(db_result) > 0: - row = db_result[0] - return { - "messages": row.get("messages"), - "response": row.get("response"), - "proxy_server_request": row.get("proxy_server_request"), - } + resolved = await _resolve_request_response_payload( + db_result[0], cold_storage_handler=ColdStorageHandler() + ) + return resolved._asdict() return None diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0b583129591..b9716c22cee 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3786,3 +3786,298 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict assert body["data"][0]["metadata"] == {} finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _FakeColdStorageLogger: + """Injectable cold storage logger that records the object key it was asked for.""" + + def __init__(self, payload): + self._payload = payload + self.requested_object_keys = [] + + async def get_proxy_server_request_from_cold_storage_with_object_key( + self, object_key + ): + self.requested_object_keys.append(object_key) + return self._payload + + +def _cold_storage_handler(payload): + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler + + logger = _FakeColdStorageLogger(payload) + return ColdStorageHandler(cold_storage_logger=logger), logger + + +@pytest.mark.parametrize( + "value, expected", + [ + (None, False), + ("", False), + (" ", False), + ("{}", False), + ("[]", False), + ("null", False), + ('{"a": 1}', True), + ({}, False), + ({"a": 1}, True), + ([], False), + ([1], True), + (5, True), + ], +) +def test_spend_log_field_has_content(value, expected): + assert spend_management_endpoints._spend_log_field_has_content(value) is expected + + +@pytest.mark.parametrize( + "metadata, expected", + [ + (None, None), + ("{}", None), + ("not-json", None), + ({"cold_storage_object_key": ""}, None), + ({"cold_storage_object_key": "k/req-1.json"}, "k/req-1.json"), + ('{"cold_storage_object_key": "k/req-2.json"}', "k/req-2.json"), + ], +) +def test_cold_storage_object_key_from_metadata(metadata, expected): + assert ( + spend_management_endpoints._cold_storage_object_key_from_metadata(metadata) + == expected + ) + + +@pytest.mark.asyncio +async def test_resolve_payload_prefers_pg_and_skips_cold_storage(): + handler, logger = _cold_storage_handler({"messages": "X", "response": "Y"}) + row = { + "messages": "{}", + "response": '{"choices": [{"message": {"content": "hi"}}]}', + "proxy_server_request": "{}", + "metadata": {"cold_storage_object_key": "k/req.json"}, + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert resolved.response == '{"choices": [{"message": {"content": "hi"}}]}' + assert logger.requested_object_keys == [] + + +@pytest.mark.asyncio +async def test_resolve_payload_fetches_from_cold_storage_when_pg_empty(): + cold_payload = { + "messages": [{"role": "user", "content": "what is 2+2"}], + "response": {"choices": [{"message": {"content": "4"}}]}, + "proxy_server_request": {"body": {"model": "gpt-4o-mini"}}, + } + handler, logger = _cold_storage_handler(cold_payload) + row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": {"cold_storage_object_key": "llm-gateway/prod/req-42.json"}, + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert logger.requested_object_keys == ["llm-gateway/prod/req-42.json"] + assert resolved.messages == cold_payload["messages"] + assert resolved.response == cold_payload["response"] + assert resolved.proxy_server_request == cold_payload["proxy_server_request"] + + +@pytest.mark.asyncio +async def test_resolve_payload_metadata_as_json_string(): + cold_payload = {"messages": "in", "response": "out", "proxy_server_request": None} + handler, logger = _cold_storage_handler(cold_payload) + row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": json.dumps({"cold_storage_object_key": "k/str-meta.json"}), + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert logger.requested_object_keys == ["k/str-meta.json"] + assert resolved.response == "out" + + +@pytest.mark.asyncio +async def test_resolve_payload_no_object_key_returns_empty_without_fetch(): + handler, logger = _cold_storage_handler({"messages": "should-not-be-used"}) + row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": {}, + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert logger.requested_object_keys == [] + assert resolved == spend_management_endpoints.RequestResponsePayload( + "{}", "{}", "{}" + ) + + +@pytest.mark.asyncio +async def test_resolve_payload_cold_storage_miss_falls_back_to_pg_values(): + handler, logger = _cold_storage_handler(None) + row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": {"cold_storage_object_key": "k/missing.json"}, + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert logger.requested_object_keys == ["k/missing.json"] + assert resolved == spend_management_endpoints.RequestResponsePayload( + "{}", "{}", "{}" + ) + + +@pytest.mark.asyncio +async def test_resolve_payload_cold_storage_exception_falls_back_to_pg_values(): + """A backend error during fetch degrades to PG values instead of bubbling a 500.""" + + class _RaisingLogger: + async def get_proxy_server_request_from_cold_storage_with_object_key( + self, object_key + ): + raise RuntimeError("cold storage backend unavailable") + + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler + + handler = ColdStorageHandler(cold_storage_logger=_RaisingLogger()) + row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": {"cold_storage_object_key": "k/boom.json"}, + } + + resolved = await spend_management_endpoints._resolve_request_response_payload( + row, cold_storage_handler=handler + ) + + assert resolved == spend_management_endpoints.RequestResponsePayload( + "{}", "{}", "{}" + ) + + +@pytest.mark.asyncio +async def test_cold_storage_handler_uses_injected_logger(): + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler + + logger = _FakeColdStorageLogger({"messages": "in", "response": "out"}) + handler = ColdStorageHandler(cold_storage_logger=logger) + + result = await handler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key="k/req.json" + ) + + assert result == {"messages": "in", "response": "out"} + assert logger.requested_object_keys == ["k/req.json"] + + +@pytest.mark.asyncio +async def test_cold_storage_handler_returns_none_when_no_logger_configured(monkeypatch): + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler + + monkeypatch.setattr(litellm, "cold_storage_custom_logger", None, raising=False) + handler = ColdStorageHandler() + + result = await handler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key="k/req.json" + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_cold_storage_handler_resolves_configured_logger_from_registry(monkeypatch): + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler + + logger = _FakeColdStorageLogger({"messages": "from-registry"}) + monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2", raising=False) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_custom_logger_for_callback_name", + lambda name: logger if name == "s3_v2" else None, + ) + handler = ColdStorageHandler() + + result = await handler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key="k/req.json" + ) + + assert result == {"messages": "from-registry"} + assert logger.requested_object_keys == ["k/req.json"] + + +def test_ui_view_request_response_reads_from_cold_storage(client, monkeypatch): + """End-to-end: a placeholder row with a cold_storage_object_key is served from + cold storage through the detail endpoint.""" + from types import SimpleNamespace + + placeholder_row = { + "messages": "{}", + "response": "{}", + "proxy_server_request": "{}", + "metadata": {"cold_storage_object_key": "k/cold.json"}, + } + + async def _query_raw(_sql, *_args): + return [placeholder_row] + + fake_prisma = SimpleNamespace(db=SimpleNamespace(query_raw=_query_raw)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) + + cold_logger = _FakeColdStorageLogger( + { + "messages": [{"role": "user", "content": "hi"}], + "response": {"choices": [{"message": {"content": "hello"}}]}, + "proxy_server_request": None, + } + ) + monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2", raising=False) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [], + ) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_custom_logger_for_callback_name", + lambda name: cold_logger if name == "s3_v2" else None, + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_1" + ) + try: + response = client.get( + "/spend/logs/ui/req-cold", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + body = response.json() + assert body["messages"] == [{"role": "user", "content": "hi"}] + assert body["response"] == {"choices": [{"message": {"content": "hello"}}]} + assert cold_logger.requested_object_keys == ["k/cold.json"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/PrettyMessagesView.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/PrettyMessagesView.test.tsx index b73dcafcdc3..e7295ed7a72 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/PrettyMessagesView.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/PrettyMessagesView.test.tsx @@ -27,6 +27,17 @@ describe("PrettyMessagesView", () => { expect(screen.getByText("Hi there!")).toBeInTheDocument(); }); + it("renders input when request is a bare messages array (cold storage payload)", () => { + const request = [{ role: "user", content: "Write me a poem" }]; + const response = { + choices: [{ message: { role: "assistant", content: "A quiet moment." } }], + }; + + render(); + expect(screen.getByText("Write me a poem")).toBeInTheDocument(); + expect(screen.getByText("A quiet moment.")).toBeInTheDocument(); + }); + it("should render the realtime pretty view for realtime API responses", () => { const request = {}; const response = { diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts index 32ae294b1ee..09b8f551c1d 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts @@ -39,18 +39,24 @@ export const ROLE_STYLES: Record = { * Parse request messages and response message from log data */ export const parseMessages = (request: any, response: any): ParsedMessages => { - // Parse request messages + // Parse request messages. `request` is either the raw request body + // ({ messages: [...] }) or, when prompts come from cold storage, the bare + // messages array itself. const requestMessages: ParsedMessage[] = []; - if (request?.messages && Array.isArray(request.messages)) { - request.messages.forEach((msg: any) => { - requestMessages.push({ - role: msg.role || "user", - content: parseMessageContent(msg.content), - toolCallId: msg.tool_call_id, - }); + const requestMessageList = Array.isArray(request) + ? request + : Array.isArray(request?.messages) + ? request.messages + : []; + + requestMessageList.forEach((msg: any) => { + requestMessages.push({ + role: msg.role || "user", + content: parseMessageContent(msg.content), + toolCallId: msg.tool_call_id, }); - } + }); // Parse response message let responseMessage: ParsedMessage | null = null;