diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 7b9de4644b4..8c331c86eac 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -5172,3 +5172,171 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool ) assert exc_info.value.status_code == 400 assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) + + +_REPLAY_TOOL: Final = { + "type": "function", + "name": "vault_token", + "parameters": {"type": "object", "properties": {"slot": {"type": "integer"}}, "required": ["slot"]}, +} +_REASONING_MARKER: Final = "SECRET-CHAIN-OF-THOUGHT" +_REPLAY_TURNS: Final = [1, 2, 5, 8] + + +def _replay_field(message, key, default=None): + if isinstance(message, dict): + return message.get(key, default) + return getattr(message, key, default) + + +def _replay_tool_call_ids(message) -> tuple: + return tuple(_replay_field(call, "id") for call in _replay_field(message, "tool_calls") or ()) + + +def _replay_visible_text(message): + content = _replay_field(message, "content") + if content is None: + return "" + return content if isinstance(content, str) else json.dumps(content) + + +def _replay_reasoning_item(turn, summary=True, encrypted=False, content_blocks=False) -> dict: + text = f"{_REASONING_MARKER}-{turn}" + return { + "type": "reasoning", + "id": f"rs_{turn}", + "summary": [{"type": "summary_text", "text": text}] if summary else [], + **({"content": [{"type": "reasoning_text", "text": text}]} if content_blocks else {}), + **({"encrypted_content": "gAAAAABopaque=="} if encrypted else {}), + } + + +def _replay_call_item(turn, suffix=""): + return { + "type": "function_call", + "id": f"fc_{turn}{suffix}", + "call_id": f"call_{turn}{suffix}", + "name": "vault_token", + "arguments": json.dumps({"slot": turn}), + "status": "completed", + } + + +def _replay_call_output_item(turn, suffix=""): + return {"type": "function_call_output", "call_id": f"call_{turn}{suffix}", "output": f"tok-{turn}{suffix}"} + + +def _replay_assistant_text_item(turn): + return { + "type": "message", + "id": f"msg_{turn}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": f"Fetching slot {turn}.", "annotations": []}], + } + + +_REPLAY_SHAPES: Final = { + "reasoning_summary": lambda t: [ + _replay_reasoning_item(t), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "reasoning_empty_summary": lambda t: [ + _replay_reasoning_item(t, summary=False), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "reasoning_encrypted_only": lambda t: [ + _replay_reasoning_item(t, summary=False, encrypted=True), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "reasoning_content_blocks": lambda t: [ + _replay_reasoning_item(t, summary=False, content_blocks=True), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "assistant_text_beside_the_call": lambda t: [ + _replay_assistant_text_item(t), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "reasoning_and_assistant_text": lambda t: [ + _replay_reasoning_item(t), + _replay_assistant_text_item(t), + _replay_call_item(t), + _replay_call_output_item(t), + ], + "no_reasoning_item": lambda t: [_replay_call_item(t), _replay_call_output_item(t)], + "user_speaks_again_each_turn": lambda t: [ + _replay_reasoning_item(t), + _replay_call_item(t), + _replay_call_output_item(t), + {"role": "user", "content": f"Now do slot {t + 1}."}, + ], + "parallel_calls": lambda t: [ + _replay_reasoning_item(t), + _replay_call_item(t), + _replay_call_item(t, "b"), + _replay_call_output_item(t), + _replay_call_output_item(t, "b"), + ], +} + + +def _replay_transcript(turns, shape) -> tuple: + opening = ({"role": "user", "content": "Collect vault tokens one at a time, starting at slot 1."},) + return opening + tuple(item for turn in range(1, turns + 1) for item in shape(turn)) + + +def _bridge_replay(transcript) -> dict: + return LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="openai/kimi-k3", + input=list(transcript), + responses_api_request={"tools": [_REPLAY_TOOL], "max_output_tokens": 500}, + custom_llm_provider="openai", + ) + + +def _expected_replay_call_ids(turns, shape_name) -> frozenset: + suffixes = ("", "b") if shape_name == "parallel_calls" else ("",) + return frozenset(f"call_{turn}{suffix}" for turn in range(1, turns + 1) for suffix in suffixes) + + +@pytest.mark.parametrize("shape_name", sorted(_REPLAY_SHAPES)) +@pytest.mark.parametrize("turns", _REPLAY_TURNS) +class TestBridgeMultiTurnReplay: + """A stateless client appends the bridge's own output back onto input every turn""" + + def test_every_replayed_call_reaches_the_provider_as_a_tool_call(self, shape_name, turns): + messages = _bridge_replay(_replay_transcript(turns, _REPLAY_SHAPES[shape_name]))["messages"] + sent = frozenset(call_id for message in messages for call_id in _replay_tool_call_ids(message)) + missing = _expected_replay_call_ids(turns, shape_name) - sent + assert not missing, f"{sorted(missing)} were replayed but reach the provider as no tool call" + + def test_reasoning_never_becomes_visible_content(self, shape_name, turns): + for message in _bridge_replay(_replay_transcript(turns, _REPLAY_SHAPES[shape_name]))["messages"]: + role = _replay_field(message, "role") + assert _REASONING_MARKER not in _replay_visible_text(message), f"reasoning is visible {role} content" + + def test_the_provider_sees_exactly_the_user_turns_the_client_sent(self, shape_name, turns): + transcript = _replay_transcript(turns, _REPLAY_SHAPES[shape_name]) + expected = sum(1 for item in transcript if item.get("role") == "user") + got = sum(1 for m in _bridge_replay(transcript)["messages"] if _replay_field(m, "role") == "user") + assert got == expected, f"{got} user messages reach the provider, but the client sent {expected}" + + def test_no_tool_result_is_orphaned(self, shape_name, turns): + messages = _bridge_replay(_replay_transcript(turns, _REPLAY_SHAPES[shape_name]))["messages"] + for index, message in enumerate(messages): + if _replay_field(message, "role") != "tool": + continue + call_id = _replay_field(message, "tool_call_id") + announced = frozenset( + announced_id for earlier in messages[:index] for announced_id in _replay_tool_call_ids(earlier) + ) + assert call_id in announced, f"tool result {call_id} arrives before any message announces it" + + def test_the_tools_survive_the_whole_transcript(self, shape_name, turns): + request = _bridge_replay(_replay_transcript(turns, _REPLAY_SHAPES[shape_name])) + assert len(request.get("tools") or []) == 1, "the tool definition is gone from the request"