167 lines
5.3 KiB
Python
167 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from deerflow.agents.middlewares.title_middleware import (
|
|
TitleMiddleware,
|
|
finalize_thread_title,
|
|
)
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
|
|
@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
|