mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 19:16:17 +00:00
b6503e9a35
* feat(knowledge): integrate RAGFlow retrieval and management * test(knowledge): cover merged listing tool * feat(knowledge): add per-message retrieval scope * chore(docs): remove unrelated document * docs(knowledge): add interaction screenshots * feat(knowledge): simplify scope selector trigger * docs(knowledge): refresh selector screenshot * feat(knowledge): defer standalone management * docs(knowledge): show chat-only scope UI * fix(knowledge): honor scope on clarification replies * fix(knowledge): harden scoped replay validation * docs(knowledge): clarify replay scope precedence * fix(knowledge): keep provider settings on tools * fix(config): preserve tools-only knowledge settings * fix(knowledge): submit custom assistant identity * refactor(knowledge): trim PR scope changes * fix(knowledge): sanitize document scope display * feat(knowledge): enable scope selection in main chat * fix(knowledge): emphasize active scope icon without button frame * fix(knowledge): close context scrubbing and refresh e2e checks * fix(knowledge): preserve idempotent canonical retries * fix(knowledge): accept promptless conversation runs * style(knowledge): format backend regression tests * chore(knowledge): trim PR scope and fix frontend format * fix(knowledge): remove shared-scope notice * fix(knowledge): remove scope persistence notice * docs(knowledge): include main chat in catalog scope * fix(knowledge): preserve scope recovery and upgrades * fix(config): preserve LightRAG knowledge upgrades --------- Co-authored-by: foreleven <for-eleven@hotmail.com>
425 lines
16 KiB
Python
425 lines
16 KiB
Python
import asyncio
|
|
from enum import Enum
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
|
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
|
from deerflow.subagents import batch_service as service_module
|
|
from deerflow.subagents.batch_runtime import BatchSubmitRequest
|
|
from deerflow.subagents.batch_service import SubagentBatchService
|
|
from deerflow.subagents.capacity import SubagentExecutionCapacity
|
|
|
|
|
|
class FakeStatus(Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
@property
|
|
def is_terminal(self) -> bool:
|
|
return self in {FakeStatus.COMPLETED, FakeStatus.FAILED}
|
|
|
|
|
|
def _request(**overrides) -> BatchSubmitRequest:
|
|
values = {
|
|
"user_id": "user-1",
|
|
"thread_id": "thread-1",
|
|
"run_id": "run-1",
|
|
"tool_call_id": "call-1",
|
|
"submission_key": "run-1:call-1",
|
|
"title": "Records",
|
|
"subagent_type": "general-purpose",
|
|
"items": [{"key": "record-1", "prompt": "Process record 1"}],
|
|
"max_live_items": None,
|
|
"max_running_items": None,
|
|
"execution_spec": {
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "General purpose",
|
|
"system_prompt": "Work carefully.",
|
|
},
|
|
"parent_model": "model-a",
|
|
"knowledge_scope": {
|
|
"version": 1,
|
|
"mode": "selected",
|
|
"dataset_ids": ["dataset-1"],
|
|
},
|
|
},
|
|
}
|
|
values.update(overrides)
|
|
return BatchSubmitRequest(**values)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_keeps_batch_running_limit_separate_from_one_process_capacity() -> None:
|
|
repository = SimpleNamespace(create_batch=AsyncMock(return_value={"id": "batch-1"}))
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(max_running_items_per_batch=32),
|
|
runtime_config=SubagentRuntimeConfig(max_running=3),
|
|
)
|
|
|
|
result = await service.submit(_request(max_live_items=20, max_running_items=10))
|
|
|
|
assert result == {"id": "batch-1"}
|
|
assert repository.create_batch.await_args.kwargs["max_running_items"] == 10
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_item_marks_real_running_then_persists_terminal_result(monkeypatch) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.RUNNING,
|
|
result=None,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
)
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.marked_running = False
|
|
self.finalized = None
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": _request().execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def mark_item_running(self, *_args, **_kwargs):
|
|
self.marked_running = True
|
|
result.status = FakeStatus.COMPLETED
|
|
result.result = "done"
|
|
return True
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def finalize_item(self, *_args, **kwargs):
|
|
self.finalized = kwargs
|
|
return True
|
|
|
|
execution_capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1))
|
|
executor_kwargs = {}
|
|
|
|
class Executor:
|
|
def __init__(self, **kwargs) -> None:
|
|
executor_kwargs.update(kwargs)
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", lambda _execution_id: result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
execution_capacity=execution_capacity,
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.gather(*list(service._executions.values()))
|
|
|
|
assert repository.marked_running is True
|
|
assert repository.finalized is not None
|
|
assert repository.finalized["succeeded"] is True
|
|
assert repository.finalized["result"] == "done"
|
|
assert executor_kwargs["execution_capacity"] is execution_capacity
|
|
assert executor_kwargs["knowledge_scope"] == {
|
|
"version": 1,
|
|
"mode": "selected",
|
|
"dataset_ids": ["dataset-1"],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(monkeypatch) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.PENDING,
|
|
result=None,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
)
|
|
reads = 0
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.finalized = None
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": _request().execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def mark_item_running(self, *_args, **_kwargs):
|
|
raise AssertionError("a task that completes between polls need not expose running")
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
# Exactly one pre-launch revalidation is expected; the poll loop
|
|
# must not renew for a task that completes between polls.
|
|
self.renews = getattr(self, "renews", 0) + 1
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def finalize_item(self, *_args, **kwargs):
|
|
self.finalized = kwargs
|
|
return True
|
|
|
|
class Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
def read_result(_execution_id):
|
|
nonlocal reads
|
|
reads += 1
|
|
if reads > 1:
|
|
result.status = FakeStatus.COMPLETED
|
|
result.result = "fast result"
|
|
return result
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", read_result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(poll_interval_seconds=0.1, lease_seconds=120),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*list(service._executions.values())),
|
|
timeout=1,
|
|
)
|
|
|
|
assert repository.finalized is not None
|
|
assert repository.finalized["result"] == "fast result"
|
|
assert repository.renews == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_admission_failure_requeues_instead_of_finalizing(monkeypatch) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.FAILED,
|
|
result=None,
|
|
error="Process-wide subagent capacity is full",
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
admission_failure=True,
|
|
)
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.requeued = None
|
|
self.finalized = False
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": _request().execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def requeue_item_after_admission_failure(self, item_id, **kwargs):
|
|
self.requeued = (item_id, kwargs)
|
|
return True
|
|
|
|
async def finalize_item(self, *_args, **_kwargs):
|
|
self.finalized = True
|
|
return True
|
|
|
|
class Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", lambda _execution_id: result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.gather(*list(service._executions.values()))
|
|
|
|
assert repository.requeued is not None
|
|
assert repository.requeued[0] == "item-1"
|
|
assert repository.finalized is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) -> None:
|
|
"""A batch cancelled while tool assembly is blocked must not launch.
|
|
|
|
Regression (review of 249dba82): ``_execute_item()`` called
|
|
``executor.execute_async()`` unconditionally after assembly, so
|
|
``cancel_batch()`` landing while assembly was blocked in the worker thread
|
|
terminalized the durable item but could not stop the not-yet-started
|
|
execution, and the orphaned launch still invoked the model.
|
|
"""
|
|
import threading
|
|
from datetime import UTC, datetime
|
|
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
|
from deerflow.persistence.subagent_batches import SubagentBatchRepository
|
|
|
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
|
try:
|
|
repository = SubagentBatchRepository(get_session_factory())
|
|
await repository.create_batch(
|
|
batch_id="batch-1",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
tool_call_id="call-1",
|
|
submission_key="run-1:call-1",
|
|
title="Cancelled during assembly",
|
|
subagent_type="general-purpose",
|
|
items=[{"key": "record-1", "prompt": "Process record 1"}],
|
|
max_live_items=2,
|
|
max_running_items=1,
|
|
max_attempts=2,
|
|
execution_spec={
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "test",
|
|
"system_prompt": "sys",
|
|
},
|
|
},
|
|
)
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(lease_seconds=60),
|
|
runtime_config=SubagentRuntimeConfig(max_running=3),
|
|
app_config=SimpleNamespace(),
|
|
)
|
|
claimed = await repository.claim_items(
|
|
now=datetime.now(UTC),
|
|
lease_owner=service._lease_owner,
|
|
lease_seconds=60,
|
|
limit=10,
|
|
)
|
|
assert len(claimed) == 1
|
|
item = claimed[0]
|
|
assert item["batch"]["id"] == "batch-1"
|
|
|
|
assembly_started = threading.Event()
|
|
assembly_release = threading.Event()
|
|
|
|
def _blocking_assembly(**_kwargs):
|
|
# Real blocking wait in the assembly-pool worker: parks assembly
|
|
# until the test has cancelled the batch.
|
|
assembly_started.set()
|
|
assembly_release.wait(timeout=15)
|
|
return []
|
|
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", _blocking_assembly)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_a, **_k: "test-model")
|
|
monkeypatch.setattr(service_module, "request_cancel_background_task", lambda _execution_id: None)
|
|
|
|
launched: list[str] = []
|
|
|
|
class _Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
launched.append(task_id or "generated")
|
|
return task_id or "generated"
|
|
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", _Executor)
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"get_background_task_result",
|
|
lambda _execution_id: SimpleNamespace(
|
|
status=FakeStatus.COMPLETED,
|
|
result="done",
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=[],
|
|
),
|
|
)
|
|
|
|
item_task = asyncio.create_task(service._execute_item(item))
|
|
assert await asyncio.to_thread(assembly_started.wait, 15)
|
|
|
|
cancelled = await service.cancel_batch(batch_id="batch-1", user_id="user-1")
|
|
assert cancelled is not None
|
|
|
|
assembly_release.set()
|
|
await asyncio.wait_for(item_task, timeout=10)
|
|
|
|
assert launched == [], "cancelled work must not launch after assembly"
|
|
batch = await repository.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["counts"]["cancelled"] == 1
|
|
finally:
|
|
await close_engine()
|