diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py index 3baeebcd74a..1f3bded035b 100644 --- a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py +++ b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py @@ -547,7 +547,7 @@ def test_r21_responses_stream_input_item_without_content_after_an_assistant_item assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received (row,) = _success_rows(marker, 1) assert row["call_type"] == "aresponses" and row["end_user"] == marker, row - assert row["request_id"] == _inner_response_id(str(completed["response"]["id"])), (row, completed) + assert row["request_id"] == str(completed["response"]["id"]), (row, completed) def test_r22_passthrough_converse_forwards_a_message_without_content_verbatim(gateway: Gateway) -> None: diff --git a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py index a1432ffd7a0..19aa11a8e36 100644 --- a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py +++ b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py @@ -1,7 +1,6 @@ import asyncio import base64 import json -import os import re import threading import uuid @@ -26,8 +25,6 @@ from integration._support.upstream import _aws_event_frame, aws_event_stream_fra from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with - pytestmark: Final = pytest.mark.shares_worker_state _MODEL_ID: Final = "global.moonshotai.kimi-k3" @@ -43,7 +40,6 @@ _REJECTION: Final = "structured output schema uses unsupported regex negative lo _THROTTLED: Final = "Too many requests, please wait before trying again" _UNKNOWN_TYPE: Final = "somethingBedrockAddedLater" _UNKNOWN_ONLY: Final = "none of its 1 events carried a known event type" -_SIGNING_KEY: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") _JSON_HEADERS: Final = MappingProxyType({":content-type": "application/json", ":message-type": "event"}) _USAGE: Final[dict[str, JsonValue]] = {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}} _RESPONSE: Final = json.dumps( @@ -277,17 +273,9 @@ def _responses_text(events: Iterable[dict[str, JsonValue]]) -> str: ) -def _inner_response_id(identity: str) -> str: - managed: Final = decrypt_if_encrypted_with(identity.removeprefix("resp_"), _SIGNING_KEY) - assert managed is not None, identity - issued: Final = managed.split(";", 1)[0].rsplit("response_id:", 1)[1] - decoded: Final = base64.b64decode(issued.removeprefix("resp_")).decode() - return decoded.rsplit("response_id:", 1)[1] - - def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str: (completed,) = tuple(event for event in events if event.get("type") == "response.completed") - return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"])) + return str(_JSON.validate_python(completed["response"])["id"]) def _only_received(wire: Wire) -> tuple[str, dict[str, JsonValue]]: diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 4441582fb6b..a6714979c8e 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -456,7 +456,7 @@ def _inner_response_id(identity: str) -> str: def _completed_response_id(events: Iterable[dict[str, JsonValue]]) -> str: (completed,) = tuple(event for event in events if event.get("type") == "response.completed") - return _inner_response_id(str(_JSON.validate_python(completed["response"])["id"])) + return str(_JSON.validate_python(completed["response"])["id"]) def _chat_text(body: Mapping[str, JsonValue]) -> str: @@ -911,7 +911,7 @@ async def test_k06_openai_sdk_async_responses_stream_reaches_kimi_without_a_cach received: Final = _only_received(rig.wire) assert received.model == _KIMI_US and received.streaming, received assert received.cache_points == 0, received.body - _success_row(_inner_response_id(completed[0].id)) + _success_row(completed[0].id) _LOCATIONS: Final = ( diff --git a/tests/integration/providers/test_gemini_thinking_replay_chaos.py b/tests/integration/providers/test_gemini_thinking_replay_chaos.py index f8bb52ef60a..fb3092adfc3 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_chaos.py +++ b/tests/integration/providers/test_gemini_thinking_replay_chaos.py @@ -233,6 +233,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None: assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text +def _spend_id(served: _Served) -> str: + if served.call.endpoint == "responses" and served.call.stream: + payloads: Final = ( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in served.text.splitlines() + if line.startswith("data: {") + ) + (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed") + return str(_JSON_OBJECT.validate_python(completed["response"])["id"]) + return _response_id(served.call.marker) + + async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signature_free(gateway: Gateway) -> None: calls: Final = _calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0) with wire_server(_echo) as wire, gateway.scenario() as scenario: @@ -244,7 +256,7 @@ async def test_concurrent_thinking_replays_across_endpoints_reach_gemini_signatu _assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) landed: Final = _spend_request_ids(model, 30) assert landed >= { - _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream) + _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream) } diff --git a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py index ead83ed66aa..c185c968be9 100644 --- a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py +++ b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py @@ -301,6 +301,18 @@ def _assert_answered_with_its_own_marker(served: _Served) -> None: assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text +def _spend_id(served: _Served) -> str: + if served.call.endpoint == "responses" and served.call.stream: + payloads: Final = ( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in served.text.splitlines() + if line.startswith("data: {") + ) + (completed,) = tuple(event for event in payloads if event.get("type") == "response.completed") + return str(_JSON_OBJECT.validate_python(completed["response"])["id"]) + return _response_id(served.call.marker) + + async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_its_own_form(gateway: Gateway) -> None: calls: Final = _calls( 30, ("chat", "messages", "responses"), lambda index: index % 2 == 0, lambda index: (index // 2) % 2 == 0 @@ -314,7 +326,7 @@ async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_ _assert_each_in_its_own_form(wire.drain(), calls) landed: Final = _spend_rows_by_model(model, 30) assert landed >= { - _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream) + _spend_id(item) for item in served if not (item.call.endpoint == "messages" and item.call.stream) }