diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9b0d1e60839..52985d49c7e 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: return None if texts is None else tuple(texts) +def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]: + return tuple(texts or ()) + + def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) @@ -152,7 +156,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): and returned ``None`` for, is re-scanned by the same translation so its texts reach the client through the translation's ended-stream write-back. A replacement whose scanned texts do not line up with the originals, or whose tool calls - differ from them, is undeliverable, so the executor releases the original chunks.""" + differ from them, is undeliverable, so the executor releases the original chunks. A stream + that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as + long as the hook left the tool calls alone.""" def __init__( self, @@ -185,15 +191,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail): rewrite: Final = hooked if replacement is None else replacement if rewrite is None: return inputs - scanned: Final = _text_snapshot(inputs.get("texts")) rescanned: Final = await self._rescan(rewrite, logging_obj) - if scanned is None or rescanned is None: + if rescanned is None: raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten: Final = rescanned.get("texts") - if rewritten is None or len(rewritten) != len(scanned): + if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if not rewritten: + return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} return rewritten_inputs diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 2af318c88cd..98674843c73 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1367,6 +1367,44 @@ class _LegacyScanningTranslation: return response +class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation): + """Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from + the inputs, while the non-streaming scan of the same response sends an empty list.""" + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]}) + await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +def _tool_only_chunk(): + return {"text": "", "tool_call": _chunk()["tool_call"]} + + def _native(text): return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} @@ -1375,11 +1413,17 @@ def _legacy_replacement(*texts, tool_calls=None): return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} -async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): - return await _run_legacy_streaming_steps(monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error) +async def _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None +): + return await _run_legacy_streaming_steps( + monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation + ) -async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail="block", on_error="next"): +async def _run_legacy_streaming_steps( + monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None +): monkeypatch.setattr(litellm, "callbacks", list(guardrails)) return await PipelineExecutor.execute_steps( steps=[ @@ -1397,7 +1441,7 @@ async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail=" call_type="completion", policy_name="p", streaming_chunks=chunks, - endpoint_translation=_LegacyScanningTranslation(), + endpoint_translation=_LegacyScanningTranslation() if translation is None else translation, ) @@ -1513,6 +1557,38 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( assert chunks == [_chunk()] +@pytest.mark.asyncio +async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert chunks == [_tool_only_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_tool_only_chunk()] + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world"))