mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
test(headroom): cover stream conversion after deployment-level compression
This commit is contained in:
parent
4d4cf40334
commit
2affd800ec
1 changed files with 26 additions and 0 deletions
|
|
@ -2014,6 +2014,32 @@ async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_co
|
|||
assert result["messages"] == EXPECTED_MESSAGES
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_compression(
|
||||
guardrail: HeadroomGuardrail,
|
||||
):
|
||||
kwargs = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [dict(m) for m in ORIGINAL_MESSAGES],
|
||||
"stream": True,
|
||||
"guardrails": ["headroom"],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
|
||||
):
|
||||
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
|
||||
|
||||
assert result is not None
|
||||
assert has_headroom_retrieve_tool(result["tools"])
|
||||
assert result["stream"] is False
|
||||
assert result[HEADROOM_CONVERTED_STREAM_KEY] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end(
|
||||
guardrail: HeadroomGuardrail,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue