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 class FakeEmbeddingClient: dimensions = 3 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_dimensions": 3, "vector": [1.0, 0.0, 0.0], }, { "wiki_slug": "wiki/vector", "section_index": 1, "section_content": "无关片段。", "embedding_model": "bge-embedding-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(), ) == []