scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(grayswan): merge response text and tool calls into one assistant monitor message

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-29 22:27:24 +00:00
parent 8e93f092af
commit 8265510c2a
3 changed files with 77 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py
View file

@ -275,10 +275,20 @@ class GraySwanGuardrail(CustomGuardrail):
# Convert texts to messages format for GraySwan API
# Use "user" role for request content, "assistant" for response content
role: Final = "assistant" if input_type == "response" else "user"
merged_tail: Final = (
_MonitorMessage(role="assistant", content=texts[-1], tool_calls=response_tool_calls)
if texts and response_tool_calls
else None
)
messages: Final = (
*context,
*(_MonitorMessage(role=role, content=text) for text in texts),
*((_MonitorMessage(role="assistant", tool_calls=response_tool_calls),) if response_tool_calls else ()),
*(_MonitorMessage(role=role, content=text) for text in (texts[:-1] if merged_tail else texts)),
*((merged_tail,) if merged_tail else ()),
*(
(_MonitorMessage(role="assistant", tool_calls=response_tool_calls),)
if response_tool_calls and not texts
else ()
),
)
# Get dynamic params from request metadata

74
tests/integration/observability/test_grayswan_wire.py
View file

@ -65,25 +65,27 @@ def _grayswan_config(
on_flagged_action: str = "monitor",
streaming_end_of_stream_only: bool = False,
) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["guardrails"] = [
{
"guardrail_name": identity,
"litellm_params": {
"guardrail": "grayswan",
"mode": mode,
"default_on": True,
"api_base": vendor_url,
"api_key": _VENDOR_KEY,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
"optional_params": {
"on_flagged_action": on_flagged_action,
"violation_threshold": 0.5,
"policy_id": "synthetic-policy",
config: Final = {
**yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()),
"guardrails": [
{
"guardrail_name": identity,
"litellm_params": {
"guardrail": "grayswan",
"mode": mode,
"default_on": True,
"api_base": vendor_url,
"api_key": _VENDOR_KEY,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
"optional_params": {
"on_flagged_action": on_flagged_action,
"violation_threshold": 0.5,
"policy_id": "synthetic-policy",
},
},
},
}
]
}
],
}
path: Final = tmp_path / f"{identity}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@ -142,9 +144,10 @@ def test_post_call_sends_request_conversation_and_tools(gateway: Gateway, tmp_pa
request_messages: Final = [dict(message) for message in _REQUEST_MESSAGES]
request_tools: Final = [dict(tool) for tool in _TOOLS]
with wire_server(_vendor()) as vendor, wire_server(
_chat_provider({"role": "assistant", "content": response_text})
) as upstream:
with (
wire_server(_vendor()) as vendor,
wire_server(_chat_provider({"role": "assistant", "content": response_text})) as upstream,
):
config_path: Final = _grayswan_config(tmp_path, identity, vendor.url, "post_call")
with owned_proxy(gateway, tmp_path, {}, config=config_path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key=_PROVIDER_KEY)
@ -173,9 +176,10 @@ def test_post_call_scans_tool_call_only_response_and_blocks(gateway: Gateway, tm
"function": {"name": "send_email", "arguments": '{"to": "cfo@example.com", "body": "wire funds"}'},
}
with wire_server(_vendor(violation=1.0)) as vendor, wire_server(
_chat_provider({"role": "assistant", "content": None, "tool_calls": [tool_call]})
) as upstream:
with (
wire_server(_vendor(violation=1.0)) as vendor,
wire_server(_chat_provider({"role": "assistant", "content": None, "tool_calls": [tool_call]})) as upstream,
):
config_path: Final = _grayswan_config(tmp_path, identity, vendor.url, "post_call", on_flagged_action="block")
with owned_proxy(gateway, tmp_path, {}, config=config_path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key=_PROVIDER_KEY)
@ -241,9 +245,7 @@ def test_post_call_sends_anthropic_messages_conversation(gateway: Gateway, tmp_p
{"role": "user", "content": user_text},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": "toolu_inbox", "name": "read_inbox", "input": {}}
],
"content": [{"type": "tool_use", "id": "toolu_inbox", "name": "read_inbox", "input": {}}],
},
{
"role": "user",
@ -386,10 +388,13 @@ def test_post_call_streams_end_of_stream_with_conversation(gateway: Gateway, tmp
assert "streamed " in response.text and "summary" in response.text, response.text
(body,) = _monitor_bodies(vendor)
messages: Final = body["messages"]
assert messages == [*([dict(message) for message in _REQUEST_MESSAGES]), {
"role": "assistant",
"content": response_text,
}], body
assert messages == [
*([dict(message) for message in _REQUEST_MESSAGES]),
{
"role": "assistant",
"content": response_text,
},
], body
def test_pre_call_payload_shape_unchanged(gateway: Gateway, tmp_path: Path) -> None:
@ -397,9 +402,10 @@ def test_pre_call_payload_shape_unchanged(gateway: Gateway, tmp_path: Path) -> N
system_text: Final = "You are a mail assistant."
user_text: Final = f"summarize my inbox {identity}"
with wire_server(_vendor()) as vendor, wire_server(
_chat_provider({"role": "assistant", "content": "permitted"})
) as upstream:
with (
wire_server(_vendor()) as vendor,
wire_server(_chat_provider({"role": "assistant", "content": "permitted"})) as upstream,
):
config_path: Final = _grayswan_config(tmp_path, identity, vendor.url, "pre_call")
with owned_proxy(gateway, tmp_path, {}, config=config_path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key=_PROVIDER_KEY)

25
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py
View file

@ -747,6 +747,31 @@ async def test_post_call_scan_only_tool_results_scopes_context_and_tools() -> No
assert "tools" not in payload
@pytest.mark.asyncio
async def test_post_call_merges_response_text_and_tool_calls_into_one_message() -> None:
guardrail = _post_call_guardrail()
client = _CapturingClient()
guardrail.async_handler = client
tool_call = {
"id": "call_send",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "cfo@example.com"}'},
}
await guardrail.apply_guardrail(
inputs={"texts": ["response text"], "tool_calls": [tool_call]},
request_data={**_REQUEST_DATA, "litellm_logging_obj": _LoggingObj("acompletion")},
input_type="response",
logging_obj=_LoggingObj("acompletion"),
)
messages = list(client.calls[0]["json"]["messages"])
assert messages == [
*_REQUEST_DATA["messages"],
{"role": "assistant", "content": "response text", "tool_calls": (tool_call,)},
]
@pytest.mark.asyncio
async def test_post_call_prefers_request_route_over_logging_call_type() -> None:
guardrail = _post_call_guardrail()