test(guardrails): cover streamed tool-call name rewrites on chat and Messages
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled

Skipping the name write-back in either handler left every test green; a
guardrail that renames a tool call now has a regression test on both the
chat chunk path and the Anthropic SSE path
This commit is contained in:
mateo-berri 2026-09-09 14:21:40 -07:00
parent ddedb4867b
commit 0a053d2c81
2 changed files with 46 additions and 0 deletions

View file

@ -367,6 +367,29 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
assert '"stop_reason": "tool_use"' in raw
assert "persim" not in raw
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self):
class RenameTool(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
for tool_call in inputs.get("tool_calls", []):
tool_call.function.name = "lookup_fruit_reviewed"
return inputs
handler = AnthropicMessagesHandler()
chunks = self._ended_tool_use_sse_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=RenameTool(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
raw = b"".join(chunks).decode()
assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw
assert '"name": "lookup_fruit"' not in raw
assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"}
@pytest.mark.asyncio
async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()

View file

@ -1171,6 +1171,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
assert chunks[3].choices[0].delta.tool_calls is None
assert chunks[3].choices[0].finish_reason == "tool_calls"
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self):
class RenameTool(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
for tool_call in inputs.get("tool_calls", []):
tool_call["function"]["name"] = "lookup_fruit_reviewed"
return inputs
handler = OpenAIChatCompletionsHandler()
chunks = self._ended_tool_call_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=RenameTool(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]]
assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None]
assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"}
assert fragments[0].id == "call_1"
@pytest.mark.asyncio
async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()