194 lines
7.1 KiB
Python
194 lines
7.1 KiB
Python
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(),
|
|
) == []
|