From 8265510c2a8d6526dfd31cb6f82900420565c288 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 29 Sep 2026 22:27:24 +0000 Subject: [PATCH] fix(grayswan): merge response text and tool calls into one assistant monitor message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/grayswan/grayswan.py | 14 +++- .../observability/test_grayswan_wire.py | 74 ++++++++++--------- .../guardrail_hooks/test_grayswan.py | 25 +++++++ 3 files changed, 77 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 200fd785fa8..b14b3f1a0af 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -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 diff --git a/tests/integration/observability/test_grayswan_wire.py b/tests/integration/observability/test_grayswan_wire.py index 1ebd07b1cea..9c5a2d9be68 100644 --- a/tests/integration/observability/test_grayswan_wire.py +++ b/tests/integration/observability/test_grayswan_wire.py @@ -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) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py index 8fa0484947f..a3d7f9e73ef 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -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()