feat: add skill and assistant knowledge bases

2026-09-02 00:49:13 +08:00
1 parent dd0ee282df
commit fbd8dec48f
@@ -0,0 +1,96 @@
+import pytest
+import pytest_asyncio
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository
+from deerflow.persistence.assistant_knowledge.model import (
+    AssistantKnowledgeEntityContributionRow,
+    AssistantKnowledgeEntityRow,
+)
+from deerflow.persistence.base import Base
+
+
+@pytest_asyncio.fixture
+async def repository():
+    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+    async with engine.begin() as connection:
+        await connection.run_sync(Base.metadata.create_all)
+    yield AssistantKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False))
+    await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_single_base_entity_alignment_search_and_revision(repository: AssistantKnowledgeRepository) -> None:
+    first = await repository.initialize(created_by="admin")
+    second = await repository.initialize(created_by="someone-else")
+    assert first["id"] == second["id"]
+
+    package = {
+        "source": {"type": "skill"},
+        "wiki_pages": [{"slug": "skills/demo/overview", "title": "示例知识", "content": "张三负责测试。"}],
+        "entities": [{"local_id": "e1", "type": "person", "name": "张三", "aliases": ["Zhang San"]}],
+        "relations": [],
+        "documents": [],
+        "chunks": [],
+        "graph": {"nodes": [], "edges": []},
+    }
+    await repository.import_package(
+        package=package,
+        source_type="skill",
+        source_key="demo",
+        source_name="示例技能",
+        trigger="skill_direct",
+        created_by="admin",
+        skill_name="demo",
+    )
+    await repository.import_package(
+        package=package,
+        source_type="weknora",
+        source_key="kb-1",
+        source_name="普通库",
+        trigger="manual_from_weknora",
+        created_by="admin",
+    )
+
+    results = await repository.list_pages(query="张三")
+    alias_results = await repository.list_pages(query="Zhang San")
+    revisions = await repository.list_revisions("skills/demo/overview")
+    assert {row["page_type"] for row in results} == {"source", "entity"}
+    assert [row["page_type"] for row in alias_results] == ["entity"]
+    assert len(revisions) == 2
+
+
+@pytest.mark.asyncio
+async def test_delete_source_keeps_entity_until_last_contribution_is_removed(
+    repository: AssistantKnowledgeRepository,
+) -> None:
+    package = {
+        "source": {},
+        "wiki_pages": [],
+        "entities": [{"local_id": "e1", "type": "concept", "name": "共享实体"}],
+        "relations": [],
+    }
+    for source_key in ("one", "two"):
+        await repository.import_package(
+            package=package,
+            source_type="skill",
+            source_key=source_key,
+            source_name=source_key,
+            trigger="skill_direct",
+            created_by="admin",
+            skill_name=source_key,
+        )
+
+    sources = await repository.list_sources()
+    await repository.delete_source(next(row["id"] for row in sources if row["source_key"] == "one"))
+    async with repository._sf() as session:
+        entity = (await session.execute(select(AssistantKnowledgeEntityRow))).scalar_one()
+        contributions = (await session.execute(select(AssistantKnowledgeEntityContributionRow))).scalars().all()
+        assert entity.status == "active"
+        assert len(contributions) == 1
+
+    await repository.delete_source(next(row["id"] for row in sources if row["source_key"] == "two"))
+    async with repository._sf() as session:
+        entity = (await session.execute(select(AssistantKnowledgeEntityRow))).scalar_one()
+        assert entity.status == "deleted"
@@ -0,0 +1,62 @@
+import pytest
+
+from deerflow.skill_knowledge.graph_writer import (
+    EDGE_QUERY,
+    NODE_QUERY,
+    execute_plan,
+    parameterized_statements,
+    plan_from_ir,
+)
+
+
+class RecordingTransaction:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, dict]] = []
+
+    async def run(self, query: str, parameters: dict) -> None:
+        self.calls.append((query, parameters))
+
+
+def _plan():
+    return plan_from_ir(
+        {
+            "entities": [
+                {"local_id": "a", "type": "person", "name": "张三", "confidence": 1},
+                {"local_id": "b", "type": "system", "name": "知识库", "confidence": 1},
+            ],
+            "relations": [
+                {
+                    "source_local_id": "a",
+                    "target_local_id": "b",
+                    "predicate": "uses",
+                    "confidence": 0.9,
+                    "evidence": [{"file": "SKILL.md", "line": 2}],
+                }
+            ],
+        },
+        target_type="weknora",
+        target_id="kb-1",
+        binding_id="binding-1",
+        snapshot_id="snapshot-1",
+        job_id="job-1",
+    )
+
+
+def test_direct_graph_plan_uses_only_fixed_parameterized_templates() -> None:
+    statements = parameterized_statements(_plan())
+    assert len(statements) == 3
+    assert {row["statement"] for row in statements} == {NODE_QUERY, EDGE_QUERY}
+    assert all("parameters" in row for row in statements)
+    assert all("张三" not in row["statement"] for row in statements)
+
+
+@pytest.mark.asyncio
+async def test_direct_graph_writer_obeys_explicit_enable_switch() -> None:
+    transaction = RecordingTransaction()
+    disabled = await execute_plan(transaction, _plan(), enabled=False)
+    assert disabled == {"nodes": 0, "edges": 0}
+    assert transaction.calls == []
+
+    enabled = await execute_plan(transaction, _plan(), enabled=True)
+    assert enabled == {"nodes": 2, "edges": 1}
+    assert len(transaction.calls) == 3
@@ -0,0 +1,95 @@
+import pytest
+import pytest_asyncio
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+from deerflow.persistence.base import Base
+from deerflow.persistence.skill_knowledge import SkillKnowledgeRepository
+
+
+@pytest_asyncio.fixture
+async def repository():
+    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+    async with engine.begin() as connection:
+        await connection.run_sync(Base.metadata.create_all)
+    yield SkillKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False))
+    await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_matrix_job_creates_one_item_and_binding_per_pair(repository: SkillKnowledgeRepository) -> None:
+    job = await repository.create_job(
+        skill_names=["a", "b", "c"],
+        targets=[
+            {"target_type": "weknora", "target_mode": "wiki", "target_id": "kb-1", "target_name": "一"},
+            {"target_type": "assistant", "target_mode": "assistant", "target_id": "base", "target_name": "知识梳理"},
+        ],
+        created_by="admin",
+    )
+    detail = await repository.get_job(job["id"])
+    bindings = await repository.list_bindings()
+
+    assert detail is not None
+    assert detail["total_count"] == 6
+    assert len(detail["items"]) == 6
+    assert len(bindings) == 6
+
+
+@pytest.mark.asyncio
+async def test_detach_delete_only_removes_binding_contributions(repository: SkillKnowledgeRepository) -> None:
+    job = await repository.create_job(
+        skill_names=["a", "b"],
+        targets=[{"target_type": "assistant", "target_mode": "assistant", "target_id": "base", "target_name": "知识梳理"}],
+        created_by="admin",
+        review_mode="off",
+    )
+    detail = await repository.get_job(job["id"])
+    assert detail is not None
+    binding = (await repository.list_bindings(skill_names=["a"]))[0]
+    detached = await repository.detach_binding(binding["id"], delete_contributions=True)
+    assert detached and detached["status"] == "detached_deleted"
+    assert len(await repository.list_bindings()) == 1
+
+
+@pytest.mark.asyncio
+async def test_item_claim_is_atomic_and_records_worker_lease(repository: SkillKnowledgeRepository) -> None:
+    job = await repository.create_job(
+        skill_names=["a"],
+        targets=[
+            {
+                "target_type": "assistant",
+                "target_mode": "assistant",
+                "target_id": "base",
+                "target_name": "知识梳理",
+            }
+        ],
+        created_by="admin",
+    )
+    detail = await repository.get_job(job["id"])
+    item_id = detail["items"][0]["id"]
+
+    first = await repository.claim_item(item_id, owner="worker-a")
+    second = await repository.claim_item(item_id, owner="worker-b")
+
+    assert first and first["lease_owner"] == "worker-a"
+    assert first["attempts"] == 1
+    assert second is None
+
+
+@pytest.mark.asyncio
+async def test_binding_allows_only_one_active_item(repository: SkillKnowledgeRepository) -> None:
+    target = {
+        "target_type": "assistant",
+        "target_mode": "assistant",
+        "target_id": "base",
+        "target_name": "知识梳理",
+    }
+    first = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
+    duplicate = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
+    assert first["total_count"] == 1
+    assert duplicate["total_count"] == 0
+    assert duplicate["status"] == "completed"
+
+    detail = await repository.get_job(first["id"])
+    await repository.update_item(detail["items"][0]["id"], status="failed")
+    retry = await repository.create_job(skill_names=["a"], targets=[target], created_by="admin")
+    assert retry["total_count"] == 1
@@ -0,0 +1,52 @@
+from pathlib import Path
+from zipfile import ZipFile
+
+from deerflow.skill_knowledge import build_skill_ir, scan_skill
+
+
+def test_scanner_excludes_code_blocks_source_and_sensitive_files(tmp_path: Path) -> None:
+    (tmp_path / "SKILL.md").write_text(
+        "# 示例\n本技能使用 客户画像数据。\n```python\nPASSWORD = 'must-not-enter-ir'\n```\n",
+        encoding="utf-8",
+    )
+    (tmp_path / "worker.py").write_text("# uses hidden system\nSECRET='never-read'", encoding="utf-8")
+    (tmp_path / ".env").write_text("TOKEN=super-secret-value", encoding="utf-8")
+
+    scan = scan_skill("example", tmp_path)
+    ir = build_skill_ir(scan)
+
+    assert scan.skipped_code_count == 1
+    assert scan.blocked_files == [{"path": ".env", "reason": "疑似密钥或敏感配置文件"}]
+    assert "must-not-enter-ir" not in str(ir)
+    assert "never-read" not in str(ir)
+    assert any(relation["predicate"] == "uses" for relation in ir["relations"])
+
+
+def test_code_only_change_does_not_change_digest(tmp_path: Path) -> None:
+    (tmp_path / "SKILL.md").write_text("# Stable knowledge", encoding="utf-8")
+    code = tmp_path / "worker.py"
+    code.write_text("print(1)", encoding="utf-8")
+    before = scan_skill("example", tmp_path).source_digest
+    code.write_text("print(2)", encoding="utf-8")
+    after = scan_skill("example", tmp_path).source_digest
+    assert before == after
+
+
+def test_zip_is_inspected_and_only_safe_knowledge_members_are_read(tmp_path: Path) -> None:
+    archive = tmp_path / "materials.zip"
+    with ZipFile(archive, "w") as output:
+        output.writestr("guide.md", "# 操作手册\n仅提炼这段知识。")
+        output.writestr("worker.py", "SECRET = 'code-must-not-be-read'")
+        output.writestr("../escape.md", "path traversal")
+
+    scan = scan_skill("example", tmp_path)
+
+    assert [file.path for file in scan.files] == ["materials.zip!/guide.md"]
+    assert scan.skipped_code_count == 1
+    assert scan.blocked_files == [
+        {
+            "path": "materials.zip!/../escape.md",
+            "reason": "压缩包包含路径穿越或符号链接",
+        }
+    ]
+    assert "code-must-not-be-read" not in str(scan.manifest)