diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 317203cc2d5..98ec60d8191 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -16,6 +16,7 @@ from typing import ( Protocol, TypeAlias, TypeVar, + cast, # noqa: TID251 # preserve the mutable JSON payload at the existing untyped boundary overload, runtime_checkable, ) @@ -116,13 +117,16 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor -from litellm.proxy.route_llm_request import ( - route_request, -) +from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ProxyLogging, + StreamingHookResponse, + StreamingReasoningState, + StreamingToolCallState, _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, + streaming_reasoning_with_response, + streaming_tool_calls_with_response, ) from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict @@ -3894,6 +3898,44 @@ class ProxyBaseLLMRequestProcessing: _finalize_streaming_generator_cleanup = finalize_streaming_generator_cleanup + @staticmethod + def _streaming_chunk_response_text(chunk: StreamingHookResponse | Mapping[str, object]) -> str: + if isinstance(chunk, (ModelResponse, ModelResponseStream)): + return litellm.get_response_string(response_obj=chunk) + if isinstance(chunk, Mapping): + return str(chunk.get("content", "")) + if isinstance(chunk, str): + return "" + try: + data: Final = chunk.model_dump(mode="json", exclude_none=True) + return str(data.get("content", "")) + except Exception: + return "" + + @staticmethod + async def _apply_streaming_chunk_hook( + *, + chunk: Any, # noqa: ANN401 # streaming callbacks may replace chunks with provider-specific response types + proxy_logging_obj: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + request_data: dict[str, object], # mutable-ok: ProxyLogging requires the existing mutable request payload + str_so_far: str, + streaming_tool_calls_so_far: StreamingToolCallState, + streaming_reasoning_so_far: StreamingReasoningState = (), + ) -> tuple[StreamingHookResponse, StreamingToolCallState, StreamingReasoningState]: + return ( + await proxy_logging_obj.async_post_call_streaming_hook( + user_api_key_dict=user_api_key_dict, + response=chunk, + data=request_data, + str_so_far=str_so_far, + streaming_tool_calls_so_far=streaming_tool_calls_so_far, + streaming_reasoning_so_far=streaming_reasoning_so_far, + ), + streaming_tool_calls_with_response(streaming_tool_calls_so_far, chunk), + streaming_reasoning_with_response(streaming_reasoning_so_far, chunk), + ) + @staticmethod async def async_streaming_data_generator( response: object, @@ -3945,6 +3987,8 @@ class ProxyBaseLLMRequestProcessing: ) try: str_so_far = "" + streaming_tool_calls_so_far: StreamingToolCallState = () # rebind-ok: accumulated streaming state advances after each chunk + streaming_reasoning_so_far: StreamingReasoningState = () # rebind-ok: accumulated streaming state advances after each chunk async for chunk in guarded_stream: # ``.format(chunk)`` was previously evaluated for every chunk # regardless of log level; gate it behind the level check. @@ -3952,25 +3996,21 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.debug("async_data_generator: received streaming chunk - %s", chunk) if not fast_path: - chunk = await proxy_logging_obj.async_post_call_streaming_hook( + ( + chunk, + streaming_tool_calls_so_far, + streaming_reasoning_so_far, + ) = await ProxyBaseLLMRequestProcessing._apply_streaming_chunk_hook( + chunk=chunk, + proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, - response=chunk, - data=request_data, + request_data=cast(dict[str, object], request_data), # cast-ok: request JSON has string keys str_so_far=str_so_far, + streaming_tool_calls_so_far=streaming_tool_calls_so_far, + streaming_reasoning_so_far=streaming_reasoning_so_far, ) - if isinstance(chunk, (ModelResponse, ModelResponseStream)): - response_str = litellm.get_response_string(response_obj=chunk) - str_so_far += response_str - elif hasattr(chunk, "model_dump"): - try: - d = chunk.model_dump(mode="json", exclude_none=True) - if isinstance(d, dict): - str_so_far += str(d.get("content", "")) - except Exception: - pass - elif isinstance(chunk, dict): - str_so_far += str(chunk.get("content", "")) + str_so_far += ProxyBaseLLMRequestProcessing._streaming_chunk_response_text(chunk) model_name = request_data.get("model", "") chunk = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e035f061db9..c718d17bd4c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -866,6 +866,8 @@ from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports PrismaClient, ProxyLogging, ProxyUpdateSpend, + StreamingReasoningState, + StreamingToolCallState, _cache_user_row, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _get_docs_url, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _get_openapi_url, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -893,6 +895,8 @@ from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, prefetch_config_params, + streaming_reasoning_with_response, + streaming_tool_calls_with_response, update_spend, ) from litellm.proxy.video_endpoints.endpoints import router as video_router @@ -9848,19 +9852,26 @@ async def _apply_streaming_chunk_hooks( user_api_key_dict: UserAPIKeyAuth, request_data: dict, str_so_far: str, -) -> tuple[Any, str]: + streaming_tool_calls_so_far: StreamingToolCallState = (), + streaming_reasoning_so_far: StreamingReasoningState = (), +) -> tuple[Any, str, StreamingToolCallState, StreamingReasoningState]: + stream_chunk: Final = chunk chunk = await proxy_logging_obj.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, response=chunk, data=request_data, str_so_far=str_so_far if str_so_far else None, + streaming_tool_calls_so_far=streaming_tool_calls_so_far, + streaming_reasoning_so_far=streaming_reasoning_so_far, ) if isinstance(chunk, (ModelResponse, ModelResponseStream)): response_str: Final = litellm.get_response_string(response_obj=chunk) str_so_far += response_str - return chunk, str_so_far + updated_tool_calls: Final = streaming_tool_calls_with_response(streaming_tool_calls_so_far, stream_chunk) + updated_reasoning: Final = streaming_reasoning_with_response(streaming_reasoning_so_far, stream_chunk) + return chunk, str_so_far, updated_tool_calls, updated_reasoning def _format_streaming_sse_chunk(chunk: str | bytes) -> str | bytes: @@ -10128,6 +10139,8 @@ async def async_data_generator( # Previously "".join(str_so_far_parts) was called every chunk, re-joining # the entire accumulated response. String += is O(n) amortized total. _str_so_far: str = "" + _streaming_tool_calls_so_far: StreamingToolCallState = () # rebind-ok: accumulated streaming state advances after each chunk + _streaming_reasoning_so_far: StreamingReasoningState = () # rebind-ok: accumulated streaming state advances after each chunk # Separate iterator-level vs per-chunk hook decisions. The iterator # wrap is needed when any callback overrides # ``async_post_call_streaming_iterator_hook`` or has @@ -10167,11 +10180,18 @@ async def async_data_generator( chunk = cast(Any, item) # cast-ok: sentinel already handled above, item is a real chunk here if needs_per_chunk_hook: ### CALL HOOKS ### - modify outgoing data - chunk, _str_so_far = await _apply_streaming_chunk_hooks( + ( + chunk, + _str_so_far, + _streaming_tool_calls_so_far, + _streaming_reasoning_so_far, + ) = await _apply_streaming_chunk_hooks( chunk=chunk, user_api_key_dict=user_api_key_dict, request_data=request_data, str_so_far=_str_so_far, + streaming_tool_calls_so_far=_streaming_tool_calls_so_far, + streaming_reasoning_so_far=_streaming_reasoning_so_far, ) # Mid-stream fallbacks surface metadata on individual chunks rather than diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4743970a526..4f48376f3c2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,6 +19,7 @@ from collections.abc import ( Awaitable, Callable, Coroutine, + Iterator, Mapping, Sequence, ) @@ -255,7 +256,7 @@ from litellm.repositories.verification_token_repository import ( from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ChatCompletionReasoningSummaryTextBlock, ResponsesAPIResponse from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -263,7 +264,12 @@ from litellm.types.mcp import ( ) from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult -from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + LLMResponseTypes, + LoggedLiteLLMParams, +) from litellm.utils import ( add_custom_logger_callback_to_specific_event, ) @@ -1102,6 +1108,198 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +StreamingHookResponse: TypeAlias = str | ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream + + +class _StreamingHookResponseText(str): + """Marks the exact text object passed to a per-chunk streaming hook.""" + + +class _StructuredStreamingGuardrailText(_StreamingHookResponseText): + pass + + +@dataclass(frozen=True, slots=True) +class _StreamingToolCallFragment: + choice_index: int + tool_index: int + name: str + arguments: str + + @property + def key(self) -> tuple[int, int]: + return self.choice_index, self.tool_index + + +StreamingToolCallState: TypeAlias = tuple[_StreamingToolCallFragment, ...] + + +@dataclass(frozen=True, slots=True) +class _StreamingReasoningFragment: + choice_index: int + text: str + + +StreamingReasoningState: TypeAlias = tuple[_StreamingReasoningFragment, ...] + + +def _streaming_reasoning_text(delta: Delta) -> str: + reasoning: Final = delta.reasoning_content if hasattr(delta, "reasoning_content") else None + thinking_blocks: Final = delta.thinking_blocks if hasattr(delta, "thinking_blocks") else None + reasoning_items: Final = delta.reasoning_items if hasattr(delta, "reasoning_items") else None + thinking: Final = tuple( + _thinking_block_text(block) for block in thinking_blocks or () if block["type"] == "thinking" + ) + summaries: Final = tuple(_reasoning_summary_text(item.get("summary") or ()) for item in reasoning_items or ()) + return "".join((reasoning or "", *thinking, *summaries)) + + +def _thinking_block_text(block: Mapping[str, object]) -> str: + text: Final = block.get("thinking") + return text if isinstance(text, str) else "" + + +def _streaming_reasoning_fragments(response: ModelResponseStream) -> Iterator[_StreamingReasoningFragment]: + for choice in response.choices: + if text := _streaming_reasoning_text(choice.delta): + yield _StreamingReasoningFragment(choice_index=choice.index, text=text) + + +def _reasoning_summary_text(summary: Sequence[ChatCompletionReasoningSummaryTextBlock]) -> str: + return "".join(block.get("text", "") for block in summary) + + +def streaming_reasoning_with_response( + reasoning: Sequence[_StreamingReasoningFragment], response: object +) -> StreamingReasoningState: + if not isinstance(response, ModelResponseStream): + return tuple(reasoning) + fragments: Final = (*reasoning, *_streaming_reasoning_fragments(response)) + indices: Final = tuple(dict.fromkeys(fragment.choice_index for fragment in fragments)) + return tuple( + _StreamingReasoningFragment( + choice_index=index, + text="".join(fragment.text for fragment in fragments if fragment.choice_index == index), + ) + for index in indices + ) + + +def _streaming_hook_response_text(*, response_str: str, str_so_far: str | None, response: object) -> str: + complete_response: Final = str_so_far + response_str if str_so_far is not None else response_str + if complete_response == "" and isinstance(response, (ModelResponse, ModelResponseStream)): + return _StreamingHookResponseText(complete_response) + return complete_response + + +def _streaming_tool_call_fragments(response: ModelResponseStream) -> Iterator[_StreamingToolCallFragment]: + for choice in response.choices: + for tool_call in choice.delta.tool_calls or (): + if isinstance(tool_call, ChatCompletionDeltaToolCall): + yield _StreamingToolCallFragment( + choice_index=choice.index, + tool_index=tool_call.index, + name=tool_call.function.name or "", + arguments=tool_call.function.arguments or "", + ) + else: + yield _StreamingToolCallFragment( + choice_index=choice.index, + tool_index=tool_call.index, + name=tool_call.custom.name or "", + arguments=tool_call.custom.input or "", + ) + if choice.delta.function_call is not None: + yield _StreamingToolCallFragment( + choice_index=choice.index, + tool_index=-1, + name=choice.delta.function_call.name or "", + arguments=choice.delta.function_call.arguments or "", + ) + + +def _assembled_streaming_tool_calls( + fragments: Sequence[_StreamingToolCallFragment], +) -> StreamingToolCallState: + keys: Final = tuple( + fragment.key + for position, fragment in enumerate(fragments) + if fragment.key not in tuple(previous.key for previous in fragments[:position]) + ) + return tuple( + _StreamingToolCallFragment( + choice_index=key[0], + tool_index=key[1], + name="".join(fragment.name for fragment in fragments if fragment.key == key), + arguments="".join(fragment.arguments for fragment in fragments if fragment.key == key), + ) + for key in keys + ) + + +def streaming_tool_calls_with_response( + tool_calls: Sequence[_StreamingToolCallFragment], response: object +) -> StreamingToolCallState: + if not isinstance(response, ModelResponseStream): + return tuple(tool_calls) + fragments: Final = _streaming_tool_call_fragments(response) + return _assembled_streaming_tool_calls((*tool_calls, *fragments)) + + +def _streaming_guardrail_response_text( + *, + complete_response: str, + response: object, + streaming_tool_calls_so_far: Sequence[_StreamingToolCallFragment], + streaming_reasoning_so_far: Sequence[_StreamingReasoningFragment], +) -> str: + if not isinstance(response, ModelResponseStream): + return complete_response + tool_calls: Final = streaming_tool_calls_with_response(streaming_tool_calls_so_far, response) + tool_fields: Final = tuple( + "tool_call:" + + json.dumps( + (tool_call.choice_index, tool_call.tool_index, tool_call.name, tool_call.arguments), + ensure_ascii=False, + separators=(",", ":"), + ) + for tool_call in tool_calls + ) + reasoning: Final = streaming_reasoning_with_response(streaming_reasoning_so_far, response) + reasoning_fields: Final = tuple( + "reasoning:" + json.dumps((fragment.choice_index, fragment.text), ensure_ascii=False, separators=(",", ":")) + for fragment in reasoning + ) + structured_fields: Final = (*tool_fields, *reasoning_fields) + if not structured_fields: + return complete_response + return _StructuredStreamingGuardrailText("\n".join((str(complete_response), *structured_fields))) + + +def _is_unchanged_structured_streaming_hook_response( + *, callback_response: object, complete_response: str, response_str: str, response: object +) -> bool: + if not isinstance(response, (ModelResponse, ModelResponseStream)): + return False + if isinstance(complete_response, _StructuredStreamingGuardrailText): + return callback_response == complete_response + if isinstance(complete_response, _StreamingHookResponseText): + return callback_response is complete_response + if response_str != "" and not ( + isinstance(response, ModelResponseStream) + and any( + choice.delta.tool_calls + or choice.delta.function_call is not None + or getattr(choice.delta, "reasoning_content", None) + or getattr(choice.delta, "thinking_blocks", None) + or getattr(choice.delta, "reasoning_items", None) + for choice in response.choices + ) + ): + return False + return callback_response == complete_response + + _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) @@ -3885,6 +4083,8 @@ class ProxyLogging: response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream, user_api_key_dict: UserAPIKeyAuth, str_so_far: str | None = None, + streaming_tool_calls_so_far: Sequence[_StreamingToolCallFragment] = (), + streaming_reasoning_so_far: Sequence[_StreamingReasoningFragment] = (), ): """ Allow user to modify outgoing streaming data -> per chunk @@ -3904,6 +4104,7 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router + hook_response: StreamingHookResponse = response # rebind-ok: each callback may replace the chunk response_str: str | None = None if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) @@ -3951,22 +4152,35 @@ class ProxyLogging: else: _callback = callback if _callback is not None and isinstance(_callback, CustomLogger): - if str_so_far is not None: - complete_response = str_so_far + response_str - else: - complete_response = response_str + complete_response = _streaming_hook_response_text( + response_str=response_str, + str_so_far=str_so_far, + response=hook_response, + ) + if isinstance(_callback, CustomGuardrail): + complete_response = _streaming_guardrail_response_text( + complete_response=complete_response, + response=hook_response, + streaming_tool_calls_so_far=streaming_tool_calls_so_far, + streaming_reasoning_so_far=streaming_reasoning_so_far, + ) callback_response: ( - ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None + str | ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None ) callback_response = await _callback.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, response=complete_response, ) - if callback_response is not None: - response = callback_response + if callback_response is not None and not _is_unchanged_structured_streaming_hook_response( + callback_response=callback_response, + complete_response=complete_response, + response_str=response_str, + response=hook_response, + ): + hook_response = callback_response except Exception as e: raise e - return response + return hook_response async def async_post_call_streaming_iterator_hook( self, diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 09d71cb4ead..546a7849125 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -31,6 +31,7 @@ from pydantic import BaseModel import litellm import litellm.proxy.proxy_server as ps from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, @@ -78,6 +79,30 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr ) +def _tool_call_chunk(arguments: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-tool-call", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "type": "function", + "function": {"name": "run_command", "arguments": arguments}, + } + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + async def _async_iter(items): for it in items: yield it @@ -589,12 +614,15 @@ def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch): chunk = _simple_chunk(content="abc") - async def _passthrough(*, user_api_key_dict, response, data, str_so_far=None): + async def _passthrough( + *, user_api_key_dict, response, data, str_so_far=None, streaming_tool_calls_so_far=(), + streaming_reasoning_so_far=(), + ): return response monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough) - new_chunk, new_str = await _apply_streaming_chunk_hooks( + new_chunk, new_str, tool_calls, _ = await _apply_streaming_chunk_hooks( chunk=chunk, user_api_key_dict=_user_auth(), request_data={}, @@ -605,14 +633,40 @@ async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch): "chunk_is_basemodel": isinstance(new_chunk, ModelResponseStream), "str_so_far": new_str, "grew": len(new_str) > len("prior:"), + "tool_calls": tool_calls, } assert observed == { "chunk_is_basemodel": True, "str_so_far": "prior:abc", "grew": True, + "tool_calls": (), } +@pytest.mark.asyncio +async def test_apply_streaming_chunk_hooks_compacts_tool_call_history(monkeypatch): + async def _passthrough( + *, user_api_key_dict, response, data, str_so_far=None, streaming_tool_calls_so_far=(), + streaming_reasoning_so_far=(), + ): + return response + + monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough) + stream_state = () + + for arguments in ('{"command":"', "blocked-", "command", '"}'): + _, _, stream_state, _ = await _apply_streaming_chunk_hooks( + chunk=_tool_call_chunk(arguments), + user_api_key_dict=_user_auth(), + request_data={}, + str_so_far="", + streaming_tool_calls_so_far=stream_state, + ) + + assert len(stream_state) == 1 + assert stream_state[0].arguments == '{"command":"blocked-command"}' + + @pytest.mark.asyncio async def test_apply_streaming_chunk_hooks_hook_raises_exception(monkeypatch): async def _boom(*args, **kwargs): @@ -715,6 +769,80 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): } +@pytest.mark.asyncio +async def test_async_data_generator_preserves_tool_calls_through_per_chunk_hook( + monkeypatch, +): + class _PassThrough(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return response + + callback = _PassThrough() + monkeypatch.setattr(litellm, "callbacks", [callback]) + _patch_logging_flags(monkeypatch, needs_per_chunk=True) + + chunk = ModelResponseStream( + id="chatcmpl-tools", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call-weather", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "", + }, + }, + { + "index": 1, + "id": "call-time", + "type": "function", + "function": {"name": "get_time", "arguments": ""}, + }, + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + out = [] + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4o-mini"}, + ): + out.append(line) + + first = out[0] + assert isinstance(first, (str, bytes)) + first_text = first.decode() if isinstance(first, bytes) else first + assert first_text.startswith("data: {") + payload = json.loads(first_text.removeprefix("data: ").removesuffix("\n\n")) + assert payload["choices"][0]["delta"]["tool_calls"] == [ + { + "index": 0, + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + }, + { + "index": 1, + "id": "call-time", + "type": "function", + "function": {"name": "get_time", "arguments": ""}, + }, + ] + assert out[-1] == "data: [DONE]\n\n" + + @pytest.mark.asyncio async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch): _patch_logging_flags(monkeypatch) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 1aed60ee5e2..90554d5183c 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4824,6 +4824,48 @@ class TestAsyncStreamingDataGeneratorFastPath: ProxyLogging._callback_capabilities_cache.clear() + @pytest.mark.asyncio + async def test_apply_streaming_chunk_hook_compacts_tool_call_history(self): + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + + def _chunk(arguments: str): + return litellm.ModelResponseStream( + id="chatcmpl-tool-call", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "type": "function", + "function": {"name": "run_command", "arguments": arguments}, + } + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + stream_state = () + for arguments in ("blocked-", "command"): + _, stream_state, _ = await ProxyBaseLLMRequestProcessing._apply_streaming_chunk_hook( + chunk=_chunk(arguments), + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=MagicMock(spec=ProxyUserAPIKeyAuth), + request_data={}, + str_so_far="", + streaming_tool_calls_so_far=stream_state, + ) + + assert len(stream_state) == 1 + assert stream_state[0].arguments == "blocked-command" + class TestDisconnectGatherCleanup: def _disconnect_request(self) -> Request: diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 50f50478ad3..986c1d9ad62 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,11 +10,13 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio +import json from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from fastapi import HTTPException @@ -27,7 +29,7 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, streaming_tool_calls_with_response from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Usage @@ -384,6 +386,241 @@ async def test_async_post_call_streaming_hook_invokes_per_chunk_callback(proxy_l assert out.startswith("modified-") +@pytest.mark.parametrize("str_so_far", [None, "I will check that. "]) +@pytest.mark.parametrize("content", [None, "Checking those now. "]) +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_preserves_tool_calls_when_callback_returns_unmodified_text( + proxy_logging, make_user_api_key_auth, monkeypatch, str_so_far, content +): + class _PassThrough(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return response + + monkeypatch.setattr(litellm, "callbacks", [_PassThrough()]) + + response = litellm.ModelResponseStream( + id="chatcmpl-tools", + choices=[ + { + "index": 0, + "delta": { + "content": content, + "tool_calls": [ + { + "index": 0, + "id": "call-weather", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + }, + { + "index": 1, + "id": "call-time", + "type": "function", + "function": {"name": "get_time", "arguments": ""}, + }, + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + out = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=response, + user_api_key_dict=make_user_api_key_auth(), + str_so_far=str_so_far, + ) + + assert out is response + assert [ + tool_call.model_dump(exclude_none=True) + for tool_call in out.choices[0].delta.tool_calls + ] == [ + { + "id": "call-weather", + "function": {"arguments": "", "name": "get_weather"}, + "type": "function", + "index": 0, + }, + { + "id": "call-time", + "function": {"arguments": "", "name": "get_time"}, + "type": "function", + "index": 1, + }, + ] + + class _Replace(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return "replacement" + + monkeypatch.setattr(litellm, "callbacks", [_PassThrough(), _Replace()]) + replaced = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=response, + user_api_key_dict=make_user_api_key_auth(), + str_so_far=str_so_far, + ) + assert replaced == "replacement" + + if str_so_far is not None: + class _EquivalentCopy(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return response.encode().decode() + + monkeypatch.setattr(litellm, "callbacks", [_EquivalentCopy()]) + equivalent = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=response, + user_api_key_dict=make_user_api_key_auth(), + str_so_far=str_so_far, + ) + assert equivalent is response + + class _Suppress(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return "" + + monkeypatch.setattr(litellm, "callbacks", [_Suppress()]) + suppressed = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=response, + user_api_key_dict=make_user_api_key_auth(), + str_so_far=str_so_far, + ) + assert suppressed == "" + + +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_exposes_tool_calls_to_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class _BlockingGuardrail(CustomGuardrail): + def should_run_guardrail(self, data, event_type): + return True + + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + if "blocked-command" in response: + return "data: blocked\n\n" + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [_BlockingGuardrail(guardrail_name="tool-call-scanner")], + ) + + first_chunk = litellm.ModelResponseStream( + id="chatcmpl-guardrail-tools", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call-shell", + "type": "function", + "function": { + "name": "run_command", + "arguments": '{"command":"blocked-', + }, + } + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + second_chunk = litellm.ModelResponseStream( + id="chatcmpl-guardrail-tools", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "function": {"arguments": 'command"}'}, + } + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + out = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=second_chunk, + user_api_key_dict=make_user_api_key_auth(), + streaming_tool_calls_so_far=streaming_tool_calls_with_response((), first_chunk), + ) + + assert out == "data: blocked\n\n" + + +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_preserves_equal_guardrail_projection( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class _CopyingGuardrail(CustomGuardrail): + def should_run_guardrail(self, data, event_type): + return True + + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return response.encode().decode() + + monkeypatch.setattr( + litellm, + "callbacks", + [_CopyingGuardrail(guardrail_name="tool-call-scanner")], + ) + response = litellm.ModelResponseStream( + id="chatcmpl-guardrail-copy", + choices=[ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call-weather", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + } + ] + }, + "finish_reason": None, + } + ], + created=0, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + out = await proxy_logging.async_post_call_streaming_hook( + data={}, + response=response, + user_api_key_dict=make_user_api_key_auth(), + ) + + assert out is response + + @pytest.mark.asyncio async def test_async_post_call_streaming_hook_callback_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch): class _Per(CustomLogger): @@ -901,3 +1138,173 @@ async def test_stream_provider_failure_is_not_attributed_to_any_guardrail( with pytest.raises(RuntimeError, match="provider connection dropped"): await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _failing_provider_stream(), request_data) assert request_data["metadata"] == {} + + +@pytest.fixture +def reasoning_guardrail(proxy_logging, monkeypatch): + from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import ( + AzureContentSafetyTextModerationGuardrail, + ) + + guardrail: Final = AzureContentSafetyTextModerationGuardrail( + guardrail_name="reasoning-scanner", api_key="test-key", api_base="https://moderation.test", + event_hook=GuardrailEventHooks.post_call, default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + return guardrail + + +def _moderate_reasoning(request: httpx.Request) -> httpx.Response: + text: Final = json.loads(request.content)["text"] + assert "opaque-marker" not in text + return httpx.Response(200, json={ + "blocklistsMatch": [], + "categoriesAnalysis": [{"category": "Violence", "severity": 6 if "blocked-command" in text else 0}], + }) + + +def _reasoning_chunk(field: str, text: str, choice_index: int = 0) -> litellm.ModelResponseStream: + delta: Final = { + "reasoning_content": {"reasoning_content": text}, + "thinking_blocks": {"thinking_blocks": [ + {"type": "thinking", "thinking": text, "signature": "opaque-marker"}, + {"type": "redacted_thinking", "data": "opaque-marker"}, + ]}, + "reasoning_items": {"reasoning_items": [{ + "type": "reasoning", "id": "r1", "encrypted_content": "opaque-marker", + "summary": [{"type": "summary_text", "text": text}], + }]}, + }[field] + return litellm.ModelResponseStream(choices=[{"index": choice_index, "delta": delta}]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["reasoning_content", "thinking_blocks", "reasoning_items"]) +@pytest.mark.parametrize("text", ["blocked-command", "safe reasoning", ""]) +async def test_native_azure_moderates_plaintext_reasoning( + field, text, reasoning_guardrail, proxy_logging, make_user_api_key_auth, +): + chunk: Final = _reasoning_chunk(field, text) + async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler: + reasoning_guardrail.async_handler = handler + result: Final = await proxy_logging.async_post_call_streaming_hook( + data={}, response=chunk, user_api_key_dict=make_user_api_key_auth(), + ) + if text == "blocked-command": + assert isinstance(result, str) + assert "Azure Content Safety Guardrail" in result + assert "blocked-command" not in result + else: + assert result is chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["reasoning_content", "thinking_blocks", "reasoning_items"]) +@pytest.mark.parametrize("path", ["proxy_server", "common_request_processing"]) +@pytest.mark.parametrize("same_choice", [True, False]) +async def test_native_azure_moderates_split_reasoning_in_both_streams( + field, path, same_choice, reasoning_guardrail, proxy_logging, make_user_api_key_auth, +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.proxy_server import async_data_generator + + chunks: Final = (_reasoning_chunk(field, "blocked-"), _reasoning_chunk(field, "command", 0 if same_choice else 1)) + + async def source(): + for chunk in chunks: + yield chunk + + stream: Final = ( + async_data_generator(response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={}) + if path == "proxy_server" + else ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={}, + proxy_logging_obj=proxy_logging, + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda error: f"data: {error}\n\n", + ) + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler: + reasoning_guardrail.async_handler = handler + output: Final = tuple([part async for part in stream]) + assert len(output) >= 2 + if same_choice: + assert isinstance(output[1], str) + assert "Azure Content Safety Guardrail" in output[1] + assert '"command"' not in output[1] + else: + if path == "common_request_processing": + assert output[1] is chunks[1] + else: + assert isinstance(output[1], str) + assert json.loads(output[1].removeprefix("data: "))["choices"] == chunks[1].model_dump(exclude_none=True)["choices"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("callback_kind", ["ordinary", "disabled_guardrail"]) +@pytest.mark.parametrize("content", [None, "answer"]) +async def test_reasoning_chunks_preserve_ordinary_and_disabled_callback_behavior( + callback_kind, content, proxy_logging, make_user_api_key_auth, monkeypatch, +): + class OrdinaryCallback(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + assert "blocked-command" not in response + return response + + class DisabledGuardrail(CustomGuardrail): + def should_run_guardrail(self, data, event_type): + return False + + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + pytest.fail("Disabled guardrail ran") + + callback: Final = OrdinaryCallback() if callback_kind == "ordinary" else DisabledGuardrail(guardrail_name="disabled") + monkeypatch.setattr(litellm, "callbacks", [callback]) + chunk: Final = litellm.ModelResponseStream(choices=[{ + "index": 0, "delta": {"content": content, "reasoning_content": "blocked-command"}, + }]) + result: Final = await proxy_logging.async_post_call_streaming_hook( + data={}, response=chunk, user_api_key_dict=make_user_api_key_auth(), + ) + assert result is chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["proxy_server", "common_request_processing"]) +async def test_reasoning_stream_preserves_opaque_items_without_optional_plaintext( + path, reasoning_guardrail, proxy_logging, make_user_api_key_auth, +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.proxy_server import async_data_generator + + chunks: Final = tuple(litellm.ModelResponseStream(choices=[{ + "index": 0, "delta": {"reasoning_items": [item]}, + }]) for item in ( + {"type": "reasoning", "id": "encrypted", "encrypted_content": "opaque-marker"}, + {"type": "reasoning", "id": "empty-summary", "summary": [{"type": "summary_text"}]}, + )) + + async def source(): + for chunk in chunks: + yield chunk + + stream: Final = ( + async_data_generator(response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={}) + if path == "proxy_server" + else ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=source(), user_api_key_dict=make_user_api_key_auth(), request_data={}, + proxy_logging_obj=proxy_logging, + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda error: f"data: {error}\n\n", + ) + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(_moderate_reasoning)) as handler: + reasoning_guardrail.async_handler = handler + output: Final = tuple([part async for part in stream]) + assert len(output) >= 2 + if path == "common_request_processing": + assert output[:2] == chunks + else: + for chunk, frame in zip(chunks, output[:2]): + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: "))["choices"] == chunk.model_dump(exclude_none=True)["choices"]