diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py index 09ef1889818..f2e36137a19 100644 --- a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -40,6 +40,58 @@ def _bedrock_mantle_kwargs() -> dict: } +def _openai_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="completion", + model="gpt-5.5", + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": "gpt-5.5", + "custom_llm_provider": "openai", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +def test_completion_forwards_custom_llm_provider_to_responses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + with patch("litellm.responses", return_value=cached) as fake_responses: + result = bridge.completion(**_openai_kwargs()) + + assert result is cached + assert fake_responses.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_custom_llm_provider_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_openai_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_acompletion_forwards_aws_region_name_to_aresponses(): bridge = ResponsesToCompletionBridgeHandler()