diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 36dffa3795e..ebbf569a356 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -70,7 +70,7 @@ class BaseResponsesAPIConfig(ABC): return False def supports_streamed_retrieval(self) -> bool: - return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE) + return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI) def sign_request( self, diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 27b5924c540..771d68a0ea0 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" -_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_RESPONSES_API_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingProxyType( { "aresponses": "previous_response_id", @@ -104,7 +104,7 @@ def _rewrite_advertised_id( return event -def _is_responses_api_create_route(request_route: str | None) -> bool: +def _streams_response_ids(request_route: str | None) -> bool: if request_route is None: return False canonical: Final = ( @@ -112,7 +112,8 @@ def _is_responses_api_create_route(request_route: str | None) -> bool: if request_route.startswith(_RESPONSES_API_PROVIDER_PREFIX + "/") else request_route ) - return canonical in _RESPONSES_API_CREATE_ROUTES + parent, _, response_id = canonical.rpartition("/") + return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "") class ResponsesIDSecurity(CustomLogger): @@ -339,7 +340,7 @@ class ResponsesIDSecurity(CustomLogger): async for chunk in response: if ( isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and _is_responses_api_create_route(user_api_key_dict.request_route) + and _streams_response_ids(user_api_key_dict.request_route) and not general_settings.get("disable_responses_id_security", False) ): chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index b5ef9e95741..27e256951c8 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -192,6 +192,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_search_requested: bool = has_server_tool_search(request_tools) self._tool_search_call_ids: frozenset[str] = frozenset() + async def aclose(self) -> None: + await self.litellm_custom_stream_wrapper.aclose() + def _is_tool_search_call(self, call_id: str, fn_name: str) -> bool: if call_id in self._tool_search_call_ids: return True diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index ed90acbd2fa..5fde9d6b444 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -702,6 +702,33 @@ async def test_streaming_events_share_the_chat_completion_response_id(): assert decoded["custom_llm_provider"] == "anthropic" +class _ClosableStreamWrapper(_FakeStreamWrapper): + def __init__(self, chunks): + super().__init__(chunks) + self.closed = False + + async def aclose(self): + self.closed = True + + +@pytest.mark.asyncio +async def test_aclose_closes_the_bridged_chat_completion_stream(): + wrapper: Final = _ClosableStreamWrapper([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=wrapper, + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + await iterator.__anext__() + await iterator.aclose() + + assert wrapper.closed is True + + def test_sync_streaming_events_share_the_chat_completion_response_id(): iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) diff --git a/tests/unit/responses/test_responses_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 2a756fcde9d..3e01574c6cc 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -118,7 +118,7 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url( assert str(requests[0].request.url) == expected_url -StreamProvider: TypeAlias = Literal["openai", "azure"] +StreamProvider: TypeAlias = Literal["openai", "azure", "azure_ai"] _STORED_RESPONSE_ID: Final = "resp_background" _STREAM_HEADERS: Final = {"content-type": "text/event-stream"} @@ -168,6 +168,10 @@ def _retrieve_route(mock_router: respx.MockRouter, provider: StreamProvider) -> return mock_router.get(f"https://api.openai.com/v1/responses/{_STORED_RESPONSE_ID}") case "azure": return mock_router.get(f"https://example-resource.openai.azure.com/openai/responses/{_STORED_RESPONSE_ID}") + case "azure_ai": + return mock_router.get( + f"https://example-resource.services.ai.azure.com/openai/v1/responses/{_STORED_RESPONSE_ID}" + ) def _retrieve_kwargs(provider: StreamProvider) -> dict[str, object]: @@ -181,6 +185,12 @@ def _retrieve_kwargs(provider: StreamProvider) -> dict[str, object]: "api_base": "https://example-resource.openai.azure.com", "api_version": "2025-03-01-preview", } + case "azure_ai": + return { + "custom_llm_provider": "azure_ai", + "api_key": "sk-test", + "api_base": "https://example-resource.services.ai.azure.com", + } async def _retrieve(sync_mode: bool, **kwargs: object) -> object: @@ -197,7 +207,7 @@ async def _event_types(stream: object) -> tuple[str, ...]: @pytest.mark.parametrize("sync_mode", (True, False)) -@pytest.mark.parametrize("provider", ("openai", "azure")) +@pytest.mark.parametrize("provider", ("openai", "azure", "azure_ai")) @pytest.mark.parametrize( ("resume_kwargs", "expected_query"), ( diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index c7fd7fe247f..7e74bfd0f17 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -11,7 +11,7 @@ from fastapi import HTTPException from litellm.proxy.hooks.responses_id_security import ( ResponsesIDSecurity, - _is_responses_api_create_route, + _streams_response_ids, ) from litellm.types.llms.openai import ( GenericEvent, @@ -527,7 +527,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() -class TestIsResponsesApiCreateRoute: +class TestStreamsResponseIds: """Test the route gate that decides whether a streamed response id is encrypted.""" @pytest.mark.parametrize( @@ -536,10 +536,13 @@ class TestIsResponsesApiCreateRoute: "/v1/responses", "/responses", "/openai/v1/responses", + "/v1/responses/resp_abc", + "/responses/resp_abc", + "/openai/v1/responses/resp_abc", ], ) - def test_create_routes_match(self, route): - assert _is_responses_api_create_route(route) is True + def test_create_and_retrieve_routes_match(self, route): + assert _streams_response_ids(route) is True @pytest.mark.parametrize( "route", @@ -547,14 +550,14 @@ class TestIsResponsesApiCreateRoute: None, "/chat/completions", "/openai/v1/chat/completions", - "/v1/responses/{response_id}", - "/openai/v1/responses/{response_id}", + "/v1/responses/", + "/v1/responses/resp_abc/input_items", "/v1/responsesX", "/responsesX", ], ) - def test_non_create_routes_do_not_match(self, route): - assert _is_responses_api_create_route(route) is False + def test_other_routes_do_not_match(self, route): + assert _streams_response_ids(route) is False class TestAsyncPostCallStreamingIteratorHook: @@ -609,7 +612,13 @@ class TestAsyncPostCallStreamingIteratorHook: @pytest.mark.asyncio @pytest.mark.parametrize( "route", - ["/v1/responses", "/responses", "/openai/v1/responses"], + [ + "/v1/responses", + "/responses", + "/openai/v1/responses", + "/v1/responses/resp_rawprovider123", + "/openai/v1/responses/resp_rawprovider123", + ], ) async def test_streamed_id_encrypted_on_all_responses_routes( self, responses_id_security, route, monkeypatch