221 lines
8.3 KiB
Python
221 lines
8.3 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from app.gateway.services import (
|
|
apply_knowledge_space_strict_scope,
|
|
enforce_knowledge_retrieval_mode_run_config,
|
|
enforce_knowledge_space_strict_run_config,
|
|
knowledge_space_strict_excluded_tools,
|
|
)
|
|
|
|
|
|
def test_strict_scope_preserves_existing_and_blocks_configured_web_tools() -> None:
|
|
config = SimpleNamespace(
|
|
tools=[
|
|
SimpleNamespace(name="company_search", group="web"),
|
|
SimpleNamespace(name="calculator", group="utility"),
|
|
]
|
|
)
|
|
|
|
excluded = knowledge_space_strict_excluded_tools(["custom_disabled"], app_config=config)
|
|
|
|
assert "custom_disabled" in excluded
|
|
assert "company_search" in excluded
|
|
assert "web_search" in excluded
|
|
assert "browser_fetch_page" in excluded
|
|
assert "calculator" not in excluded
|
|
assert "llmwiki_search" not in excluded
|
|
|
|
|
|
def test_harness_strict_scope_keeps_skills_and_report_delivery_available() -> None:
|
|
from deerflow.agents.lead_agent.agent import _knowledge_space_strict_exclusions
|
|
|
|
excluded = set(_knowledge_space_strict_exclusions([]))
|
|
|
|
assert "web_search" in excluded
|
|
assert "task" in excluded
|
|
assert "skill_view" not in excluded
|
|
assert "skill_list" not in excluded
|
|
assert "search_skills" not in excluded
|
|
assert "read_file" not in excluded
|
|
assert "deep_research_report" not in excluded
|
|
|
|
|
|
def test_selected_knowledge_space_activates_fail_closed_external_retrieval_block() -> None:
|
|
context = {"excluded_tools": ["custom_disabled"]}
|
|
config = SimpleNamespace(
|
|
tools=[
|
|
SimpleNamespace(name="company_search", group="web"),
|
|
SimpleNamespace(name="calculator", group="utility"),
|
|
]
|
|
)
|
|
|
|
activated = apply_knowledge_space_strict_scope(
|
|
context,
|
|
selection_requested=True,
|
|
knowledge_base_ids=[],
|
|
assistant_knowledge_base_ids=[],
|
|
app_config=config,
|
|
)
|
|
|
|
assert activated is True
|
|
assert context["knowledge_space_strict"] is True
|
|
assert "web_search" in context["excluded_tools"]
|
|
assert "browser_fetch_page" in context["excluded_tools"]
|
|
assert "company_search" in context["excluded_tools"]
|
|
assert "llmwiki_search" not in context["excluded_tools"]
|
|
assert "calculator" not in context["excluded_tools"]
|
|
|
|
|
|
def test_selected_scope_overrides_forged_client_runtime_config() -> None:
|
|
config = {
|
|
"configurable": {
|
|
"knowledge_space_strict": False,
|
|
"excluded_tools": [],
|
|
"llmwiki_knowledge_base_ids": ["forged"],
|
|
"assistant_knowledge_base_ids": ["forged-assistant"],
|
|
},
|
|
"context": {
|
|
"knowledge_space_strict": False,
|
|
"excluded_tools": ["client-disabled"],
|
|
"llmwiki_knowledge_base_ids": ["forged"],
|
|
"assistant_knowledge_base_ids": ["forged-assistant"],
|
|
},
|
|
}
|
|
|
|
enforce_knowledge_space_strict_run_config(
|
|
config,
|
|
knowledge_base_ids=["wiki-authorized"],
|
|
assistant_knowledge_base_ids=["assistant-authorized"],
|
|
excluded_tools=["server-disabled"],
|
|
app_config=SimpleNamespace(tools=[]),
|
|
)
|
|
|
|
for scope in ("configurable", "context"):
|
|
assert config[scope]["knowledge_space_strict"] is True
|
|
assert config[scope]["llmwiki_knowledge_base_ids"] == [
|
|
"wiki-authorized"
|
|
]
|
|
assert config[scope]["assistant_knowledge_base_ids"] == [
|
|
"assistant-authorized"
|
|
]
|
|
assert "web_search" in config[scope]["excluded_tools"]
|
|
assert "server-disabled" in config[scope]["excluded_tools"]
|
|
|
|
|
|
def test_retrieval_modes_mount_mutually_exclusive_knowledge_tool_surfaces() -> None:
|
|
classic = {
|
|
"configurable": {
|
|
"llmwiki_knowledge_base_ids": ["kb-1"],
|
|
"assistant_knowledge_base_ids": [],
|
|
"excluded_tools": [
|
|
"llmwiki_search",
|
|
"llmwiki_read_wiki_page",
|
|
"ordinary_knowledge_search",
|
|
],
|
|
},
|
|
"context": {
|
|
"llmwiki_knowledge_base_ids": ["kb-1"],
|
|
"assistant_knowledge_base_ids": [],
|
|
"excluded_tools": [],
|
|
},
|
|
}
|
|
enforce_knowledge_retrieval_mode_run_config(classic, "classic")
|
|
for scope in ("configurable", "context"):
|
|
excluded = set(classic[scope]["excluded_tools"])
|
|
assert "llmwiki_search" not in excluded
|
|
assert "llmwiki_read_wiki_page" not in excluded
|
|
assert "llmwiki_browse_wiki_index" not in excluded
|
|
assert "ordinary_knowledge_search" in excluded
|
|
assert "ordinary_wiki_search" in excluded
|
|
|
|
|
|
@pytest.mark.parametrize("delivery", [False, True])
|
|
def test_classic_strict_knowledge_scope_builds_lead_agent(monkeypatch, delivery) -> None:
|
|
from deerflow.agents.deep_research import retrieval
|
|
from deerflow.agents.lead_agent import agent as lead_agent
|
|
|
|
captured_tools: dict = {}
|
|
app_config = SimpleNamespace(
|
|
get_model_config=lambda _name: SimpleNamespace(supports_thinking=True),
|
|
)
|
|
monkeypatch.setattr(lead_agent, "_resolve_model_name", lambda *_args, **_kwargs: "test-model")
|
|
monkeypatch.setattr(lead_agent, "validate_agent_id", lambda _value: "roundtable-coordinator")
|
|
monkeypatch.setattr(lead_agent, "create_chat_model", lambda **_kwargs: object())
|
|
monkeypatch.setattr(lead_agent, "create_agent", lambda **kwargs: kwargs)
|
|
monkeypatch.setattr(lead_agent, "load_agent_soul", lambda _agent: "")
|
|
monkeypatch.setattr(lead_agent, "_log_es_query_routing_registration", lambda **_kwargs: None)
|
|
monkeypatch.setattr(lead_agent, "_build_middlewares", lambda *_args, **_kwargs: [])
|
|
monkeypatch.setattr(lead_agent, "apply_prompt_template", lambda **_kwargs: "test prompt")
|
|
monkeypatch.setattr(lead_agent, "_filter_available_skills_for_excluded_tools", lambda skills, _excluded: skills)
|
|
monkeypatch.setattr(retrieval, "resolve_collector_skill_visibility", lambda _agent, skills, _extra: skills)
|
|
monkeypatch.setattr(retrieval, "runtime_extra_skills", lambda _config: [])
|
|
|
|
import deerflow.tools
|
|
|
|
def fake_get_available_tools(**kwargs):
|
|
captured_tools.update(kwargs)
|
|
return []
|
|
|
|
monkeypatch.setattr(deerflow.tools, "get_available_tools", fake_get_available_tools)
|
|
|
|
result = lead_agent._make_lead_agent(
|
|
{
|
|
"configurable": {
|
|
"is_bootstrap": False,
|
|
"writing_mode": True,
|
|
"workbench_write_now": delivery,
|
|
"workbench_session_id": "test-session" if delivery else None,
|
|
"knowledge_space_strict": True,
|
|
"knowledge_retrieval_mode": "classic",
|
|
}
|
|
},
|
|
app_config=app_config,
|
|
)
|
|
|
|
assert result["system_prompt"].startswith("test prompt")
|
|
assert captured_tools["include_mcp"] is False
|
|
assert "web_search" in captured_tools["excluded_tools"]
|
|
assert ("deep_research_report" in {tool.name for tool in result["tools"]}) is delivery
|
|
|
|
smart = {
|
|
"configurable": {
|
|
"llmwiki_knowledge_base_ids": ["kb-1"],
|
|
"assistant_knowledge_base_ids": [],
|
|
"excluded_tools": ["ordinary_knowledge_search"],
|
|
},
|
|
"context": {
|
|
"llmwiki_knowledge_base_ids": ["kb-1"],
|
|
"assistant_knowledge_base_ids": [],
|
|
"excluded_tools": [],
|
|
},
|
|
}
|
|
enforce_knowledge_retrieval_mode_run_config(smart, "smart_reasoning")
|
|
for scope in ("configurable", "context"):
|
|
excluded = set(smart[scope]["excluded_tools"])
|
|
assert "ordinary_knowledge_search" not in excluded
|
|
assert "ordinary_wiki_search" not in excluded
|
|
assert "llmwiki_search" in excluded
|
|
assert "llmwiki_read_wiki_page" in excluded
|
|
assert "llmwiki_browse_wiki_index" in excluded
|
|
|
|
|
|
def test_classic_without_an_ordinary_space_hides_page_and_index_tools() -> None:
|
|
config = {
|
|
"configurable": {
|
|
"llmwiki_knowledge_base_ids": [],
|
|
"assistant_knowledge_base_ids": ["assistant-1"],
|
|
},
|
|
"context": {
|
|
"llmwiki_knowledge_base_ids": [],
|
|
"assistant_knowledge_base_ids": ["assistant-1"],
|
|
},
|
|
}
|
|
enforce_knowledge_retrieval_mode_run_config(config, "classic")
|
|
for scope in ("configurable", "context"):
|
|
excluded = set(config[scope]["excluded_tools"])
|
|
assert "llmwiki_search" not in excluded
|
|
assert "llmwiki_read_wiki_page" in excluded
|
|
assert "llmwiki_browse_wiki_index" in excluded
|