From 3d71adf10f3ccf86bcbae341e87d23b0b550f57c Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 3 Nov 2026 22:48:07 +0000 Subject: [PATCH] fix(spend-logs): trim classifier rows and ignore continuation keys off /v1/responses in latest_turn mode --- .../spend_tracking/spend_tracking_utils.py | 29 ++++--- .../test_spend_tracking_utils.py | 83 ++++++++++++++----- 2 files changed, 82 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a5bef26e3f7..71d1f724b9d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1262,7 +1262,7 @@ _TOOL_INPUT_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use", "mcp_ _TOOL_OUTPUT_BLOCK_TYPES: Final = frozenset({"tool_result", "mcp_tool_result", "function_call_output"}) _LATEST_TURN_STORAGE: Final = "latest_turn" _STATEFUL_REQUEST_KEYS: Final = frozenset({"previous_response_id", "conversation"}) -_LATEST_TURN_KEPT_KEYS: Final = frozenset({"model", "classifier_input"}) +_LATEST_TURN_KEPT_KEYS: Final = frozenset({"model"}) _RESPONSES_CALL_TYPES: Final = frozenset({"responses", "aresponses"}) _REQUEST_BODY_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) _TURNS_KEYS: Final = ("messages", "input") @@ -1328,8 +1328,9 @@ def _latest_user_message_text(turns: object) -> str | None: return _user_content_text(message_field(latest, "content")) -def _continues_stored_conversation(request_body: Mapping[str, object]) -> bool: - return any(request_body.get(key) for key in _STATEFUL_REQUEST_KEYS) +def _continues_stored_conversation(request_body: Mapping[str, object], call_type: object) -> bool: + """Only `/v1/responses` resumes a stored conversation, so the same keys on a chat body do not exempt it.""" + return call_type in _RESPONSES_CALL_TYPES and any(request_body.get(key) for key in _STATEFUL_REQUEST_KEYS) def _replayed_by_responses_bridge(standard_payload: Mapping[str, object] | None) -> bool: @@ -1343,8 +1344,8 @@ def _replayed_by_responses_bridge(standard_payload: Mapping[str, object] | None) return isinstance(response, Mapping) and "choices" in response -def _trimmed_turn(request_body: Mapping[str, object]) -> Mapping[str, object]: - if _continues_stored_conversation(request_body): +def _trimmed_turn(request_body: Mapping[str, object], call_type: object) -> Mapping[str, object]: + if _continues_stored_conversation(request_body, call_type): return request_body kept: Final = {key: value for key, value in request_body.items() if key in _LATEST_TURN_KEPT_KEYS} turns_key: Final = next((key for key in _TURNS_KEYS if key in request_body), None) @@ -1354,13 +1355,19 @@ def _trimmed_turn(request_body: Mapping[str, object]) -> Mapping[str, object]: return MappingProxyType({**kept, turns_key: [{"role": "user", "content": text}]}) -def _latest_turn_body(request_body: Mapping[str, object]) -> Mapping[str, object]: - trimmed: Final = _trimmed_turn(request_body) +def _latest_turn_body(request_body: Mapping[str, object], call_type: object) -> Mapping[str, object]: + """A classifier row stores the trimmed originating request, since the composed classifier message quotes the caller's history.""" originating: Final = request_body.get("originating_request_masked") - if not isinstance(originating, Mapping) or trimmed is request_body: - return trimmed + if not isinstance(originating, Mapping): + return _trimmed_turn(request_body, call_type) + originating_body: Final = _REQUEST_BODY_ADAPTER.validate_python(originating) + model: Final = {key: value for key, value in request_body.items() if key == "model"} return MappingProxyType( - {**trimmed, "originating_request_masked": _trimmed_turn(_REQUEST_BODY_ADAPTER.validate_python(originating))} + { + **_trimmed_turn(originating_body, None), + **model, + "originating_request_masked": _trimmed_turn(originating_body, None), + } ) @@ -1369,7 +1376,7 @@ def _stored_request_body( ) -> Mapping[str, object]: if not stores_latest_turn_only() or _replayed_by_responses_bridge(standard_payload): return request_body - return _latest_turn_body(request_body) + return _latest_turn_body(request_body, None if standard_payload is None else standard_payload.get("call_type")) def _is_request_body_credential(key: str, value: object) -> bool: diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index d5d9131780c..0b387ca7486 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7281,25 +7281,61 @@ def test_latest_turn_mode_stores_only_the_model_and_the_latest_user_text_of_a_st assert _stored_request(body) == expected -@pytest.mark.parametrize( - "body", - [ - { - "model": "gpt-4o-mini", - "previous_response_id": "resp_earlier", - "instructions": "system prompt", - "input": [{"type": "function_call_output", "call_id": "call_1", "output": "tool out"}], - }, - {"model": "gpt-4o-mini", "conversation": "conv_1", "input": "next question"}, - ], +_CONTINUATION_BODIES: Final[tuple[dict[str, object], ...]] = ( + { + "model": "gpt-4o-mini", + "previous_response_id": "resp_earlier", + "instructions": "system prompt", + "input": [{"type": "function_call_output", "call_id": "call_1", "output": "tool out"}], + }, + {"model": "gpt-4o-mini", "conversation": "conv_1", "input": "next question"}, ) -def test_latest_turn_mode_keeps_the_whole_request_when_it_continues_a_stored_conversation( +_NATIVE_RESPONSES_PAYLOAD: Final = {"call_type": "aresponses", "response": {"object": "response", "output": []}} + + +@pytest.mark.parametrize("body", _CONTINUATION_BODIES) +def test_latest_turn_mode_keeps_the_whole_request_when_a_responses_call_continues_a_stored_conversation( monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object] ) -> None: from litellm.proxy import proxy_server monkeypatch.setattr(proxy_server, "general_settings", _LATEST_TURN_SETTINGS) - assert _stored_request(body) == body + assert _stored_request_with_payload(body, _NATIVE_RESPONSES_PAYLOAD) == body + + +@pytest.mark.parametrize( + ("body", "call_type", "expected"), + [ + ( + { + "model": "gpt-4o-mini", + "previous_response_id": "resp_earlier", + "messages": [*_PRIOR_CHAT_TURNS, {"role": "user", "content": "latest"}], + }, + "acompletion", + {"model": "gpt-4o-mini", "messages": _user_turn("latest")}, + ), + ( + { + "model": "claude", + "conversation": "conv_1", + "system": "system prompt", + "messages": [*_PRIOR_CHAT_TURNS, {"role": "user", "content": "latest"}], + }, + "anthropic_messages", + {"model": "claude", "messages": _user_turn("latest")}, + ), + (_CONTINUATION_BODIES[0], None, {"model": "gpt-4o-mini"}), + ], +) +def test_latest_turn_mode_trims_a_request_whose_continuation_keys_are_not_on_a_responses_call( + monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], call_type: str | None, expected: dict[str, object] +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", _LATEST_TURN_SETTINGS) + standard_payload: Final = {"call_type": call_type, "response": {"choices": [{"message": {"content": "a"}}]}} + assert _stored_request_with_payload(body, standard_payload) == expected _BRIDGED_RESPONSES_BODY: Final = { @@ -7346,16 +7382,24 @@ def test_latest_turn_mode_keeps_the_whole_request_of_a_responses_turn_served_by_ assert _stored_request_with_payload(_BRIDGED_RESPONSES_BODY, standard_payload) == expected -def test_latest_turn_mode_keeps_the_classifier_audit_with_its_originating_request_trimmed( +def test_latest_turn_mode_stores_a_classifier_row_as_the_trimmed_originating_request( monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy import proxy_server originating: Final = { "model": "gpt-4o-mini", - "messages": [*_PRIOR_CHAT_TURNS, {"role": "user", "content": "route me"}], + "messages": [ + {"role": "system", "content": "caller system prompt"}, + *_PRIOR_CHAT_TURNS, + {"role": "user", "content": "route me"}, + ], + } + composed_ask: Final = "Caller system prompt, quoted as task context:\ncaller system prompt\nRecent conversation:\n[1] earlier\nCurrent ask:\nroute me" + classifier_input: Final = { + "model": "classifier", + "messages": [{"role": "system", "content": "rubric"}, {"role": "user", "content": composed_ask}], } - classifier_input: Final = {"model": "classifier", "messages": [{"role": "user", "content": "route me"}]} standard_payload: Final = { "call_type": "acompletion", "classifier_input": classifier_input, @@ -7363,12 +7407,13 @@ def test_latest_turn_mode_keeps_the_classifier_audit_with_its_originating_reques "response": {"choices": [{"message": {"content": "simple"}}]}, } monkeypatch.setattr(proxy_server, "general_settings", _LATEST_TURN_SETTINGS) - assert _stored_request_with_payload(classifier_input, standard_payload) == { + stored: Final = _stored_request_with_payload(classifier_input, standard_payload) + assert stored == { "model": "classifier", - "classifier_input": classifier_input, - "originating_request_masked": {"model": "gpt-4o-mini", "messages": _user_turn("route me")}, "messages": _user_turn("route me"), + "originating_request_masked": {"model": "gpt-4o-mini", "messages": _user_turn("route me")}, } + assert "caller system prompt" not in json.dumps(stored) @pytest.mark.parametrize("value", ["latest_trun", "Latest_Turn", "full", True])