mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
test: cover generic HTTP streaming provider header forwarding
Add sync and async regression tests for the BaseLLMHTTPHandler streaming path, which forwards provider response headers for the ~30 providers that ride the generic handler and had no coverage. Also drop redundant setup prose from the moonshot invoke test docstring.
This commit is contained in:
parent
726db1a4c1
commit
722c650bfd
2 changed files with 76 additions and 4 deletions
|
|
@ -209,10 +209,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest):
|
|||
endpoint with the messages body. Iteration of the stream itself is
|
||||
not exercised here — moonshot streaming delegates to the OpenAI
|
||||
parser and is covered by the OpenAI test suite.
|
||||
|
||||
Patch ``make_sync_call`` at its import site in
|
||||
``base_invoke_transformation`` so we observe the exact kwargs it is
|
||||
called with at stream-wrapper construction time.
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
|
|
|
|||
|
|
@ -2071,3 +2071,79 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
|
|||
retry_authorization = posts[1]["headers"]["Authorization"]
|
||||
assert retry_authorization.startswith("AWS4-HMAC-SHA256")
|
||||
assert retry_authorization != first_attempt_headers["Authorization"]
|
||||
|
||||
|
||||
_GENERIC_STREAM_SSE = (
|
||||
b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,'
|
||||
b'"model":"test-model","choices":[{"index":0,"delta":{"content":"hi"},'
|
||||
b'"finish_reason":null}]}\n\n'
|
||||
b"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
|
||||
def _generic_stream_upstream_response() -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={
|
||||
"x-request-id": "generic-req-123",
|
||||
"x-ratelimit-remaining-requests": "42",
|
||||
},
|
||||
content=_GENERIC_STREAM_SSE,
|
||||
request=httpx.Request("POST", "https://fake-vllm.test/v1/chat/completions"),
|
||||
)
|
||||
|
||||
|
||||
def test_generic_http_handler_sync_streaming_forwards_provider_response_headers():
|
||||
"""
|
||||
Regression test for the generic BaseLLMHTTPHandler streaming path used by
|
||||
~30 providers (deepseek, groq, hosted_vllm, databricks, openrouter, ...).
|
||||
|
||||
The sync `completion()` streaming branch builds the CustomStreamWrapper from
|
||||
`make_sync_call`, which returns the upstream response headers alongside the
|
||||
stream. Those headers must reach the caller as `llm_provider-*` entries in
|
||||
`_hidden_params["additional_headers"]`, which is what the proxy merges into
|
||||
the client-facing response headers.
|
||||
"""
|
||||
mock_client = Mock(spec=HTTPHandler)
|
||||
mock_client.post = Mock(return_value=_generic_stream_upstream_response())
|
||||
|
||||
response = litellm.completion(
|
||||
model="hosted_vllm/test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_base="https://fake-vllm.test/v1",
|
||||
api_key="sk-test",
|
||||
stream=True,
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
additional_headers = response._hidden_params["additional_headers"]
|
||||
assert additional_headers["llm_provider-x-request-id"] == "generic-req-123"
|
||||
assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42"
|
||||
|
||||
assert "".join([chunk.choices[0].delta.content or "" for chunk in response]) == "hi"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_http_handler_async_streaming_forwards_provider_response_headers():
|
||||
"""
|
||||
Companion to the sync test above for `acompletion_stream_function`, which
|
||||
builds its CustomStreamWrapper from `make_async_call_stream_helper`.
|
||||
"""
|
||||
mock_client = AsyncMock(spec=AsyncHTTPHandler)
|
||||
mock_client.post = AsyncMock(return_value=_generic_stream_upstream_response())
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="hosted_vllm/test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_base="https://fake-vllm.test/v1",
|
||||
api_key="sk-test",
|
||||
stream=True,
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
additional_headers = response._hidden_params["additional_headers"]
|
||||
assert additional_headers["llm_provider-x-request-id"] == "generic-req-123"
|
||||
assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42"
|
||||
|
||||
collected = [chunk async for chunk in response]
|
||||
assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue