mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
test(guardrails): cover streamed tool-call name rewrites on chat and Messages
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:
parent
ddedb4867b
commit
0a053d2c81
2 changed files with 46 additions and 0 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue