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;