scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-29 01:42:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge 8ecd31d76c into 2dccc0dc79

Browse source
This commit is contained in:
IvanShang 2026-09-23 14:41:30 +00:00 • committed by GitHub
commit cd295525a8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 9 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -891,7 +891,9 @@ class HeadroomGuardrail(CustomGuardrail):
if not has_headroom_retrieve_tool(effective.get("tools")):
return base_result
return { # mutable-ok: the hook contract is a plain dict the router merges into the request kwargs
**effective,
**{ # mutable-ok: removes stream-only options
key: value for key, value in effective.items() if key != "stream_options"
},
"stream": False,
HEADROOM_CONVERTED_STREAM_KEY: True,
}

7
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
View file

@ -2155,7 +2155,7 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple
tools: Optional[list],
expect_conversion: bool,
):
kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools}
kwargs = {"model": "gpt-4o", "stream": stream, "stream_options": {"include_usage": True}, "tools": tools}
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type)
@ -2163,12 +2163,15 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple
assert result is kwargs
assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs
assert kwargs["stream"] is stream
assert kwargs["stream_options"] == {"include_usage": True}
return
assert result is not None
assert result["stream"] is False
assert "stream_options" not in result
assert result[HEADROOM_CONVERTED_STREAM_KEY] is True
assert kwargs["stream"] is True
assert kwargs["stream_options"] == {"include_usage": True}
@pytest.mark.asyncio
@ -2223,6 +2226,7 @@ async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_c
"model": "gpt-4o",
"messages": [dict(m) for m in ORIGINAL_MESSAGES],
"stream": True,
"stream_options": {"include_usage": True},
"guardrails": ["headroom"],
"metadata": {},
}
@ -2238,6 +2242,7 @@ async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_c
assert result is not None
assert has_headroom_retrieve_tool(result["tools"])
assert result["stream"] is False
assert "stream_options" not in result
assert result[HEADROOM_CONVERTED_STREAM_KEY] is True