372 lines
12 KiB
Python
372 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import deerflow.agents.middlewares.title_middleware as title_middleware_module
|
|
from deerflow.agents.middlewares.title_middleware import (
|
|
TitleMiddleware,
|
|
finalize_thread_title,
|
|
)
|
|
from deerflow.config.title_config import TitleConfig
|
|
|
|
|
|
class _FakeCheckpointer:
|
|
def __init__(self, checkpoint: dict, metadata: dict | None = None):
|
|
self._tuple = SimpleNamespace(
|
|
checkpoint=checkpoint,
|
|
metadata=metadata or {},
|
|
config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
|
)
|
|
self.aput_calls: list[tuple[dict, dict, dict, dict]] = []
|
|
self._version = 0
|
|
|
|
async def aget_tuple(self, _config):
|
|
return self._tuple
|
|
|
|
def get_next_version(self, _current, _channel):
|
|
self._version += 1
|
|
return self._version
|
|
|
|
async def aput(self, config, checkpoint, metadata, versions):
|
|
self.aput_calls.append((config, checkpoint, metadata, versions))
|
|
self._tuple = SimpleNamespace(
|
|
checkpoint=checkpoint,
|
|
metadata=metadata,
|
|
config={
|
|
"configurable": {
|
|
"thread_id": config["configurable"]["thread_id"],
|
|
"checkpoint_ns": config["configurable"]["checkpoint_ns"],
|
|
"checkpoint_id": "new-ckpt",
|
|
}
|
|
},
|
|
)
|
|
return self._tuple.config
|
|
|
|
|
|
class _FakeThreadStore:
|
|
def __init__(self):
|
|
self.updated: list[tuple[str, str, str | None]] = []
|
|
|
|
async def update_display_name(self, thread_id: str, title: str, *, user_id=None):
|
|
self.updated.append((thread_id, title, user_id))
|
|
|
|
|
|
def _source_state() -> dict:
|
|
return {
|
|
"messages": [
|
|
SimpleNamespace(type="human", content="please analyze the first topic"),
|
|
SimpleNamespace(type="ai", content="this is the first answer"),
|
|
],
|
|
"title": "please analyze the first topic",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
|
|
def _dict_source_state() -> dict:
|
|
return {
|
|
"messages": [
|
|
{"type": "human", "content": "please analyze the first topic"},
|
|
{"type": "ai", "content": "this is the first answer"},
|
|
],
|
|
"title": "please analyze the first topic",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
|
|
def _knowledge_base_source_state(question: str = "这个库讲了什么") -> dict:
|
|
return {
|
|
"messages": [
|
|
SimpleNamespace(
|
|
type="human",
|
|
content=question,
|
|
additional_kwargs={
|
|
"llmwiki_knowledge_base_context": [
|
|
{"id": "kb-1", "name": "对话沉淀"}
|
|
]
|
|
},
|
|
),
|
|
SimpleNamespace(type="ai", content="这个知识库汇总了对话资料。"),
|
|
],
|
|
"title": question,
|
|
"title_provisional": True,
|
|
}
|
|
|
|
|
|
class _FakeTitleModel:
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
self.extra_body: dict = {}
|
|
self.prompt = ""
|
|
|
|
async def ainvoke(self, prompt, config):
|
|
self.prompt = prompt
|
|
return SimpleNamespace(content=self.content)
|
|
|
|
|
|
def _title_config() -> TitleConfig:
|
|
return TitleConfig(
|
|
enabled=True,
|
|
max_words=6,
|
|
max_chars=60,
|
|
prompt_template="User: {user_msg}\nAssistant: {assistant_msg}\nMax: {max_words}",
|
|
)
|
|
|
|
|
|
def test_admin_question_title_switch_finalizes_question_before_first_model_call(monkeypatch) -> None:
|
|
monkeypatch.setattr(
|
|
title_middleware_module,
|
|
"get_conversation_context_settings",
|
|
lambda: SimpleNamespace(use_question_as_title=True),
|
|
)
|
|
middleware = TitleMiddleware(
|
|
title_config=_title_config().model_copy(update={"enabled": False})
|
|
)
|
|
state = {
|
|
"messages": [SimpleNamespace(type="human", content="请帮我分析这份合同")],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace())
|
|
|
|
assert result == {
|
|
"title": "请帮我分析这份合同",
|
|
"title_provisional": False,
|
|
}
|
|
|
|
|
|
def test_admin_question_title_switch_off_keeps_async_title_generation(monkeypatch) -> None:
|
|
monkeypatch.setattr(
|
|
title_middleware_module,
|
|
"get_conversation_context_settings",
|
|
lambda: SimpleNamespace(use_question_as_title=False),
|
|
)
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
state = {
|
|
"messages": [SimpleNamespace(type="human", content="请帮我分析这份合同")],
|
|
}
|
|
|
|
result = middleware.before_model(state, runtime=SimpleNamespace())
|
|
|
|
assert result == {
|
|
"title": "请帮我分析这份合同",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_thread_title_updates_checkpoint_and_display_name(monkeypatch) -> None:
|
|
async def _fake_generate(self, state):
|
|
assert state["title_provisional"] is True
|
|
return {"title": "first topic analysis", "title_provisional": False}
|
|
|
|
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
|
|
|
|
checkpointer = _FakeCheckpointer({"channel_values": _source_state()})
|
|
thread_store = _FakeThreadStore()
|
|
|
|
updated = await finalize_thread_title(
|
|
checkpointer=checkpointer,
|
|
thread_store=thread_store,
|
|
thread_id="thread-1",
|
|
source_state=_source_state(),
|
|
)
|
|
|
|
assert updated is True
|
|
assert len(checkpointer.aput_calls) == 1
|
|
_, written_checkpoint, _, new_versions = checkpointer.aput_calls[0]
|
|
assert written_checkpoint["channel_values"]["title"] == "first topic analysis"
|
|
assert written_checkpoint["channel_values"]["title_provisional"] is False
|
|
assert set(new_versions) == {"title", "title_provisional"}
|
|
assert thread_store.updated == [("thread-1", "first topic analysis", None)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_thread_title_skips_if_current_thread_already_finalized(monkeypatch) -> None:
|
|
async def _fake_generate(self, state):
|
|
return {"title": "first topic analysis", "title_provisional": False}
|
|
|
|
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
|
|
|
|
checkpointer = _FakeCheckpointer(
|
|
{
|
|
"channel_values": {
|
|
**_source_state(),
|
|
"title": "manually renamed title",
|
|
"title_provisional": False,
|
|
}
|
|
}
|
|
)
|
|
thread_store = _FakeThreadStore()
|
|
|
|
updated = await finalize_thread_title(
|
|
checkpointer=checkpointer,
|
|
thread_store=thread_store,
|
|
thread_id="thread-1",
|
|
source_state=_source_state(),
|
|
)
|
|
|
|
assert updated is False
|
|
assert checkpointer.aput_calls == []
|
|
assert thread_store.updated == []
|
|
|
|
|
|
def test_sync_after_model_keeps_provisional_title_for_background_finalize() -> None:
|
|
middleware = TitleMiddleware()
|
|
|
|
result = middleware.after_model(_source_state(), runtime=SimpleNamespace())
|
|
|
|
assert result is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_after_model_keeps_provisional_title_for_background_finalize() -> None:
|
|
middleware = TitleMiddleware()
|
|
|
|
result = await middleware.aafter_model(_source_state(), runtime=SimpleNamespace())
|
|
|
|
assert result is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_thread_title_accepts_serialized_dict_messages(monkeypatch) -> None:
|
|
async def _fake_generate(self, state):
|
|
return {"title": "first topic analysis", "title_provisional": False}
|
|
|
|
monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)
|
|
|
|
checkpointer = _FakeCheckpointer({"channel_values": _dict_source_state()})
|
|
|
|
updated = await finalize_thread_title(
|
|
checkpointer=checkpointer,
|
|
thread_store=None,
|
|
thread_id="thread-1",
|
|
source_state=_dict_source_state(),
|
|
)
|
|
|
|
assert updated is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deictic_knowledge_base_question_rejects_invented_model_title(
|
|
monkeypatch,
|
|
) -> None:
|
|
model = _FakeTitleModel("Rust异步编程库简介")
|
|
monkeypatch.setattr(
|
|
title_middleware_module,
|
|
"create_chat_model",
|
|
lambda **_kwargs: model,
|
|
)
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
|
|
result = await middleware._agenerate_title_result(
|
|
_knowledge_base_source_state()
|
|
)
|
|
|
|
assert result == {
|
|
"title": "对话沉淀：知识库概览",
|
|
"title_provisional": False,
|
|
}
|
|
assert "服务端已验证的当前知识库名称" in model.prompt
|
|
assert '"对话沉淀"' in model.prompt
|
|
assert "禁止猜测或改写为其他库名" in model.prompt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deictic_knowledge_base_question_keeps_title_with_selected_name(
|
|
monkeypatch,
|
|
) -> None:
|
|
model = _FakeTitleModel("对话沉淀内容概览")
|
|
monkeypatch.setattr(
|
|
title_middleware_module,
|
|
"create_chat_model",
|
|
lambda **_kwargs: model,
|
|
)
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
|
|
result = await middleware._agenerate_title_result(
|
|
_knowledge_base_source_state("请介绍本知识库的主要内容")
|
|
)
|
|
|
|
assert result == {
|
|
"title": "对话沉淀内容概览",
|
|
"title_provisional": False,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ordinary_question_title_generation_is_unchanged(monkeypatch) -> None:
|
|
model = _FakeTitleModel("量子计算发展概览")
|
|
monkeypatch.setattr(
|
|
title_middleware_module,
|
|
"create_chat_model",
|
|
lambda **_kwargs: model,
|
|
)
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
state = {
|
|
"messages": [
|
|
SimpleNamespace(type="human", content="请介绍量子计算的发展"),
|
|
SimpleNamespace(type="ai", content="量子计算经历了多个发展阶段。"),
|
|
],
|
|
"title": "请介绍量子计算的发展",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
result = await middleware._agenerate_title_result(state)
|
|
|
|
assert result == {
|
|
"title": "量子计算发展概览",
|
|
"title_provisional": False,
|
|
}
|
|
assert "服务端已验证的当前知识库名称" not in model.prompt
|
|
|
|
|
|
def test_serialized_user_message_uses_knowledge_base_aware_local_title() -> None:
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
state = {
|
|
"messages": [
|
|
{
|
|
"type": "human",
|
|
"content": "当前库主要说了什么？",
|
|
"additional_kwargs": {
|
|
"llmwiki_knowledge_base_context": [
|
|
{"id": "kb-1", "name": "政策资料"}
|
|
]
|
|
},
|
|
},
|
|
{"type": "ai", "content": "该库收录了政策资料。"},
|
|
],
|
|
"title": "当前库主要说了什么？",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
assert middleware._generate_title_result(state) == {
|
|
"title": "政策资料：知识库概览",
|
|
"title_provisional": False,
|
|
}
|
|
|
|
|
|
def test_title_prompt_uses_last_non_empty_assistant_body_without_tool_calls() -> None:
|
|
middleware = TitleMiddleware(title_config=_title_config())
|
|
state = {
|
|
"messages": [
|
|
SimpleNamespace(type="human", content="请分析材料"),
|
|
SimpleNamespace(
|
|
type="ai",
|
|
content="准备读取文件",
|
|
tool_calls=[{"name": "read_file", "args": {}}],
|
|
),
|
|
SimpleNamespace(type="tool", content="工具返回内容"),
|
|
SimpleNamespace(type="ai", content="这是基于材料形成的最终结论。"),
|
|
SimpleNamespace(type="ai", content="<think>仅内部思考</think>"),
|
|
],
|
|
"title": "请分析材料",
|
|
"title_provisional": True,
|
|
}
|
|
|
|
prompt, _ = middleware._build_title_prompt(state)
|
|
|
|
assert "这是基于材料形成的最终结论。" in prompt
|
|
assert "准备读取文件" not in prompt
|
|
assert "仅内部思考" not in prompt
|