459 lines
16 KiB
Python
459 lines
16 KiB
Python
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"]
|