144 lines
4.9 KiB
Python
144 lines
4.9 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
|
|
from deerflow.config.skill_evolution_config import SkillEvolutionConfig
|
|
|
|
|
|
def test_review_worker_config_defaults():
|
|
cfg = SkillEvolutionConfig()
|
|
assert cfg.review_worker.enabled is False
|
|
assert cfg.review_worker.min_tool_calls == 5
|
|
assert cfg.review_worker.timeout_seconds == 120
|
|
assert cfg.review_worker.model_name is None
|
|
assert cfg.review_worker.max_messages == 60
|
|
|
|
|
|
def test_review_worker_config_from_dict():
|
|
cfg = SkillEvolutionConfig.model_validate({
|
|
"enabled": True,
|
|
"review_worker": {"enabled": True, "min_tool_calls": 5, "model_name": "my-model"},
|
|
})
|
|
assert cfg.review_worker.enabled is True
|
|
assert cfg.review_worker.min_tool_calls == 5
|
|
assert cfg.review_worker.model_name == "my-model"
|
|
|
|
|
|
def _make_ai_message(tool_calls_data: list[dict]):
|
|
msg = MagicMock()
|
|
msg.__class__.__name__ = "AIMessage"
|
|
tcs = []
|
|
for tc in tool_calls_data:
|
|
m = MagicMock()
|
|
m.name = tc["name"]
|
|
m.args = tc.get("args", {})
|
|
tcs.append(m)
|
|
msg.tool_calls = tcs
|
|
return msg
|
|
|
|
|
|
def test_extract_tool_call_count_counts_all():
|
|
from deerflow.skills.review_worker import extract_tool_call_count
|
|
msgs = [
|
|
HumanMessage(content="do something"),
|
|
_make_ai_message([{"name": "bash"}, {"name": "read_file"}]),
|
|
_make_ai_message([{"name": "write_file"}]),
|
|
]
|
|
assert extract_tool_call_count(msgs) == 3
|
|
|
|
|
|
def test_extract_tool_call_count_empty():
|
|
from deerflow.skills.review_worker import extract_tool_call_count
|
|
assert extract_tool_call_count([]) == 0
|
|
|
|
|
|
def test_internal_write_handshake_is_excluded_from_review_count_and_prompt():
|
|
from deerflow.skills.review_worker import _format_messages_for_review, extract_tool_call_count
|
|
|
|
messages = [
|
|
HumanMessage(content="写报告"),
|
|
AIMessage(
|
|
content="好的，我来写。",
|
|
tool_calls=[
|
|
{
|
|
"id": "begin-1",
|
|
"name": "begin_write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
ToolMessage(
|
|
content='{"status":"ready"}',
|
|
name="begin_write_file",
|
|
tool_call_id="begin-1",
|
|
),
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "write-1",
|
|
"name": "write_file",
|
|
"args": {"path": "/mnt/user-data/outputs/report.md", "content": "# 报告"},
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
]
|
|
|
|
assert extract_tool_call_count(messages) == 1
|
|
review_text = _format_messages_for_review(messages, max_messages=20)
|
|
assert "begin_write_file" not in review_text
|
|
assert '"status":"ready"' not in review_text
|
|
assert "write_file()" in review_text
|
|
|
|
|
|
def test_maybe_start_review_disabled_returns_false():
|
|
from deerflow.skills.review_worker import maybe_start_review
|
|
cfg = MagicMock()
|
|
cfg.skill_evolution.enabled = False
|
|
cfg.skill_evolution.review_worker.enabled = True
|
|
cfg.skill_evolution.review_worker.min_tool_calls = 2
|
|
msgs = [_make_ai_message([{"name": "bash"}, {"name": "read_file"}, {"name": "write_file"}])]
|
|
assert maybe_start_review("thread-1", msgs, cfg) is False
|
|
|
|
|
|
def test_maybe_start_review_too_few_tool_calls_returns_false():
|
|
from deerflow.skills.review_worker import maybe_start_review
|
|
cfg = MagicMock()
|
|
cfg.skill_evolution.enabled = True
|
|
cfg.skill_evolution.review_worker.enabled = True
|
|
cfg.skill_evolution.review_worker.min_tool_calls = 5
|
|
msgs = [_make_ai_message([{"name": "bash"}])]
|
|
assert maybe_start_review("thread-4", msgs, cfg) is False
|
|
|
|
|
|
def test_maybe_start_review_starts_task():
|
|
from deerflow.skills import review_worker as rw
|
|
# Reset module state so thread-3 has no lock
|
|
rw._locks.pop("thread-3", None)
|
|
rw._pending.pop("thread-3", None)
|
|
|
|
cfg = MagicMock()
|
|
cfg.skill_evolution.enabled = True
|
|
cfg.skill_evolution.review_worker.enabled = True
|
|
cfg.skill_evolution.review_worker.min_tool_calls = 2
|
|
cfg.skill_evolution.review_worker.timeout_seconds = 30
|
|
cfg.skill_evolution.review_worker.max_messages = 20
|
|
cfg.skill_evolution.review_worker.model_name = None
|
|
|
|
msgs = [_make_ai_message([{"name": "bash"}, {"name": "read_file"}, {"name": "write_file"}])]
|
|
|
|
with patch("deerflow.skills.review_worker.asyncio") as mock_asyncio:
|
|
fake_lock = MagicMock()
|
|
fake_lock.locked.return_value = False
|
|
mock_asyncio.Lock.return_value = fake_lock
|
|
mock_asyncio.create_task = MagicMock()
|
|
|
|
result = rw.maybe_start_review("thread-3", msgs, cfg)
|
|
|
|
assert result is True
|
|
mock_asyncio.create_task.assert_called_once()
|
|
scheduled_coro = mock_asyncio.create_task.call_args.args[0]
|
|
scheduled_coro.close()
|