mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +00:00
103 lines
5.3 KiB
Python
103 lines
5.3 KiB
Python
"""PII redaction must cover malformed and legacy calls at memory admission."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, messages_from_dict, messages_to_dict
|
|
from langchain_openai.chat_models.base import _convert_dict_to_message
|
|
from langgraph.runtime import Runtime
|
|
|
|
from deerflow.agents.memory import summarization_hook
|
|
from deerflow.agents.middlewares import memory_middleware
|
|
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
|
|
from deerflow.agents.middlewares.pii_redaction_middleware import redact_text
|
|
from deerflow.agents.middlewares.summarization_middleware import SummarizationEvent
|
|
from deerflow.config.memory_config import MemoryConfig
|
|
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
|
|
|
_EMAIL = "alice@example.com"
|
|
_SECRET = "memory-regression-deployment-secret"
|
|
|
|
|
|
def _attempt(kind: str) -> AIMessage:
|
|
if kind == "invalid-checkpoint":
|
|
# A persisted parsed-invalid view need not retain provider-raw calls.
|
|
message = AIMessage(content="Preparing contact lookup", invalid_tool_calls=[{"name": "lookup", "id": "call-1", "args": '{"email":"alice@example.com",', "error": "Invalid arguments for alice@example.com"}])
|
|
elif kind == "legacy":
|
|
message = _convert_dict_to_message({"role": "assistant", "content": "Preparing contact lookup", "function_call": {"name": "lookup", "arguments": '{"email":"alice@example.com"}'}})
|
|
else:
|
|
message = _convert_dict_to_message({"role": "assistant", "content": "Preparing contact lookup", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": '{"email":"alice@example.com",'}}]})
|
|
assert isinstance(message, AIMessage) and not message.tool_calls and message.invalid_tool_calls
|
|
assert isinstance(message, AIMessage)
|
|
return messages_from_dict(messages_to_dict([message]))[0]
|
|
|
|
|
|
def _capture(monkeypatch: pytest.MonkeyPatch, entry: str, messages: list, pii_config: PiiRedactionConfig | None) -> list:
|
|
recorded = []
|
|
|
|
def record(_thread_id, queued, **_kwargs):
|
|
recorded.append(queued)
|
|
|
|
manager = MagicMock()
|
|
manager.add.side_effect = record
|
|
manager.aadd = AsyncMock(side_effect=record)
|
|
manager.add_nowait.side_effect = record
|
|
monkeypatch.setattr(memory_middleware, "get_memory_manager", lambda: manager)
|
|
monkeypatch.setattr(summarization_hook, "get_memory_manager", lambda: manager)
|
|
monkeypatch.setattr(summarization_hook, "get_memory_config", lambda: MemoryConfig(enabled=True))
|
|
runtime = Runtime(context={"thread_id": "thread-1", "user_id": "alice"})
|
|
middleware = MemoryMiddleware(memory_config=MemoryConfig(enabled=True), pii_redaction_config=pii_config)
|
|
if entry == "after-agent":
|
|
middleware.after_agent({"messages": messages}, runtime)
|
|
elif entry == "async-after-agent":
|
|
asyncio.run(middleware.aafter_agent({"messages": messages}, runtime))
|
|
else:
|
|
event = SummarizationEvent(messages_to_summarize=tuple(messages), preserved_messages=(), thread_id="thread-1", agent_name=None, runtime=runtime)
|
|
summarization_hook.memory_flush_hook(event, pii_redaction_config=pii_config)
|
|
assert len(recorded) == 1
|
|
return recorded[0]
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["provider-invalid", "invalid-checkpoint", "legacy"])
|
|
@pytest.mark.parametrize("entry", ["after-agent", "async-after-agent", "compaction"])
|
|
def test_invalid_and_legacy_call_views_are_redacted_before_memory_admission(monkeypatch, kind: str, entry: str) -> None:
|
|
attempt = _attempt(kind)
|
|
messages = [HumanMessage(f"Look up {_EMAIL}"), attempt, AIMessage("Lookup was not executed")]
|
|
original = messages_to_dict(messages)
|
|
config = PiiRedactionConfig(enabled=True, token_secret=_SECRET)
|
|
|
|
queued = _capture(monkeypatch, entry, messages, config)
|
|
|
|
assert _EMAIL not in str(messages_to_dict(queued))
|
|
token = redact_text(_EMAIL, config)
|
|
assert token in queued[0].content
|
|
if kind == "legacy":
|
|
assert queued[1].additional_kwargs["function_call"]["name"] == "lookup"
|
|
assert token in queued[1].additional_kwargs["function_call"]["arguments"]
|
|
else:
|
|
assert not queued[1].tool_calls
|
|
assert queued[1].invalid_tool_calls[0]["id"] == "call-1"
|
|
assert queued[1].invalid_tool_calls[0]["name"] == "lookup"
|
|
assert token in queued[1].invalid_tool_calls[0]["args"]
|
|
assert token in queued[1].invalid_tool_calls[0]["error"]
|
|
assert queued[1].content == attempt.content
|
|
assert queued[2] is messages[2]
|
|
assert messages_to_dict(messages) == original
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["provider-invalid", "invalid-checkpoint", "legacy"])
|
|
@pytest.mark.parametrize("entry", ["after-agent", "async-after-agent", "compaction"])
|
|
@pytest.mark.parametrize("policy", ["no-config", "disabled", "email-disabled"])
|
|
def test_call_views_follow_existing_redaction_opt_in_and_detector_policy(monkeypatch, kind: str, entry: str, policy: str) -> None:
|
|
config = None if policy == "no-config" else PiiRedactionConfig(enabled=policy != "disabled", redact_email=policy != "email-disabled", token_secret=_SECRET)
|
|
messages = [_attempt(kind)]
|
|
original = messages_to_dict(messages)
|
|
|
|
queued = _capture(monkeypatch, entry, messages, config)
|
|
|
|
assert queued[0] is messages[0]
|
|
assert messages_to_dict(queued) == original
|