"""End-to-end demo: SafetyFinishReasonMiddleware on the real DeerFlow lead-agent. What it proves -------------- - The real ``make_lead_agent`` / ``DeerFlowClient`` pipeline is built (full 18-middleware chain, sandbox, tools, etc.). - A model that returns ``finish_reason='content_filter'`` + ``tool_calls`` triggers SafetyFinishReasonMiddleware. - LangChain's tool router never invokes ``write_file`` — the truncated arguments do **not** reach the sandbox. - A ``safety_termination`` custom event is emitted on the stream and the final AIMessage carries the observability stamp. Run from backend/ directory: PYTHONPATH=. uv run python scripts/e2e_safety_termination_demo.py """ from __future__ import annotations import sys from typing import Any from langchain_core.language_models import BaseChatModel from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, ChatResult # --------------------------------------------------------------------------- # Fake provider that mimics Moonshot's content_filter behaviour # --------------------------------------------------------------------------- class _ContentFilteredFakeModel(BaseChatModel): """First call returns finish_reason=content_filter + truncated write_file tool_call. Subsequent calls return a normal stop response so the agent can terminate (the middleware should make a second call unnecessary by clearing tool_calls, but we keep this safety net in case loop-detection or anything else triggers another model invocation).""" call_count: int = 0 @property def _llm_type(self) -> str: return "fake-content-filtered" def bind_tools(self, tools, **kwargs): return self def _generate(self, messages, stop=None, run_manager=None, **kwargs): self.call_count += 1 if self.call_count == 1: msg = AIMessage( content="# 政经周报\n- **会晤时间**：2026年5月12日—13日，特朗普访问中国，与", tool_calls=[ { "id": "call_truncated_write", "name": "write_file", "args": { "path": "/mnt/user-data/outputs/political-economic-news-weekly-may-16-2026.md", "content": "# 政经周报\n- **会晤时间**：2026年5月12日—13日，特朗普访问中国，与", }, } ], response_metadata={ "finish_reason": "content_filter", "model_name": "kimi-k2.6", "model_provider": "openai", }, ) else: msg = AIMessage( content="(secondary call, should not be needed)", response_metadata={"finish_reason": "stop", "model_name": "kimi-k2.6"}, ) return ChatResult(generations=[ChatGeneration(message=msg)]) async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs): return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs) # --------------------------------------------------------------------------- # Driver # --------------------------------------------------------------------------- def main() -> int: # Inject the fake model BEFORE constructing the client. Both the # client module and the lead-agent module bind ``create_chat_model`` # at import time via ``from deerflow.models import create_chat_model``, # so we patch both attribute slots — the source-of-truth patch on # ``factory.create_chat_model`` doesn't propagate back into already- # imported names. import deerflow.agents.lead_agent.agent as lead_agent_module import deerflow.client as client_module fake = _ContentFilteredFakeModel() originals = { "lead": lead_agent_module.create_chat_model, "client": client_module.create_chat_model, } def fake_create_chat_model(*args, **kwargs): return fake lead_agent_module.create_chat_model = fake_create_chat_model client_module.create_chat_model = fake_create_chat_model from deerflow.client import DeerFlowClient try: client = DeerFlowClient() print("\n=== Streaming a turn through the real lead-agent ===") events: list[dict[str, Any]] = [] for event in client.stream( "帮我整理一下最近一周政经新闻，写到 /mnt/user-data/outputs/political-economic-news-weekly-may-16-2026.md", thread_id="e2e-safety-1", ): events.append({"type": event.type, "data": event.data}) # ---- Assertions ---- safety_event = next( (e for e in events if e["type"] == "custom" and isinstance(e["data"], dict) and e["data"].get("type") == "safety_termination"), None, ) final_values = next( (e for e in reversed(events) if e["type"] == "values"), None, ) tool_messages = [e for e in events if e["type"] == "messages-tuple" and isinstance(e["data"], dict) and e["data"].get("type") == "tool"] ai_tool_call_messages = [e for e in events if e["type"] == "messages-tuple" and isinstance(e["data"], dict) and e["data"].get("type") == "ai" and e["data"].get("tool_calls")] print(f"\n[stats] total stream events: {len(events)}") print(f"[stats] model call count: {fake.call_count}") print(f"[stats] tool messages on stream: {len(tool_messages)}") print(f"[stats] AI messages carrying tool_calls: {len(ai_tool_call_messages)}") print("\n[event] safety_termination custom event:") if safety_event is None: print(" *** NOT FOUND ***") return 1 for k, v in safety_event["data"].items(): print(f" {k}: {v}") print("\n[state] final AIMessage from last values snapshot:") if final_values is None: print(" *** no values snapshot ***") return 1 # `values` event carries `_serialize_message` dicts, not Message objects. final_messages = final_values["data"].get("messages") or [] last_ai = next((m for m in reversed(final_messages) if isinstance(m, dict) and m.get("type") == "ai"), None) if last_ai is None: print(" *** no AIMessage in final state ***") print(f" message types seen: {[m.get('type') if isinstance(m, dict) else type(m).__name__ for m in final_messages]}") return 1 tool_calls = last_ai.get("tool_calls") or [] additional_kwargs = last_ai.get("additional_kwargs") or {} response_metadata = last_ai.get("response_metadata") or {} content = last_ai.get("content") print(f" tool_calls (must be empty): {tool_calls}") print(f" additional_kwargs.safety_termination: {additional_kwargs.get('safety_termination')}") content_preview = (content if isinstance(content, str) else str(content))[:200] print(f" content[:200]: {content_preview!r}") print(f" response_metadata.finish_reason: {response_metadata.get('finish_reason')}") # NOTE: `client._serialize_message` does not include `response_metadata` # in the values-event payload (client-layer behaviour, unrelated to the # middleware). The middleware *does* preserve finish_reason on the # AIMessage object — see test_safety_finish_reason_middleware.py:: # TestMessageRewrite::test_preserves_response_metadata_finish_reason. # Here we assert on the observability stamp, which carries the same # evidence and is in the serialized payload. stamp = additional_kwargs.get("safety_termination") or {} failures = [] if tool_calls: failures.append("final AIMessage still has tool_calls — middleware did NOT clear them") if not stamp: failures.append("final AIMessage missing safety_termination observability stamp") if tool_messages: failures.append(f"tool node was invoked: {len(tool_messages)} ToolMessage(s) on stream") if stamp.get("reason_value") != "content_filter": failures.append(f"safety_termination.reason_value was {stamp.get('reason_value')!r}, expected 'content_filter'") if safety_event is None: failures.append("safety_termination custom event was not emitted on the stream") if failures: print("\n=== FAIL ===") for f in failures: print(f" - {f}") return 1 print("\n=== PASS ===") print(" - tool_calls cleared on final AIMessage") print(" - tool node never invoked (no ToolMessage on stream)") print(" - safety_termination custom event emitted") print(" - observability stamp written to additional_kwargs") print(" - response_metadata.finish_reason preserved for downstream SSE") return 0 finally: lead_agent_module.create_chat_model = originals["lead"] client_module.create_chat_model = originals["client"] if __name__ == "__main__": sys.exit(main())