scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

test(proxy): type the per-chunk streaming hook test helpers

Browse source
This commit is contained in:
yucheng 2026-10-27 17:16:34 -07:00
parent c05907f867
commit 01ab5bb7a1
3 changed files with 33 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/unit/proxy/test_common_request_processing.py
View file

@ -4973,16 +4973,18 @@ class TestAsyncStreamingDataGeneratorFastPath:
hook_spy.assert_not_awaited()
@pytest.mark.asyncio
async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(self, monkeypatch):
async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")])
ProxyLogging._callback_capabilities_cache.clear()
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
hook_spy = AsyncMock(side_effect=lambda **kw: kw["response"])
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=MagicMock())
hook_spy: Final = AsyncMock(side_effect=lambda **kw: kw["response"])
monkeypatch.setattr(proxy_logging_obj, "async_post_call_streaming_hook", hook_spy)
chunks = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"]
out = [
chunks: Final = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"]
out: Final = [
c
async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=self._aiter(chunks),

2
tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py
View file

@ -40,7 +40,7 @@ class _OverridesPerChunk(CustomLogger):
class _PerChunkGuardrail(CustomGuardrail):
async def async_post_call_streaming_hook(self, *args, **kwargs):
async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
return None

42
tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py
View file

@ -10,7 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from collections.abc import AsyncGenerator, AsyncIterator, Callable
from datetime import datetime
from typing import Any, Dict, Final, List
from unittest.mock import AsyncMock, MagicMock
@ -389,24 +389,26 @@ class _RequestOnlyGuardrail(CustomGuardrail):
super().__init__(guardrail_name="request-only", event_hook=GuardrailEventHooks.pre_call, default_on=True)
self.asked: List[GuardrailEventHooks] = []
def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool:
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
self.asked.append(event_type)
return super().should_run_guardrail(data=data, event_type=event_type)
@pytest.mark.asyncio
async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardrail(
proxy_logging, make_user_api_key_auth, monkeypatch
):
guardrail = _RequestOnlyGuardrail()
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _RequestOnlyGuardrail()
monkeypatch.setattr(litellm, "callbacks", [guardrail])
chunk = ModelResponseStream(
chunk: Final = ModelResponseStream(
id="rid",
choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}],
created=0,
model="gpt-4o-mini",
)
out = await proxy_logging.async_post_call_streaming_hook(
out: Final = await proxy_logging.async_post_call_streaming_hook(
data={"messages": [{"role": "user", "content": "hi"}]},
response=chunk,
user_api_key_dict=make_user_api_key_auth(),
@ -416,23 +418,25 @@ async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardr
class _PerChunkRewriter(CustomLogger):
async def async_post_call_streaming_hook(self, user_api_key_dict, response): # type: ignore[override]
async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> str:
return f"modified-{response}"
@pytest.mark.asyncio
async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_beside_a_per_chunk_callback(
proxy_logging, make_user_api_key_auth, monkeypatch
):
guardrail = _RequestOnlyGuardrail()
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _RequestOnlyGuardrail()
monkeypatch.setattr(litellm, "callbacks", [guardrail, _PerChunkRewriter()])
chunk = ModelResponseStream(
chunk: Final = ModelResponseStream(
id="rid",
choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}],
created=0,
model="gpt-4o-mini",
)
out = await proxy_logging.async_post_call_streaming_hook(
out: Final = await proxy_logging.async_post_call_streaming_hook(
data={"messages": [{"role": "user", "content": "hi"}]},
response=chunk,
user_api_key_dict=make_user_api_key_auth(),
@ -474,8 +478,10 @@ async def _say_hi_through_chunk_hooks(proxy_logging: ProxyLogging, user_api_key_
@pytest.mark.asyncio
async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_guardrail(
proxy_logging, make_user_api_key_auth, monkeypatch
):
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _RequestOnlyGuardrail()
monkeypatch.setattr(litellm, "callbacks", [guardrail, _instance_rewriter()])
assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"
@ -484,8 +490,10 @@ async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_gua
@pytest.mark.asyncio
async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_class_hook(
proxy_logging, make_user_api_key_auth, monkeypatch
):
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "callbacks", [_ClassChunkLogger(), _instance_rewriter()])
assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"