@
fix(memory): 修复 Hindsight provider 生命周期的三处缺陷
HindsightProvider 被 build_memory_manager 每轮重新构造,以下三个问题
都源于这个无状态特性未被正确处理:
1. document_id 每轮漂移 —— 原 _reset_retain_batch 用微秒时间戳拼接,
导致同一 thread 的每轮 retain 落到不同文档,文档数随对话轮数线性膨胀。
改为按 thread_id 确定性派生 (thread-{thread_id}),同一 thread 累积成一份。
2. retain_every_n_turns 失效 —— 进程内 _turn_counter 每轮被 initialize
重置,retain_every_n_turns > 1 时永不触发。新增 turn_index 参数贯穿
sync_turn/sync_all,由 MemoryMiddleware 按对话历史 (count_user_turns)
算出并传入;_turn_counter 仅作缺省退路。
3. aiohttp ClientSession 泄漏 —— 缓存的客户端从不 shutdown,且 _run_coro
每次新建事件循环使缓存客户端失效,触发 "Unclosed client session"。
改为客户端即开即关:_run_with_client 在 finally 保证关闭,_client_factory
作为可覆盖的测试注入点。
测试:test_hindsight_provider.py 全量重写 (含 TestDocumentId /
TestRetainEveryNTurnsAcrossRecreation / TestClientLifecycle),
新增 test_memory_message_processing.py,test_memory_manager.py 补
turn_index 透传用例。
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@
This commit is contained in:
xuhd committed
1 parent
8f88427d4e
commit
dba821a70e
8 files changed
+770
-61
No files matched your search
+11
-2
@@ -153,11 +153,20 @@ class MemoryManager:
|
||||
assistant_content: str,
|
||||
*,
|
||||
thread_id: str = "",
|
||||
turn_index: int | None = None,
|
||||
) -> None:
|
||||
"""把一轮完成的对话同步到所有 provider。"""
|
||||
"""把一轮完成的对话同步到所有 provider。
|
||||
|
||||
``turn_index`` 是当前对话的轮次序号,透传给各 provider 用于轮次节流。
|
||||
"""
|
||||
for provider in self._providers:
|
||||
try:
|
||||
provider.sync_turn(user_content, assistant_content, thread_id=thread_id)
|
||||
provider.sync_turn(
|
||||
user_content,
|
||||
assistant_content,
|
||||
thread_id=thread_id,
|
||||
turn_index=turn_index,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("记忆 provider '%s' sync_turn 失败:%s", provider.name, exc)
|
||||
|
||||
|
||||
+10
@@ -29,6 +29,16 @@ def extract_message_text(message: Any) -> str:
|
||||
return str(content)
|
||||
|
||||
|
||||
def count_user_turns(filtered_messages: list[Any]) -> int:
|
||||
"""统计已过滤消息里的用户轮次数 (= human 消息条数)。
|
||||
|
||||
供 MemoryMiddleware 算出当前轮次序号传给 provider 做轮次节流 ——
|
||||
provider 每轮被重建、无法靠自身计数,必须由这里按对话历史推算。
|
||||
入参应是 :func:`filter_messages_for_memory` 的输出。
|
||||
"""
|
||||
return sum(1 for msg in filtered_messages if getattr(msg, "type", None) == "human")
|
||||
|
||||
|
||||
def filter_messages_for_memory(messages: list[Any]) -> list[Any]:
|
||||
"""Keep only user inputs and final assistant responses for memory updates."""
|
||||
filtered = []
|
||||
|
||||
@@ -117,11 +117,17 @@ class MemoryProvider(ABC):
|
||||
assistant_content: str,
|
||||
*,
|
||||
thread_id: str = "",
|
||||
turn_index: int | None = None,
|
||||
) -> None:
|
||||
"""每轮对话结束后把对话持久化到后端。
|
||||
|
||||
实现应当非阻塞 —— 如果后端有延迟,排队到后台处理而不是直接调用。
|
||||
默认 no-op。
|
||||
|
||||
``turn_index`` 是当前对话的轮次序号 (从 1 开始,由编排层按对话历史算出)。
|
||||
provider 通常被每轮重新构造、无法靠自身计数,需要按轮次节流写入
|
||||
(例如 "每 N 轮 retain 一次") 的实现应当用它而非进程内计数器。
|
||||
缺省 ``None`` 表示调用方未提供。
|
||||
"""
|
||||
|
||||
# ---- LLM 工具 --------------------------------------------------------
|
||||
|
||||
+122
-57
@@ -17,6 +17,21 @@ HTTP 客户端方式接入一个用户自起的 Hindsight Docker 实例 (``local
|
||||
|
||||
召回:``prefetch`` 并发查两个 bank,合并结果。
|
||||
|
||||
document_id 与轮次计数
|
||||
======================
|
||||
本 provider 被 ``build_memory_manager`` 每轮重新构造 —— 是无状态的。因此:
|
||||
- ``document_id`` 由 ``thread_id`` **确定性**派生,同一 thread 的所有轮次 retain
|
||||
都落到同一个 Hindsight 文档,累积而非每轮碎片化成新文档。
|
||||
- ``retain_every_n_turns`` 的轮次判定依赖外部传入的 ``turn_index``
|
||||
(由 ``MemoryMiddleware`` 按对话历史算出);未传时退回进程内 ``_turn_counter``。
|
||||
|
||||
客户端生命周期
|
||||
==============
|
||||
``_run_coro`` 每次操作新建一个事件循环。aiohttp ``ClientSession`` 是
|
||||
loop-bound 的,跨 loop 复用会泄漏并失效。因此本 provider **不缓存客户端** ——
|
||||
每次记忆操作在 :meth:`_run_with_client` 里即开即关,保证 session 在其所属的
|
||||
loop 内被正常关闭。
|
||||
|
||||
降级
|
||||
====
|
||||
``hindsight-client`` 是可选依赖。未安装或未配置 ``api_url`` 时 :meth:`is_available`
|
||||
@@ -90,7 +105,6 @@ class HindsightProvider(MemoryProvider):
|
||||
|
||||
def __init__(self, config: HindsightMemoryConfig | None = None) -> None:
|
||||
self._config = config or HindsightMemoryConfig()
|
||||
self._client: Any = None
|
||||
|
||||
self._user_id: str = ""
|
||||
self._agent_id: str = "default"
|
||||
@@ -129,7 +143,7 @@ class HindsightProvider(MemoryProvider):
|
||||
base_dir: str = "",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""解析双 bank 名称,准备 retain 批次状态。客户端按需懒创建。"""
|
||||
"""解析双 bank 名称,准备 retain 批次状态。客户端每次操作即开即关。"""
|
||||
self._user_id = user_id or "default"
|
||||
self._agent_id = (agent_id or "default").lower()
|
||||
self._thread_id = thread_id or ""
|
||||
@@ -154,41 +168,63 @@ class HindsightProvider(MemoryProvider):
|
||||
return name.strip("-") or "deerflow"
|
||||
|
||||
def _reset_retain_batch(self) -> None:
|
||||
"""重置 retain 批次状态 (新的 document_id + 轮次清零)。"""
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||||
self._document_id = f"{self._thread_id or 'thread'}-{ts}"
|
||||
"""重置 retain 批次状态。
|
||||
|
||||
``document_id`` 由 ``thread_id`` **确定性**派生 —— provider 虽然被每轮
|
||||
重新构造,同一 thread 的所有轮次仍落到同一个 Hindsight 文档,累积成一份
|
||||
而不是碎片化成一轮一个文档。无 thread 时 (例如 memory 工具镜像) 用统一
|
||||
的兜底 id。轮次计数清零,仅作为 ``turn_index`` 缺失时的退路。
|
||||
"""
|
||||
self._document_id = f"thread-{self._thread_id}" if self._thread_id else "thread-default"
|
||||
self._turn_counter = 0
|
||||
|
||||
def _get_client(self) -> Any:
|
||||
"""懒创建并缓存 Hindsight HTTP 客户端。"""
|
||||
if self._client is None:
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"base_url": self._config.api_url,
|
||||
"timeout": 30.0,
|
||||
}
|
||||
if self._config.api_key:
|
||||
kwargs["api_key"] = self._config.api_key
|
||||
self._client = Hindsight(**kwargs)
|
||||
return self._client
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""关闭客户端 (如果客户端暴露了 close)。"""
|
||||
client = self._client
|
||||
self._client = None
|
||||
if client is None:
|
||||
return
|
||||
"""无需清理 —— 客户端在每次操作内即开即关,不缓存。"""
|
||||
|
||||
# ---- 客户端生命周期 --------------------------------------------------
|
||||
|
||||
def _client_factory(self) -> Any:
|
||||
"""创建一个新的 Hindsight HTTP 客户端。
|
||||
|
||||
每次记忆操作都新建一个、用完即关,**不缓存** —— aiohttp ClientSession
|
||||
是 loop-bound 的,而 ``_run_coro`` 每次新建事件循环。测试可覆盖本方法
|
||||
注入 fake 客户端。
|
||||
"""
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"base_url": self._config.api_url,
|
||||
"timeout": 30.0,
|
||||
}
|
||||
if self._config.api_key:
|
||||
kwargs["api_key"] = self._config.api_key
|
||||
return Hindsight(**kwargs)
|
||||
|
||||
@staticmethod
|
||||
async def _aclose_client(client: Any) -> None:
|
||||
"""关闭客户端 (兼容暴露 ``aclose`` 或 ``close`` 的实现)。"""
|
||||
close = getattr(client, "aclose", None) or getattr(client, "close", None)
|
||||
if close is None:
|
||||
return
|
||||
try:
|
||||
result = close()
|
||||
if asyncio.iscoroutine(result):
|
||||
_run_coro(result, timeout=10)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
await result
|
||||
except Exception as exc: # noqa: BLE001 - 关闭失败不影响结果
|
||||
logger.debug("Hindsight 客户端关闭失败:%s", exc)
|
||||
|
||||
async def _run_with_client(self, op: Any) -> Any:
|
||||
"""在当前事件循环内创建客户端、执行 ``op(client)``、保证关闭客户端。
|
||||
|
||||
``op`` 是一个接收 client、返回 awaitable 的协程函数。无论成功、异常还是
|
||||
被取消 (超时),``finally`` 都会关闭客户端 —— 杜绝 ``Unclosed client session``。
|
||||
"""
|
||||
client = self._client_factory()
|
||||
try:
|
||||
return await op(client)
|
||||
finally:
|
||||
await self._aclose_client(client)
|
||||
|
||||
# ---- 系统提示词块 -----------------------------------------------------
|
||||
|
||||
def system_prompt_block(self) -> str:
|
||||
@@ -226,13 +262,11 @@ class HindsightProvider(MemoryProvider):
|
||||
query = query[: self._config.recall_max_input_chars]
|
||||
|
||||
try:
|
||||
user_results = self._recall_bank(self._user_bank, query)
|
||||
work_results = self._recall_bank(self._work_bank, query)
|
||||
merged = self._recall_both(query)
|
||||
except Exception as exc: # noqa: BLE001 - 召回失败不能阻断对话
|
||||
logger.debug("Hindsight 召回失败 (非致命):%s", exc)
|
||||
return ""
|
||||
|
||||
merged = self._merge_results(user_results, work_results)
|
||||
if not merged:
|
||||
return ""
|
||||
lines = "\n".join(f"- {text}" for text in merged)
|
||||
@@ -242,18 +276,34 @@ class HindsightProvider(MemoryProvider):
|
||||
f"{lines}"
|
||||
)
|
||||
|
||||
def _recall_bank(self, bank_id: str, query: str) -> list[tuple[float, str]]:
|
||||
"""召回单个 bank,返回 ``[(score, text), ...]``。"""
|
||||
client = self._get_client()
|
||||
resp = _run_coro(
|
||||
client.arecall(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=self._config.recall_budget,
|
||||
max_tokens=self._config.recall_max_tokens,
|
||||
),
|
||||
timeout=15,
|
||||
def _recall_both(self, query: str) -> list[str]:
|
||||
"""并发召回 user + work 两个 bank,合并去重。一次操作共用一个客户端。"""
|
||||
|
||||
async def _op(client: Any) -> tuple[Any, Any]:
|
||||
return await asyncio.gather(
|
||||
client.arecall(
|
||||
bank_id=self._user_bank,
|
||||
query=query,
|
||||
budget=self._config.recall_budget,
|
||||
max_tokens=self._config.recall_max_tokens,
|
||||
),
|
||||
client.arecall(
|
||||
bank_id=self._work_bank,
|
||||
query=query,
|
||||
budget=self._config.recall_budget,
|
||||
max_tokens=self._config.recall_max_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
user_resp, work_resp = _run_coro(self._run_with_client(_op), timeout=15)
|
||||
return self._merge_results(
|
||||
self._extract_results(user_resp),
|
||||
self._extract_results(work_resp),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _extract_results(resp: Any) -> list[tuple[float, str]]:
|
||||
"""从一个 recall 响应里抽出 ``[(score, text), ...]``。"""
|
||||
results = getattr(resp, "results", None) or []
|
||||
out: list[tuple[float, str]] = []
|
||||
for r in results:
|
||||
@@ -285,15 +335,28 @@ class HindsightProvider(MemoryProvider):
|
||||
assistant_content: str,
|
||||
*,
|
||||
thread_id: str = "",
|
||||
turn_index: int | None = None,
|
||||
) -> None:
|
||||
"""自动 retain 一轮对话到 work bank。
|
||||
|
||||
``auto_retain=False`` 时不写。按 ``retain_every_n_turns`` 控制频率。
|
||||
|
||||
轮次判定优先用调用方传入的 ``turn_index`` (由 MemoryMiddleware 按对话历史
|
||||
算出) —— 因为本 provider 每轮被重建、进程内 ``_turn_counter`` 会被重置,
|
||||
单靠它在 ``retain_every_n_turns > 1`` 时会永远不触发。``turn_index`` 缺失
|
||||
时才退回 ``_turn_counter``。
|
||||
"""
|
||||
if not self._config.auto_retain:
|
||||
return
|
||||
self._turn_counter += 1
|
||||
if self._turn_counter % self._config.retain_every_n_turns != 0:
|
||||
|
||||
if turn_index is not None:
|
||||
counter = turn_index
|
||||
else:
|
||||
self._turn_counter += 1
|
||||
counter = self._turn_counter
|
||||
|
||||
every = max(1, self._config.retain_every_n_turns)
|
||||
if counter % every != 0:
|
||||
return
|
||||
|
||||
turn_text = (
|
||||
@@ -330,8 +393,7 @@ class HindsightProvider(MemoryProvider):
|
||||
logger.debug("Hindsight on_memory_write 镜像失败:%s", exc)
|
||||
|
||||
def _retain(self, bank_id: str, content: str, *, kind: str) -> None:
|
||||
"""向某 bank 写入一条记忆。"""
|
||||
client = self._get_client()
|
||||
"""向某 bank 写入一条记忆。客户端即开即关。"""
|
||||
metadata = {
|
||||
"retained_at": _utc_now_iso(),
|
||||
"kind": kind,
|
||||
@@ -340,33 +402,36 @@ class HindsightProvider(MemoryProvider):
|
||||
}
|
||||
if self._thread_id:
|
||||
metadata["thread_id"] = self._thread_id
|
||||
_run_coro(
|
||||
client.aretain(
|
||||
|
||||
async def _op(client: Any) -> Any:
|
||||
return await client.aretain(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context=self._config.bank_retain_mission or None,
|
||||
metadata=metadata,
|
||||
document_id=self._document_id,
|
||||
retain_async=self._config.retain_async,
|
||||
),
|
||||
timeout=15,
|
||||
)
|
||||
)
|
||||
|
||||
_run_coro(self._run_with_client(_op), timeout=15)
|
||||
|
||||
def reflect(self, query: str, *, use_user_bank: bool = True) -> str:
|
||||
"""对某 bank 做跨记忆综合 (LLM 推理)。供 hindsight_reflect 工具调用。"""
|
||||
bank = self._user_bank if use_user_bank else self._work_bank
|
||||
client = self._get_client()
|
||||
resp = _run_coro(
|
||||
client.areflect(bank_id=bank, query=query, budget=self._config.recall_budget),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
async def _op(client: Any) -> Any:
|
||||
return await client.areflect(
|
||||
bank_id=bank,
|
||||
query=query,
|
||||
budget=self._config.recall_budget,
|
||||
)
|
||||
|
||||
resp = _run_coro(self._run_with_client(_op), timeout=30)
|
||||
return getattr(resp, "text", "") or ""
|
||||
|
||||
def recall(self, query: str) -> list[str]:
|
||||
"""召回双 bank 并返回文本列表。供 hindsight_recall 工具调用。"""
|
||||
user_results = self._recall_bank(self._user_bank, query)
|
||||
work_results = self._recall_bank(self._work_bank, query)
|
||||
return self._merge_results(user_results, work_results)
|
||||
return self._recall_both(query)
|
||||
|
||||
def retain_to_user_bank(self, content: str) -> None:
|
||||
"""主动写入 user bank。供 hindsight_retain 工具调用。"""
|
||||
|
||||
+11
-2
@@ -27,7 +27,7 @@ from langgraph.config import get_config
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.agents.memory.manager import build_memory_manager
|
||||
from deerflow.agents.memory.message_processing import filter_messages_for_memory
|
||||
from deerflow.agents.memory.message_processing import count_user_turns, filter_messages_for_memory
|
||||
from deerflow.config.memory_config import get_effective_memory_config, get_memory_config
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
|
||||
@@ -221,6 +221,10 @@ class MemoryMiddleware(AgentMiddleware[MemoryMiddlewareState]):
|
||||
if not user_text or not assistant_text:
|
||||
return None
|
||||
|
||||
# 轮次序号按对话历史算出 —— provider 每轮被重建、无法自行计数,
|
||||
# retain_every_n_turns 的节流必须靠这个外部传入的值。
|
||||
turn_index = count_user_turns(filtered_messages)
|
||||
|
||||
thread_id = self._resolve_thread_id(runtime)
|
||||
try:
|
||||
manager = build_memory_manager(
|
||||
@@ -229,7 +233,12 @@ class MemoryMiddleware(AgentMiddleware[MemoryMiddlewareState]):
|
||||
thread_id=thread_id,
|
||||
config=effective_config,
|
||||
)
|
||||
manager.sync_all(user_text, assistant_text, thread_id=thread_id or "")
|
||||
manager.sync_all(
|
||||
user_text,
|
||||
assistant_text,
|
||||
thread_id=thread_id or "",
|
||||
turn_index=turn_index,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - 写入失败不能阻断对话
|
||||
logger.warning("Hindsight retain 失败:%s", exc)
|
||||
|
||||
|
||||
@@ -0,0 +1,363 @@
|
||||
"""Memory V2 - HindsightProvider 测试 (mock 客户端,不连真实 Hindsight)。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deerflow.agents.memory.providers.hindsight import HindsightProvider
|
||||
from deerflow.config.memory_config import HindsightMemoryConfig
|
||||
|
||||
|
||||
class _FakeResult:
|
||||
def __init__(self, text: str, score: float):
|
||||
self.text = text
|
||||
self.score = score
|
||||
|
||||
|
||||
class _FakeRecallResponse:
|
||||
def __init__(self, results):
|
||||
self.results = results
|
||||
|
||||
|
||||
class _FakeReflectResponse:
|
||||
def __init__(self, text):
|
||||
self.text = text
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
"""模拟 hindsight_client.Hindsight 的异步 API。"""
|
||||
|
||||
def __init__(self):
|
||||
self.retained: list[dict] = []
|
||||
self.recall_by_bank: dict[str, list[_FakeResult]] = {}
|
||||
self.reflect_text = "综合结论"
|
||||
self.closed = False
|
||||
|
||||
async def arecall(self, *, bank_id, query, budget, max_tokens):
|
||||
return _FakeRecallResponse(self.recall_by_bank.get(bank_id, []))
|
||||
|
||||
async def aretain(self, *, bank_id, content, context=None, metadata=None, document_id=None, retain_async=True):
|
||||
self.retained.append(
|
||||
{
|
||||
"bank_id": bank_id,
|
||||
"content": content,
|
||||
"metadata": metadata,
|
||||
"document_id": document_id,
|
||||
}
|
||||
)
|
||||
|
||||
async def areflect(self, *, bank_id, query, budget):
|
||||
return _FakeReflectResponse(self.reflect_text)
|
||||
|
||||
async def aclose(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _make_provider(client: _FakeClient | None = None, **config_kwargs) -> HindsightProvider:
|
||||
config = HindsightMemoryConfig(api_url="http://localhost:8765", **config_kwargs)
|
||||
provider = HindsightProvider(config)
|
||||
provider.initialize(user_id="u1", agent_id="agent-x", thread_id="t1", base_dir="")
|
||||
if client is not None:
|
||||
# 覆盖客户端工厂,注入 fake —— 绕过真实 hindsight_client 导入。
|
||||
provider._client_factory = lambda: client
|
||||
return provider
|
||||
|
||||
|
||||
class TestAvailability:
|
||||
def test_unavailable_when_package_missing(self, monkeypatch):
|
||||
# 显式模拟 hindsight-client 未安装 (不依赖测试环境是否真装了)
|
||||
monkeypatch.setattr(
|
||||
"deerflow.agents.memory.providers.hindsight.hindsight_client_available",
|
||||
lambda: False,
|
||||
)
|
||||
provider = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
assert provider.is_available() is False
|
||||
|
||||
def test_unavailable_when_no_api_url(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"deerflow.agents.memory.providers.hindsight.hindsight_client_available",
|
||||
lambda: True,
|
||||
)
|
||||
provider = HindsightProvider(HindsightMemoryConfig(api_url=""))
|
||||
assert provider.is_available() is False
|
||||
|
||||
def test_available_when_package_and_url_present(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"deerflow.agents.memory.providers.hindsight.hindsight_client_available",
|
||||
lambda: True,
|
||||
)
|
||||
provider = HindsightProvider(HindsightMemoryConfig(api_url="http://localhost:8765"))
|
||||
assert provider.is_available() is True
|
||||
|
||||
|
||||
class TestBankResolution:
|
||||
def test_user_and_work_banks(self):
|
||||
provider = _make_provider()
|
||||
assert provider._user_bank == "deerflow-user-u1"
|
||||
assert provider._work_bank == "deerflow-work-u1-agent-x"
|
||||
|
||||
def test_empty_placeholder_collapses(self):
|
||||
provider = HindsightProvider(
|
||||
HindsightMemoryConfig(
|
||||
api_url="http://x",
|
||||
work_bank_id_template="deerflow-{thread_id}-work",
|
||||
)
|
||||
)
|
||||
provider.initialize(user_id="u1", agent_id="a", thread_id="", base_dir="")
|
||||
# 空 thread_id 占位符塌缩,不留多余连字符
|
||||
assert "--" not in provider._work_bank
|
||||
assert provider._work_bank == "deerflow-work"
|
||||
|
||||
|
||||
class TestDocumentId:
|
||||
"""document_id 必须由 thread_id 确定性派生 —— provider 每轮重建后仍稳定。"""
|
||||
|
||||
def test_deterministic_from_thread_id(self):
|
||||
# 同一 thread,两个独立构造的 provider 必须算出同一个 document_id
|
||||
p1 = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
p1.initialize(user_id="u", agent_id="a", thread_id="t-123", base_dir="")
|
||||
p2 = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
p2.initialize(user_id="u", agent_id="a", thread_id="t-123", base_dir="")
|
||||
assert p1._document_id == p2._document_id
|
||||
assert "t-123" in p1._document_id
|
||||
|
||||
def test_different_threads_get_different_doc(self):
|
||||
p1 = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
p1.initialize(user_id="u", agent_id="a", thread_id="t-1", base_dir="")
|
||||
p2 = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
p2.initialize(user_id="u", agent_id="a", thread_id="t-2", base_dir="")
|
||||
assert p1._document_id != p2._document_id
|
||||
|
||||
def test_no_thread_id_uses_fallback(self):
|
||||
p = HindsightProvider(HindsightMemoryConfig(api_url="http://x"))
|
||||
p.initialize(user_id="u", agent_id="a", thread_id="", base_dir="")
|
||||
assert p._document_id == "thread-default"
|
||||
|
||||
def test_retain_writes_thread_document_id(self):
|
||||
# sync_turn 落盘时带的 document_id 应是 thread 确定性 id,不含时间戳
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client) # thread_id="t1"
|
||||
provider.sync_turn("用户说啥", "助手答啥")
|
||||
assert client.retained[0]["document_id"] == "thread-t1"
|
||||
|
||||
def test_recreated_providers_share_one_document(self):
|
||||
# 模拟真实场景:同一 thread 连续 3 轮,每轮 provider 重建,
|
||||
# 三轮 retain 必须落到同一个文档。
|
||||
doc_ids: set[str] = set()
|
||||
for turn in range(1, 4):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client) # 每轮重建,thread_id 固定 t1
|
||||
provider.sync_turn(f"u{turn}", f"a{turn}")
|
||||
doc_ids.add(client.retained[0]["document_id"])
|
||||
assert doc_ids == {"thread-t1"}
|
||||
|
||||
|
||||
class TestPrefetchRecall:
|
||||
def test_prefetch_merges_two_banks(self):
|
||||
client = _FakeClient()
|
||||
client.recall_by_bank["deerflow-user-u1"] = [_FakeResult("用户事实", 0.9)]
|
||||
client.recall_by_bank["deerflow-work-u1-agent-x"] = [_FakeResult("工作笔记", 0.7)]
|
||||
provider = _make_provider(client)
|
||||
result = provider.prefetch("查询")
|
||||
assert "用户事实" in result
|
||||
assert "工作笔记" in result
|
||||
# 高 score 在前
|
||||
assert result.index("用户事实") < result.index("工作笔记")
|
||||
|
||||
def test_prefetch_dedupes(self):
|
||||
client = _FakeClient()
|
||||
client.recall_by_bank["deerflow-user-u1"] = [_FakeResult("同一条", 0.9)]
|
||||
client.recall_by_bank["deerflow-work-u1-agent-x"] = [_FakeResult("同一条", 0.5)]
|
||||
provider = _make_provider(client)
|
||||
result = provider.prefetch("查询")
|
||||
assert result.count("同一条") == 1
|
||||
|
||||
def test_prefetch_empty_query_returns_empty(self):
|
||||
provider = _make_provider(_FakeClient())
|
||||
assert provider.prefetch("") == ""
|
||||
|
||||
def test_prefetch_skipped_in_tools_mode(self):
|
||||
client = _FakeClient()
|
||||
client.recall_by_bank["deerflow-user-u1"] = [_FakeResult("x", 0.9)]
|
||||
provider = _make_provider(client, memory_mode="tools")
|
||||
assert provider.prefetch("查询") == ""
|
||||
|
||||
def test_prefetch_skipped_when_auto_recall_off(self):
|
||||
client = _FakeClient()
|
||||
client.recall_by_bank["deerflow-user-u1"] = [_FakeResult("x", 0.9)]
|
||||
provider = _make_provider(client, auto_recall=False)
|
||||
assert provider.prefetch("查询") == ""
|
||||
|
||||
def test_prefetch_no_results_returns_empty(self):
|
||||
provider = _make_provider(_FakeClient())
|
||||
assert provider.prefetch("查询") == ""
|
||||
|
||||
def test_recall_returns_list(self):
|
||||
client = _FakeClient()
|
||||
client.recall_by_bank["deerflow-user-u1"] = [_FakeResult("甲", 0.9)]
|
||||
provider = _make_provider(client)
|
||||
assert provider.recall("查询") == ["甲"]
|
||||
|
||||
|
||||
class TestSyncTurnRetain:
|
||||
def test_sync_turn_retains_to_work_bank(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.sync_turn("用户说啥", "助手答啥")
|
||||
assert len(client.retained) == 1
|
||||
assert client.retained[0]["bank_id"] == "deerflow-work-u1-agent-x"
|
||||
assert "用户说啥" in client.retained[0]["content"]
|
||||
|
||||
def test_sync_turn_skipped_when_auto_retain_off(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client, auto_retain=False)
|
||||
provider.sync_turn("u", "a")
|
||||
assert client.retained == []
|
||||
|
||||
def test_retain_every_n_turns_internal_counter(self):
|
||||
# 不传 turn_index 时退回进程内计数器 (同一 provider 复用的场景)
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client, retain_every_n_turns=2)
|
||||
provider.sync_turn("u1", "a1") # 第 1 轮,缓冲
|
||||
assert client.retained == []
|
||||
provider.sync_turn("u2", "a2") # 第 2 轮,触发
|
||||
assert len(client.retained) == 1
|
||||
|
||||
|
||||
class TestRetainEveryNTurnsAcrossRecreation:
|
||||
"""provider 每轮重建时,retain_every_n_turns 的节流必须靠外部 turn_index。"""
|
||||
|
||||
def test_internal_counter_breaks_across_recreation(self):
|
||||
# 文档化 bug 根因:每个新 provider 计数器从 0→1,1 % 3 != 0 → 永不触发
|
||||
retained_total = 0
|
||||
for turn in range(1, 7): # 6 轮
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client, retain_every_n_turns=3)
|
||||
provider.sync_turn(f"u{turn}", f"a{turn}") # 不传 turn_index
|
||||
retained_total += len(client.retained)
|
||||
assert retained_total == 0
|
||||
|
||||
def test_explicit_turn_index_drives_retain(self):
|
||||
# 修复:turn_index 由编排层按对话历史算出并传入,节流恢复正常
|
||||
retained_total = 0
|
||||
for turn in range(1, 7): # 6 轮,retain_every_n_turns=3
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client, retain_every_n_turns=3)
|
||||
provider.sync_turn(f"u{turn}", f"a{turn}", turn_index=turn)
|
||||
retained_total += len(client.retained)
|
||||
# 第 3、6 轮触发 → 共 2 次
|
||||
assert retained_total == 2
|
||||
|
||||
def test_turn_index_every_one_retains_each_turn(self):
|
||||
# retain_every_n_turns=1 时每轮都写
|
||||
for turn in range(1, 4):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client, retain_every_n_turns=1)
|
||||
provider.sync_turn("u", "a", turn_index=turn)
|
||||
assert len(client.retained) == 1
|
||||
|
||||
|
||||
class TestClientLifecycle:
|
||||
"""客户端必须即开即关 —— 杜绝 aiohttp 'Unclosed client session'。"""
|
||||
|
||||
def test_client_closed_after_recall(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.recall("查询")
|
||||
assert client.closed is True
|
||||
|
||||
def test_client_closed_after_retain(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.sync_turn("u", "a")
|
||||
assert client.closed is True
|
||||
|
||||
def test_client_closed_after_reflect(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.reflect("问题")
|
||||
assert client.closed is True
|
||||
|
||||
def test_client_closed_even_on_failure(self):
|
||||
class _BrokenClient(_FakeClient):
|
||||
async def arecall(self, **kwargs):
|
||||
raise RuntimeError("Hindsight 不可达")
|
||||
|
||||
client = _BrokenClient()
|
||||
provider = _make_provider(client)
|
||||
provider.prefetch("查询") # 失败被吞
|
||||
assert client.closed is True
|
||||
|
||||
def test_shutdown_is_noop(self):
|
||||
# 无缓存客户端,shutdown 是 no-op,不应抛异常
|
||||
provider = HindsightProvider(HindsightMemoryConfig())
|
||||
provider.shutdown()
|
||||
|
||||
|
||||
class TestOnMemoryWrite:
|
||||
def test_user_target_goes_to_user_bank(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.on_memory_write("add", "user", "用户偏好")
|
||||
assert client.retained[0]["bank_id"] == "deerflow-user-u1"
|
||||
|
||||
def test_memory_target_goes_to_work_bank(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.on_memory_write("add", "memory", "工作笔记")
|
||||
assert client.retained[0]["bank_id"] == "deerflow-work-u1-agent-x"
|
||||
|
||||
def test_remove_action_not_mirrored(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.on_memory_write("remove", "user", "某条")
|
||||
assert client.retained == []
|
||||
|
||||
def test_empty_content_not_mirrored(self):
|
||||
client = _FakeClient()
|
||||
provider = _make_provider(client)
|
||||
provider.on_memory_write("add", "user", " ")
|
||||
assert client.retained == []
|
||||
|
||||
|
||||
class TestReflect:
|
||||
def test_reflect_returns_text(self):
|
||||
client = _FakeClient()
|
||||
client.reflect_text = "这是综合结论"
|
||||
provider = _make_provider(client)
|
||||
assert provider.reflect("问题") == "这是综合结论"
|
||||
|
||||
|
||||
class TestSessionSwitch:
|
||||
def test_session_switch_resets_work_bank_and_batch(self):
|
||||
provider = _make_provider(_FakeClient())
|
||||
old_doc = provider._document_id
|
||||
provider.on_session_switch("t2")
|
||||
assert provider._thread_id == "t2"
|
||||
assert provider._turn_counter == 0
|
||||
assert provider._document_id != old_doc
|
||||
assert provider._document_id == "thread-t2"
|
||||
|
||||
|
||||
class TestMisc:
|
||||
def test_name(self):
|
||||
assert HindsightProvider(HindsightMemoryConfig()).name == "hindsight"
|
||||
|
||||
def test_get_tool_schemas_empty(self):
|
||||
# Hindsight 工具走 LangChain BaseTool,不经 manager
|
||||
assert HindsightProvider(HindsightMemoryConfig()).get_tool_schemas() == []
|
||||
|
||||
def test_context_mode_no_system_prompt(self):
|
||||
provider = _make_provider(memory_mode="context")
|
||||
assert provider.system_prompt_block() == ""
|
||||
|
||||
def test_tools_mode_has_system_prompt(self):
|
||||
provider = _make_provider(memory_mode="tools")
|
||||
assert "hindsight_recall" in provider.system_prompt_block()
|
||||
|
||||
def test_recall_failure_returns_empty(self):
|
||||
class _BrokenClient(_FakeClient):
|
||||
async def arecall(self, **kwargs):
|
||||
raise RuntimeError("Hindsight 不可达")
|
||||
|
||||
provider = _make_provider(_BrokenClient())
|
||||
# 召回失败不抛异常,返回空串
|
||||
assert provider.prefetch("查询") == ""
|
||||
@@ -0,0 +1,210 @@
|
||||
"""Memory V2 - MemoryManager 与 MemoryProvider 抽象基类测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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_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_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
|
||||
|
||||
|
||||
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")
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Memory V2 - message_processing 辅助函数测试。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from deerflow.agents.memory.message_processing import count_user_turns, filter_messages_for_memory
|
||||
|
||||
|
||||
class TestCountUserTurns:
|
||||
def test_empty(self):
|
||||
assert count_user_turns([]) == 0
|
||||
|
||||
def test_counts_only_human_messages(self):
|
||||
messages = [
|
||||
HumanMessage(content="问题一"),
|
||||
AIMessage(content="回答一"),
|
||||
HumanMessage(content="问题二"),
|
||||
AIMessage(content="回答二"),
|
||||
]
|
||||
assert count_user_turns(messages) == 2
|
||||
|
||||
def test_ignores_ai_messages(self):
|
||||
messages = [AIMessage(content="a"), AIMessage(content="b")]
|
||||
assert count_user_turns(messages) == 0
|
||||
|
||||
def test_works_on_filtered_output(self):
|
||||
# count_user_turns 设计上吃 filter_messages_for_memory 的输出
|
||||
raw = [
|
||||
HumanMessage(content="第一轮"),
|
||||
AIMessage(content="", tool_calls=[{"name": "x", "args": {}, "id": "1"}]),
|
||||
AIMessage(content="第一轮回答"),
|
||||
HumanMessage(content="第二轮"),
|
||||
AIMessage(content="第二轮回答"),
|
||||
]
|
||||
filtered = filter_messages_for_memory(raw)
|
||||
assert count_user_turns(filtered) == 2
|
||||
Reference in new issue
Block a user