276 lines
9.9 KiB
Python
276 lines
9.9 KiB
Python
"""Memory V2 - MemoryManager 与 MemoryProvider 抽象基类测试。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
from deerflow.agents.memory.manager import MemoryManager
|
|
from deerflow.agents.memory.provider import MemoryProvider
|
|
|
|
|
|
class _FakeProvider(MemoryProvider):
|
|
"""测试用的可配置 provider。"""
|
|
|
|
def __init__(self, name: str, *, prompt_block: str = "", prefetch_text: str = ""):
|
|
self._name = name
|
|
self._prompt_block = prompt_block
|
|
self._prefetch_text = prefetch_text
|
|
self.calls: list[str] = []
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self._name
|
|
|
|
def is_available(self) -> bool:
|
|
return True
|
|
|
|
def initialize(self, **kwargs: Any) -> None:
|
|
self.calls.append("initialize")
|
|
|
|
def shutdown(self) -> None:
|
|
self.calls.append("shutdown")
|
|
|
|
def system_prompt_block(self) -> str:
|
|
return self._prompt_block
|
|
|
|
def prefetch(self, query: str, *, thread_id: str = "") -> str:
|
|
self.calls.append(f"prefetch:{query}")
|
|
return self._prefetch_text
|
|
|
|
def sync_turn(self, user_content: str, assistant_content: str, *, thread_id: str = "", turn_index: int | None = None) -> None:
|
|
self.calls.append("sync_turn")
|
|
self.last_turn_index = turn_index
|
|
|
|
def get_tool_schemas(self) -> list[dict[str, Any]]:
|
|
return []
|
|
|
|
def on_session_switch(self, new_thread_id: str, **kwargs: Any) -> None:
|
|
self.calls.append(f"on_session_switch:{new_thread_id}")
|
|
|
|
def on_memory_write(self, action, target, content, metadata=None) -> None:
|
|
self.calls.append(f"on_memory_write:{action}:{target}")
|
|
|
|
|
|
class _ExplodingProvider(_FakeProvider):
|
|
"""每个 hook 都抛异常,用于验证 manager 的隔离性。"""
|
|
|
|
def prefetch(self, query: str, *, thread_id: str = "") -> str:
|
|
raise RuntimeError("boom")
|
|
|
|
def system_prompt_block(self) -> str:
|
|
raise RuntimeError("boom")
|
|
|
|
def sync_turn(self, user_content, assistant_content, *, thread_id: str = "", turn_index: int | None = None) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
def initialize(self, **kwargs: Any) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
|
|
class TestRegistration:
|
|
def test_builtin_always_accepted(self):
|
|
manager = MemoryManager()
|
|
assert manager.add_provider(_FakeProvider("builtin")) is True
|
|
|
|
def test_one_external_accepted(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin"))
|
|
assert manager.add_provider(_FakeProvider("hindsight")) is True
|
|
|
|
def test_second_external_rejected(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin"))
|
|
manager.add_provider(_FakeProvider("hindsight"))
|
|
assert manager.add_provider(_FakeProvider("mem0")) is False
|
|
assert len(manager.providers) == 2
|
|
|
|
def test_knowledge_does_not_consume_external_slot(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin"))
|
|
manager.add_provider(_FakeProvider("knowledge"))
|
|
assert manager.add_provider(_FakeProvider("hindsight")) is True
|
|
assert manager.add_provider(_FakeProvider("mem0")) is False
|
|
assert [p.name for p in manager.providers] == ["builtin", "knowledge", "hindsight"]
|
|
|
|
def test_get_provider(self):
|
|
manager = MemoryManager()
|
|
builtin = _FakeProvider("builtin")
|
|
manager.add_provider(builtin)
|
|
assert manager.get_provider("builtin") is builtin
|
|
assert manager.get_provider("missing") is None
|
|
|
|
|
|
class TestSystemPromptAndPrefetch:
|
|
def test_build_system_prompt_concatenates(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin", prompt_block="块A"))
|
|
manager.add_provider(_FakeProvider("hindsight", prompt_block="块B"))
|
|
result = manager.build_system_prompt()
|
|
assert "块A" in result and "块B" in result
|
|
|
|
def test_build_system_prompt_skips_empty(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin", prompt_block=""))
|
|
assert manager.build_system_prompt() == ""
|
|
|
|
def test_build_system_prompt_can_exclude_only_legacy_knowledge(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin", prompt_block="用户记忆"))
|
|
manager.add_provider(
|
|
_FakeProvider("knowledge", prompt_block="/mnt/knowledge/routing.md")
|
|
)
|
|
manager.add_provider(_FakeProvider("hindsight", prompt_block="Hindsight 说明"))
|
|
|
|
result = manager.build_system_prompt(
|
|
excluded_provider_names={"knowledge"}
|
|
)
|
|
|
|
assert "用户记忆" in result
|
|
assert "Hindsight 说明" in result
|
|
assert "/mnt/knowledge/routing.md" not in result
|
|
|
|
def test_prefetch_all_concatenates(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin", prefetch_text="召回1"))
|
|
manager.add_provider(_FakeProvider("hindsight", prefetch_text="召回2"))
|
|
result = manager.prefetch_all("查询")
|
|
assert "召回1" in result and "召回2" in result
|
|
|
|
def test_build_injection_block_combines_static_and_dynamic(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(
|
|
_FakeProvider("builtin", prompt_block="静态块", prefetch_text="动态块")
|
|
)
|
|
block = manager.build_injection_block("查询")
|
|
assert "静态块" in block and "动态块" in block
|
|
|
|
def test_lead_memory_context_keeps_memory_when_legacy_routing_is_disabled(
|
|
self,
|
|
monkeypatch,
|
|
):
|
|
from deerflow.agents.lead_agent import prompt as prompt_module
|
|
from deerflow.agents.memory import manager as manager_module
|
|
from deerflow.config import memory_config as memory_config_module
|
|
from deerflow.runtime import user_context as user_context_module
|
|
|
|
manager = MemoryManager()
|
|
manager.add_provider(_FakeProvider("builtin", prompt_block="用户偏好记忆"))
|
|
manager.add_provider(
|
|
_FakeProvider("knowledge", prompt_block="读取 /mnt/knowledge/routing.md")
|
|
)
|
|
manager.add_provider(_FakeProvider("hindsight", prompt_block="历史召回说明"))
|
|
|
|
monkeypatch.setattr(
|
|
manager_module,
|
|
"build_memory_manager",
|
|
lambda **_kwargs: manager,
|
|
)
|
|
monkeypatch.setattr(
|
|
memory_config_module,
|
|
"get_effective_memory_config",
|
|
lambda _user_id: SimpleNamespace(enabled=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
user_context_module,
|
|
"get_effective_user_id",
|
|
lambda: "user-1",
|
|
)
|
|
|
|
context = prompt_module._get_memory_context(
|
|
include_legacy_knowledge_routing=False
|
|
)
|
|
|
|
assert "用户偏好记忆" in context
|
|
assert "历史召回说明" in context
|
|
assert "/mnt/knowledge/routing.md" not in context
|
|
|
|
|
|
class TestHookDispatch:
|
|
def test_sync_all(self):
|
|
manager = MemoryManager()
|
|
p = _FakeProvider("builtin")
|
|
manager.add_provider(p)
|
|
manager.sync_all("用户说", "助手答")
|
|
assert "sync_turn" in p.calls
|
|
|
|
def test_sync_all_passes_turn_index(self):
|
|
manager = MemoryManager()
|
|
p = _FakeProvider("builtin")
|
|
manager.add_provider(p)
|
|
manager.sync_all("用户说", "助手答", turn_index=5)
|
|
assert p.last_turn_index == 5
|
|
|
|
def test_on_session_switch_all(self):
|
|
manager = MemoryManager()
|
|
p = _FakeProvider("builtin")
|
|
manager.add_provider(p)
|
|
manager.on_session_switch_all("thread-2")
|
|
assert "on_session_switch:thread-2" in p.calls
|
|
|
|
def test_on_memory_write_skips_builtin(self):
|
|
manager = MemoryManager()
|
|
builtin = _FakeProvider("builtin")
|
|
external = _FakeProvider("hindsight")
|
|
manager.add_provider(builtin)
|
|
manager.add_provider(external)
|
|
manager.on_memory_write("add", "user", "内容")
|
|
# builtin 是写入源,不应收到镜像通知
|
|
assert not any(c.startswith("on_memory_write") for c in builtin.calls)
|
|
assert "on_memory_write:add:user" in external.calls
|
|
|
|
def test_initialize_and_shutdown_all(self):
|
|
manager = MemoryManager()
|
|
p = _FakeProvider("builtin")
|
|
manager.add_provider(p)
|
|
manager.initialize_all(user_id="u1")
|
|
manager.shutdown_all()
|
|
assert "initialize" in p.calls
|
|
assert "shutdown" in p.calls
|
|
|
|
|
|
class TestFailureIsolation:
|
|
def test_exploding_provider_does_not_break_others(self):
|
|
manager = MemoryManager()
|
|
manager.add_provider(_ExplodingProvider("builtin"))
|
|
good = _FakeProvider("hindsight", prompt_block="正常块", prefetch_text="正常召回")
|
|
manager.add_provider(good)
|
|
|
|
# 一个 provider 抛异常,另一个仍正常工作
|
|
manager.initialize_all(user_id="u1")
|
|
assert manager.build_system_prompt() == "正常块"
|
|
assert manager.prefetch_all("q") == "正常召回"
|
|
manager.sync_all("u", "a") # 不抛异常
|
|
|
|
|
|
class _MinimalProvider(MemoryProvider):
|
|
"""只实现 abstractmethod 的最小 provider,用于验证基类默认 hook。"""
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "minimal"
|
|
|
|
def is_available(self) -> bool:
|
|
return True
|
|
|
|
def initialize(self, **kwargs: Any) -> None:
|
|
pass
|
|
|
|
def get_tool_schemas(self) -> list[dict[str, Any]]:
|
|
return []
|
|
|
|
|
|
class TestProviderDefaults:
|
|
"""MemoryProvider 抽象基类的默认 hook 实现。"""
|
|
|
|
def test_default_hooks_are_noop(self):
|
|
provider = _MinimalProvider()
|
|
# 未重写的 hook 应有安全的默认实现,不抛异常
|
|
assert provider.prefetch("q") == ""
|
|
assert provider.system_prompt_block() == ""
|
|
provider.queue_prefetch("q")
|
|
provider.sync_turn("u", "a")
|
|
provider.shutdown()
|
|
provider.on_session_switch("t")
|
|
provider.on_memory_write("add", "user", "c")
|