import pytest import pytest_asyncio from sqlalchemy import select from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.gateway.llmwiki_rag import _source_from_assistant_result, _source_from_result from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository from deerflow.persistence.assistant_knowledge.model import ( AssistantKnowledgeEntityContributionRow, AssistantKnowledgeEntityRow, ) from deerflow.persistence.base import Base class FakeEmbeddingClient: dimensions = 3 fingerprint = "test-bge-m3" async def embed_query(self, _query: str) -> list[float]: return [1.0, 0.0, 0.0] @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" @pytest.mark.asyncio async def test_custom_base_import_is_also_mirrored_to_global_base( repository: AssistantKnowledgeRepository, ) -> None: global_base = await repository.initialize(created_by="admin") custom_base = await repository.create_base( created_by="admin", name="专题助手库", description="只承接某个专题的 WeKnora 导出数据", ) package = { "source": {"type": "weknora"}, "wiki_pages": [{"slug": "wiki/topic", "title": "专题 Wiki", "content": "王五负责统筹。"}], "entities": [{"local_id": "e1", "type": "person", "name": "王五"}], "relations": [], } job = await repository.import_package( package=package, source_type="weknora", source_key="kb-topic", source_name="专题普通库", trigger="manual_from_weknora", created_by="admin", mapping_id="kb-topic", base_id=custom_base["id"], ) assert job["base_id"] == custom_base["id"] assert job["mirrored_global_base_id"] == global_base["id"] assert len(await repository.list_pages(base_id=custom_base["id"], query="王五")) == 2 assert len(await repository.list_pages(base_id=global_base["id"], query="王五")) == 2 bases = await repository.list_bases() assert [base["id"] for base in bases[:2]] == [global_base["id"], custom_base["id"]] @pytest.mark.asyncio async def test_imported_vectors_are_reused_for_assistant_search( repository: AssistantKnowledgeRepository, ) -> None: base = await repository.initialize(created_by="admin") job = await repository.import_package( package={ "source": {"type": "weknora"}, "wiki_pages": [{"slug": "wiki/vector", "title": "向量页", "content": "页面正文"}], "vectors": [ { "wiki_slug": "wiki/vector", "section_index": 0, "section_content": "张三负责向量化检索。", "embedding_model": "bge-embedding-m3", "embedding_fingerprint": "test-bge-m3", "embedding_dimensions": 3, "vector": [1.0, 0.0, 0.0], }, { "wiki_slug": "wiki/vector", "section_index": 1, "section_content": "无关片段。", "embedding_model": "bge-embedding-m3", "embedding_fingerprint": "test-bge-m3", "embedding_dimensions": 3, "vector": [0.0, 1.0, 0.0], }, ], }, source_type="weknora", source_key="kb-vector", source_name="向量普通库", trigger="manual_package_upload", created_by="admin", ) assert job["counts"]["vectors"] == 2 results = await repository.search_pages( base_ids=[base["id"]], query="这个查询不靠关键词", embedding_client=FakeEmbeddingClient(), ) assert results[0]["retrieval_mode"] == "vector" assert results[0]["matched_section"]["embedding_model"] == "bge-embedding-m3" assert results[0]["content_markdown"] == "张三负责向量化检索。" sources = await repository.list_sources() assert await repository.delete_source(sources[0]["id"]) assert await repository.search_pages( base_ids=[base["id"]], query="这个查询不靠关键词", embedding_client=FakeEmbeddingClient(), ) == [] @pytest.mark.asyncio async def test_reimport_retires_pages_removed_from_source_snapshot( repository: AssistantKnowledgeRepository, ) -> None: base = await repository.initialize(created_by="admin") common = { "source_type": "weknora", "source_key": "kb-refresh", "source_name": "自动同步普通库", "trigger": "vectorization_completed", "created_by": "admin", "mapping_id": "kb-refresh", } await repository.import_package( package={ "source": {"type": "weknora", "knowledge_base_id": "kb-refresh"}, "wiki_pages": [ {"slug": "wiki/kept", "title": "保留页", "content": "保留内容"}, {"slug": "wiki/removed", "title": "删除页", "content": "旧内容"}, ], }, **common, ) await repository.import_package( package={ "source": {"type": "weknora", "knowledge_base_id": "kb-refresh"}, "wiki_pages": [ {"slug": "wiki/kept", "title": "保留页", "content": "新内容"}, ], }, **common, ) pages = await repository.list_pages(base_id=base["id"], include_chunks=True) assert [page["slug"] for page in pages] == ["wiki/kept"] assert pages[0]["content_markdown"] == "新内容" @pytest.mark.asyncio async def test_custom_base_can_be_renamed_and_deleted_without_affecting_global_base( repository: AssistantKnowledgeRepository, ) -> None: global_base = await repository.initialize(created_by="admin") custom_base = await repository.create_base( created_by="admin", name="待编辑助手库", description="旧描述", ) await repository.import_package( package={ "source": {"type": "weknora"}, "wiki_pages": [ {"slug": "wiki/delete-test", "title": "删除测试", "content": "保留在总库中的内容"} ], "entities": [{"local_id": "e1", "type": "concept", "name": "删除测试实体"}], "relations": [], }, source_type="weknora", source_key="kb-delete-test", source_name="删除测试普通库", trigger="manual_from_weknora", created_by="admin", mapping_id="kb-delete-test", base_id=custom_base["id"], ) updated = await repository.update_base( custom_base["id"], name="已编辑助手库", description="新描述", ) assert updated is not None assert updated["name"] == "已编辑助手库" assert updated["description"] == "新描述" assert await repository.delete_base(custom_base["id"]) assert await repository.get_base(custom_base["id"]) is None assert all(base["id"] != custom_base["id"] for base in await repository.list_bases()) assert await repository.list_pages(base_id=global_base["id"], query="保留在总库") with pytest.raises(ValueError, match="不能删除"): await repository.delete_base(global_base["id"]) @pytest.mark.asyncio async def test_global_base_aligns_same_entity_title_and_searches_knowledge_body( repository: AssistantKnowledgeRepository, ) -> None: base = await repository.initialize(created_by="admin") first_package = { "source": {"provider": "weknora", "knowledge_base_id": "kb-one", "knowledge_base_name": "人员库"}, "wiki_pages": [ { "slug": "people/zhang-san", "title": "张 三", "page_type": "article", "content": "张三负责设备巡检。", "summary": "摘要专用词：月度综述。", } ], "entities": [ { "local_id": "one", "type": "人物", "name": "张 三", "aliases": ["Zhang San"], "source_wiki_slug": "people/zhang-san", } ], } second_package = { "source": {"provider": "weknora", "knowledge_base_id": "kb-two", "knowledge_base_name": "职责库"}, "wiki_pages": [ { "slug": "entity/person-001", "title": "Zhang San", "page_type": "concept", "content": "张三负责系统故障处置。", "summary": "另一条摘要。", } ], "entities": [ { "local_id": "two", "type": "person", "name": "张三", "aliases": ["Zhang San"], "source_wiki_slug": "entity/person-001", } ], } for package, source_key in ((first_package, "kb-one"), (second_package, "kb-two")): await repository.import_package( package=package, source_type="weknora", source_key=source_key, source_name=source_key, trigger="vectorization_completed", created_by="admin", ) pages = await repository.list_pages(base_id=base["id"]) assert len(pages) == 1 assert "张三负责设备巡检。" in pages[0]["content_markdown"] assert "张三负责系统故障处置。" in pages[0]["content_markdown"] assert len(pages[0]["origins"]) == 2 assert await repository.search_pages( base_ids=[base["id"]], query="月度综述", embedding_client=None, ) == [] results = await repository.search_pages( base_ids=[base["id"]], query="系统故障处置", embedding_client=None, ) assert len(results) == 1 assert "系统故障处置" in results[0]["matched_section"]["content"] assert "月度综述" not in results[0]["matched_section"]["content"] async with repository._sf() as session: entities = ( await session.execute( select(AssistantKnowledgeEntityRow).where( AssistantKnowledgeEntityRow.base_id == base["id"], AssistantKnowledgeEntityRow.status == "active", ) ) ).scalars().all() assert len(entities) == 1 contributions = ( await session.execute( select(AssistantKnowledgeEntityContributionRow).where( AssistantKnowledgeEntityContributionRow.entity_id == entities[0].id ) ) ).scalars().all() assert len(contributions) == 2 @pytest.mark.asyncio async def test_uploaded_package_reuses_embedded_weknora_source_identity( repository: AssistantKnowledgeRepository, ) -> None: base = await repository.initialize(created_by="admin") for upload_no, content in ((1, "旧知识"), (2, "新知识")): job = await repository.create_queued_import_job( source_type="weknora", source_key=f"upload:{upload_no}", source_name=f"upload-{upload_no}.jsonl", trigger="manual_package_upload", created_by="admin", base_id=base["id"], ) await repository.import_package_for_job( job["id"], package={ "source": { "provider": "weknora", "knowledge_base_id": "stable-kb-id", "knowledge_base_name": "同一个普通知识库", }, "wiki_pages": [ { "slug": "wiki/stable", "title": "稳定实体", "content": content, } ], }, created_by="admin", ) active_sources = await repository.list_sources(base_id=base["id"]) assert len(active_sources) == 1 assert active_sources[0]["source_key"] == "stable-kb-id" pages = await repository.list_pages(base_id=base["id"]) assert len(pages) == 1 assert pages[0]["content_markdown"] == "新知识" def test_assistant_citation_uses_matched_knowledge_instead_of_summary() -> None: source = _source_from_assistant_result( { "id": "page-1", "title": "设备知识", "summary": "只是一段摘要", "content_markdown": "完整页面正文", "matched_section": { "content": "实际命中的知识片段：设备过热时先断电。", }, "score": 0.91, }, 1, ) assert source["content"] == "实际命中的知识片段：设备过热时先断电。" assert source["snippet"] == "实际命中的知识片段：设备过热时先断电。" assert "只是一段摘要" not in source["snippet"] def test_regular_wiki_citation_does_not_fall_back_to_summary() -> None: source = _source_from_result( { "kind": "wiki_page", "wiki_page_id": "wiki-1", "title": "设备知识", "summary": "只是一段摘要", "content": "实际知识正文", "matched_sections": [], }, {"id": "kb-1", "name": "普通知识库"}, 1, ) assert source["content"] == "实际知识正文" assert source["snippet"] == "实际知识正文" assert "只是一段摘要" not in source["snippet"]