Files
deerflow-code/offline-backend-20260512/backend/tests/test_assistant_knowledge_repository.py

1313 lines
46 KiB
Python

import asyncio
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock
import pytest
import pytest_asyncio
from sqlalchemy import event, select
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from app.gateway.knowledge_transfer import sync_to_global
from app.gateway.llmwiki_rag import _source_from_assistant_result, _source_from_result
from app.gateway.routers import assistant_knowledge
from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository
from deerflow.persistence.assistant_knowledge import store as assistant_store_module
from deerflow.persistence.assistant_knowledge.model import (
AssistantKnowledgeEntityContributionRow,
AssistantKnowledgeEntityRow,
AssistantKnowledgeRelationRow,
AssistantKnowledgeWikiPageRow,
)
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]
class IncompatibleEmbeddingClient:
dimensions = 3
fingerprint = "another-model"
def __init__(self) -> None:
self.embed_query = AsyncMock(side_effect=AssertionError("incompatible vectors must be rejected before encoding"))
@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_import_retries_transient_database_lock_contention(
repository: AssistantKnowledgeRepository,
monkeypatch: pytest.MonkeyPatch,
) -> None:
attempts = 0
async def flaky_import(**_kwargs):
nonlocal attempts
attempts += 1
if attempts < 3:
raise OperationalError(
"UPDATE assistant_knowledge_base",
{},
Exception(1205, "Lock wait timeout exceeded; try restarting transaction"),
)
return {"id": "completed-job", "status": "completed"}
sleeps: list[float] = []
async def record_sleep(delay: float) -> None:
sleeps.append(delay)
monkeypatch.setattr(
repository,
"_import_package_into_base_locked",
flaky_import,
)
monkeypatch.setattr(
assistant_store_module,
"_IMPORT_CONTENTION_RETRY_DELAYS",
(0.01, 0.02, 0.03),
)
monkeypatch.setattr(asyncio, "sleep", record_sleep)
result = await repository._import_package_into_base(
base={"id": "global-base"},
source_key="mapping-1",
)
assert result == {"id": "completed-job", "status": "completed"}
assert attempts == 3
assert sleeps == [0.01, 0.02]
@pytest.mark.asyncio
async def test_import_does_not_retry_non_contention_failures(
repository: AssistantKnowledgeRepository,
monkeypatch: pytest.MonkeyPatch,
) -> None:
import_call = AsyncMock(side_effect=ValueError("invalid package"))
monkeypatch.setattr(
repository,
"_import_package_into_base_locked",
import_call,
)
with pytest.raises(ValueError, match="invalid package"):
await repository._import_package_into_base(
base={"id": "global-base"},
source_key="mapping-1",
)
import_call.assert_awaited_once()
class _KnowledgeTransferIndexStore:
async def get_index_revision(self, mapping_id: str) -> tuple[int, str, str]:
assert mapping_id == "mapping-1"
return 7, "embedding-fingerprint", "idle"
class _KnowledgeTransferAssistantStore:
def __init__(self) -> None:
self.jobs: dict[str, dict[str, Any]] = {}
self.completed_source: dict[str, Any] | None = None
self.created_jobs = 0
async def initialize(self, *, created_by: str | None) -> dict[str, Any]:
assert created_by is None
return {"id": "global-base"}
async def list_sources(self, *, base_id: str) -> list[dict[str, Any]]:
assert base_id == "global-base"
return [dict(self.completed_source)] if self.completed_source else []
async def get_job(self, job_id: str) -> dict[str, Any] | None:
return self.jobs.get(job_id)
async def create_queued_import_job(self, **kwargs: Any) -> dict[str, Any]:
assert kwargs["base_id"] == "global-base"
self.created_jobs += 1
job = {"id": f"job-{self.created_jobs}", "status": "queued"}
self.jobs[job["id"]] = job
return dict(job)
@pytest.mark.asyncio
async def test_concurrent_global_syncs_share_one_mapping_import(
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = _KnowledgeTransferAssistantStore()
started = asyncio.Event()
release = asyncio.Event()
async def export_then_import(*, job_id: str, **_kwargs: Any) -> None:
started.set()
await release.wait()
store.jobs[job_id] = {
"id": job_id,
"status": "completed",
"counts": {"wiki_pages": 1},
}
store.completed_source = {
"source_key": "mapping-1",
"metadata": {"index_revision": 7},
"last_import_job_id": job_id,
}
monkeypatch.setattr(
assistant_knowledge,
"_export_weknora_then_import_for_job",
export_then_import,
)
app = SimpleNamespace(
state=SimpleNamespace(
assistant_knowledge_store=store,
llmwiki_index_store=_KnowledgeTransferIndexStore(),
)
)
mapping = {
"id": "mapping-1",
"name": "测试知识库",
"weknora_id": "remote-1",
}
first_task = asyncio.create_task(sync_to_global(app, mapping))
await asyncio.wait_for(started.wait(), timeout=1)
second_task = asyncio.create_task(sync_to_global(app, mapping))
await asyncio.sleep(0)
assert store.created_jobs == 1
release.set()
first, second = await asyncio.gather(first_task, second_task)
assert first["status"] == "completed"
assert second["status"] == "unchanged"
assert first["job_id"] == second["job_id"]
assert store.created_jobs == 1
async def _import_relation_graph(
repository: AssistantKnowledgeRepository,
*,
base_id: str,
prefix: str,
) -> dict[str, dict]:
await repository.import_package(
package={
"source": {"provider": "weknora", "knowledge_base_id": f"kb-{prefix}"},
"wiki_pages": [
{
"slug": f"{prefix}/seed",
"title": f"{prefix} 起点",
"page_type": "article",
"content": f"{prefix} 起点正文",
},
{
"slug": f"{prefix}/target",
"title": f"{prefix} 目标",
"page_type": "article",
"content": f"{prefix} 目标正文",
},
],
"entities": [
{
"local_id": f"{prefix}-seed",
"type": "concept",
"name": f"{prefix} 起点",
"source_wiki_slug": f"{prefix}/seed",
},
{
"local_id": f"{prefix}-target",
"type": "concept",
"name": f"{prefix} 目标",
"source_wiki_slug": f"{prefix}/target",
},
],
"relations": [
{
"source_entity_id": f"{prefix}-seed",
"target_entity_id": f"{prefix}-target",
"predicate": "references",
"source_wiki_slug": f"{prefix}/seed",
}
],
},
source_type="weknora",
source_key=f"kb-{prefix}",
source_name=f"{prefix} 普通知识库",
trigger="manual_package_upload",
created_by="admin",
base_id=base_id,
mirror_to_global=False,
)
pages = await repository.list_pages(base_id=base_id)
return {str(page["slug"]): page for page in pages}
@pytest.mark.asyncio
async def test_graph_adjacency_accepts_page_id_or_slug_and_never_crosses_base(
repository: AssistantKnowledgeRepository,
) -> None:
first_base = await repository.create_base(created_by="admin", name="图谱一")
second_base = await repository.create_base(created_by="admin", name="图谱二")
first_pages = await _import_relation_graph(
repository,
base_id=first_base["id"],
prefix="first",
)
second_pages = await _import_relation_graph(
repository,
base_id=second_base["id"],
prefix="second",
)
by_slug = await repository.get_adjacent_pages(
[first_base["id"]],
{
first_base["id"]: ["first/seed"],
second_base["id"]: ["second/seed"],
},
direction="outgoing",
)
by_id = await repository.get_adjacent_pages(
[first_base["id"]],
{first_base["id"]: [first_pages["first/target"]["id"]]},
direction="incoming",
)
assert [(row["base_id"], row["neighbor_slug"]) for row in by_slug] == [
(first_base["id"], "first/target")
]
assert [(row["direction"], row["neighbor_slug"]) for row in by_id] == [
("incoming", "first/seed")
]
assert all(row["page"]["base_id"] == first_base["id"] for row in [*by_slug, *by_id])
async with repository._sf() as session:
first_entities = list(
(
await session.execute(
select(AssistantKnowledgeEntityRow).where(
AssistantKnowledgeEntityRow.base_id == first_base["id"]
)
)
).scalars()
)
second_target = (
await session.execute(
select(AssistantKnowledgeEntityRow).where(
AssistantKnowledgeEntityRow.base_id == second_base["id"],
AssistantKnowledgeEntityRow.wiki_page_id
== second_pages["second/target"]["id"],
)
)
).scalar_one()
first_seed = next(
entity
for entity in first_entities
if entity.wiki_page_id == first_pages["first/seed"]["id"]
)
session.add(
AssistantKnowledgeRelationRow(
id="cross-base-relation",
base_id=first_base["id"],
source_entity_id=first_seed.id,
target_entity_id=second_target.id,
predicate="must_not_escape",
description="malformed cross-base edge",
status="active",
)
)
await session.commit()
isolated = await repository.get_adjacent_pages(
[first_base["id"]],
{first_base["id"]: ["first/seed"]},
direction="outgoing",
)
assert [row["neighbor_slug"] for row in isolated] == ["first/target"]
@pytest.mark.asyncio
async def test_graph_adjacency_filters_inactive_relation_entity_and_page(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.create_base(created_by="admin", name="状态过滤图谱")
pages = await _import_relation_graph(repository, base_id=base["id"], prefix="active")
async def adjacent() -> list[dict]:
return await repository.get_adjacent_pages(
[base["id"]],
{base["id"]: ["active/seed"]},
direction="outgoing",
)
assert [row["neighbor_slug"] for row in await adjacent()] == ["active/target"]
async with repository._sf() as session:
relation = (
await session.execute(
select(AssistantKnowledgeRelationRow).where(
AssistantKnowledgeRelationRow.base_id == base["id"]
)
)
).scalar_one()
relation.status = "deleted"
await session.commit()
assert await adjacent() == []
async with repository._sf() as session:
relation = (
await session.execute(
select(AssistantKnowledgeRelationRow).where(
AssistantKnowledgeRelationRow.base_id == base["id"]
)
)
).scalar_one()
target_entity = (
await session.execute(
select(AssistantKnowledgeEntityRow).where(
AssistantKnowledgeEntityRow.wiki_page_id == pages["active/target"]["id"]
)
)
).scalar_one()
relation.status = "active"
target_entity.status = "deleted"
await session.commit()
assert await adjacent() == []
async with repository._sf() as session:
target_entity = (
await session.execute(
select(AssistantKnowledgeEntityRow).where(
AssistantKnowledgeEntityRow.wiki_page_id == pages["active/target"]["id"]
)
)
).scalar_one()
target_page = await session.get(
AssistantKnowledgeWikiPageRow,
pages["active/target"]["id"],
)
target_entity.status = "active"
assert target_page is not None
target_page.status = "deleted"
await session.commit()
assert await adjacent() == []
@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
@pytest.mark.parametrize(
("document_metadata", "expected_title"),
[
({"title": "设备散热操作规程", "file_name": "cooling-guide.pdf"}, "设备散热操作规程"),
({"file_name": "cooling-guide.pdf"}, "cooling-guide.pdf"),
],
)
async def test_imported_chunk_inherits_title_from_its_document(
repository: AssistantKnowledgeRepository,
document_metadata: dict[str, str],
expected_title: str,
) -> None:
base = await repository.initialize(created_by="admin")
job = await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [],
"documents": [
{
"id": "document-one",
**document_metadata,
}
],
"chunks": [
{
"id": "chunk-one",
"source_document_id": "document-one",
"content": "设备过热时应先断电并检查风道。",
}
],
},
source_type="weknora",
source_key="kb-chunk-title",
source_name="ordinary-export.jsonl",
trigger="manual_package_upload",
created_by="admin",
base_id=base["id"],
)
pages = await repository.list_pages(base_id=base["id"], include_chunks=True)
chunk = next(page for page in pages if page["page_type"] == "chunk")
assert job["counts"]["chunk_pages"] == 1
assert chunk["title"] == expected_title
assert chunk["content_markdown"].startswith(f"# {expected_title}\n")
assert "ordinary-export.jsonl 片段 1" not in chunk["content_markdown"]
@pytest.mark.asyncio
async def test_list_bases_does_not_trigger_full_library_alignment(
repository: AssistantKnowledgeRepository,
monkeypatch: pytest.MonkeyPatch,
) -> None:
base = await repository.initialize(created_by="admin")
align = AsyncMock(side_effect=AssertionError("library switcher must not align content"))
monkeypatch.setattr(repository, "_ensure_base_aligned", align)
bases = await repository.list_bases()
assert [item["id"] for item in bases] == [base["id"]]
align.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("custom", [False, True])
async def test_browsing_committed_pages_does_not_wait_for_background_import(
repository: AssistantKnowledgeRepository,
monkeypatch: pytest.MonkeyPatch,
custom: bool,
) -> None:
base = await repository.initialize(created_by="admin")
if custom:
base = await repository.create_base(created_by="admin", name="独立助手库")
await repository.import_package(
package={"wiki_pages": [{"slug": "wiki/ready", "title": "已提交内容", "content": "上次成功导入的正文"}]},
source_type="weknora",
source_key="ready-source",
source_name="已完成来源",
trigger="manual_package_upload",
created_by="admin",
base_id=base["id"],
)
# A fresh Gateway has an empty alignment cache. A different import may
# already own the write lock, but readers must still see the old snapshot.
reader = AssistantKnowledgeRepository(repository._sf)
counters = AsyncMock(side_effect=AssertionError("content reads must not scan library counters"))
monkeypatch.setattr(reader, "_base_stats_by_ids", counters)
async with reader._import_lock:
directory = await asyncio.wait_for(reader.list_pages(base_id=base["id"], include_content=False), timeout=1)
detail = await asyncio.wait_for(reader.get_page("wiki/ready", base_id=base["id"]), timeout=1)
sources = await asyncio.wait_for(reader.list_sources(base_id=base["id"]), timeout=1)
jobs = await asyncio.wait_for(reader.list_jobs(base_id=base["id"]), timeout=1)
assert [row["slug"] for row in directory] == ["wiki/ready"]
assert detail is not None and detail["content_markdown"] == "上次成功导入的正文"
assert len(sources) == 1
assert any(job["status"] == "completed" for job in jobs)
counters.assert_not_awaited()
@pytest.mark.asyncio
async def test_list_bases_loads_all_statistics_in_one_statement(
repository: AssistantKnowledgeRepository,
) -> None:
await repository.initialize(created_by="admin")
await repository.create_base(created_by="admin", name="助手库一")
await repository.create_base(created_by="admin", name="助手库二")
engine = repository._sf.kw["bind"]
statistics_statements: list[str] = []
def record_statement(_connection, _cursor, statement, _parameters, _context, _many) -> None:
if any(
table in statement
for table in (
"assistant_knowledge_sources",
"assistant_knowledge_wiki_pages",
"assistant_knowledge_entities",
"assistant_knowledge_relations",
)
):
statistics_statements.append(statement)
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
try:
bases = await repository.list_bases()
finally:
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
assert len(bases) == 3
assert all(
base[stat] == 0
for base in bases
for stat in ("source_count", "page_count", "entity_count", "relation_count")
)
assert len(statistics_statements) == 1
@pytest.mark.asyncio
async def test_update_base_does_not_scan_statistics_tables(
repository: AssistantKnowledgeRepository,
) -> None:
await repository.initialize(created_by="admin")
base = await repository.create_base(created_by="admin", name="修改前")
engine = repository._sf.kw["bind"]
statistics_statements: list[str] = []
def record_statement(_connection, _cursor, statement, _parameters, _context, _many) -> None:
if any(
table in statement
for table in (
"assistant_knowledge_sources",
"assistant_knowledge_wiki_pages",
"assistant_knowledge_entities",
"assistant_knowledge_relations",
)
):
statistics_statements.append(statement)
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
try:
updated = await repository.update_base(
base["id"],
name="修改后",
description="只修改元数据",
)
finally:
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
assert updated is not None
assert updated["name"] == "修改后"
assert updated["description"] == "只修改元数据"
assert statistics_statements == []
@pytest.mark.asyncio
async def test_page_directory_projection_omits_large_text_columns(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [
{
"slug": "wiki/large",
"title": "大目录页面",
"page_type": "article",
"content": "不应随目录列表返回的超长正文",
"summary": "目录摘要",
"category_path": ["一级", "二级"],
"tags": ["测试"],
},
{
"slug": "zzz/index",
"title": "索引",
"page_type": "index",
"content": "目录正文也不应随轻量列表返回",
},
],
},
source_type="weknora",
source_key="kb-large-directory",
source_name="大目录普通库",
trigger="manual_package_upload",
created_by="admin",
base_id=base["id"],
)
directory = await repository.list_pages(
base_id=base["id"],
include_content=False,
)
assert directory[0]["slug"] == "zzz/index"
article = next(row for row in directory if row["slug"] == "wiki/large")
assert article["category_path"] == ["一级", "二级"]
assert article["tags"] == ["测试"]
assert all("content_markdown" not in row for row in directory)
assert all("content" not in row for row in directory)
assert all("search_text" not in row for row in directory)
detail = await repository.get_page("wiki/large", base_id=base["id"])
assert detail is not None
assert detail["content_markdown"] == "不应随目录列表返回的超长正文"
assert detail["content"] == "不应随目录列表返回的超长正文"
@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_vector_search_keeps_exact_imported_raw_chunk_hit(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [
{
"slug": "wiki/vector-neighbor",
"title": "向量邻居页",
"content": "这是一条没有验收标记的 Wiki 页面。",
}
],
"vectors": [
{
"wiki_slug": "wiki/vector-neighbor",
"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],
}
],
"documents": [
{
"id": "pump-document",
"title": "星河泵站编码模型验收",
}
],
"chunks": [
{
"id": "pump-chunk",
"source_document_id": "pump-document",
"content": "TIGER-928:关闭旁通阀后等待 75 秒,再长按绿色复位键 6 秒。",
}
],
},
source_type="weknora",
source_key="kb-vector-and-raw-chunk",
source_name="向量和原文混合普通库",
trigger="manual_package_upload",
created_by="admin",
base_id=base["id"],
)
results = await repository.search_pages(
base_ids=[base["id"]],
query="TIGER-928 告警后如何复位?",
embedding_client=FakeEmbeddingClient(),
)
assert results[0]["page_type"] == "chunk"
assert results[0]["retrieval_mode"] == "keyword"
assert results[0]["retrieval_channels"] == ["keyword"]
assert "等待 75 秒" in results[0]["matched_content"]
assert "复位键 6 秒" in results[0]["matched_content"]
assert any("vector" in result.get("retrieval_channels", []) for result in results)
@pytest.mark.asyncio
async def test_exact_keyword_body_wins_when_same_page_also_has_vector_hit(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [
{
"slug": "wiki/same-page",
"title": "同页双通道证据",
"content": "TIGER-SAME-441 精确正文:关闭阀门后等待 41 秒。",
}
],
"vectors": [
{
"wiki_slug": "wiki/same-page",
"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],
}
],
},
source_type="weknora",
source_key="kb-same-page-dual-channel",
source_name="同页双通道普通库",
trigger="manual_package_upload",
created_by="admin",
base_id=base["id"],
)
results = await repository.search_pages(
base_ids=[base["id"]],
query="TIGER-SAME-441 如何处理?",
embedding_client=FakeEmbeddingClient(),
)
result = results[0]
assert result["retrieval_channels"] == ["vector", "keyword"]
assert set(result["channel_scores"]) == {"vector", "keyword"}
assert result["retrieval_mode"] == "keyword"
assert result["score"] == result["channel_scores"]["keyword"]
assert "TIGER-SAME-441" in result["matched_content"]
assert "等待 41 秒" in result["matched_section"]["content"]
assert "语义片段只说明设备维护" not in result["content_markdown"]
assert "embedding_model" not in result["matched_section"]
@pytest.mark.asyncio
async def test_overview_query_returns_representative_knowledge_without_encoding(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [
{
"slug": "wiki/index",
"title": "目录",
"page_type": "index",
"content": "这段目录文字不能被当成知识命中。",
},
{
"slug": "wiki/summary",
"title": "摘要",
"page_type": "summary",
"content": "这段摘要文字不能被当成知识命中。",
},
{
"slug": "wiki/cooling",
"title": "设备散热",
"page_type": "concept",
"content": "设备过热时应先断电,再检查风道与散热风扇。",
},
{
"slug": "wiki/inspection",
"title": "巡检制度",
"page_type": "article",
"content": "值班人员每天需要完成设备巡检并登记异常。",
},
],
},
source_type="weknora",
source_key="kb-overview",
source_name="设备运维库",
trigger="manual_package_upload",
created_by="admin",
)
embedding = FakeEmbeddingClient()
embedding.embed_query = AsyncMock(side_effect=AssertionError("overview retrieval must not start the encoder"))
results = await repository.search_pages(
base_ids=[base["id"]],
query="此知识库都讲了什么?",
embedding_client=embedding,
)
assert {row["title"] for row in results} == {"设备散热", "巡检制度"}
assert all(row["retrieval_mode"] == "overview" for row in results)
assert all(row["page_type"] not in {"directory", "index", "summary"} for row in results)
assert all(row["matched_section"]["content"] for row in results)
embedding.embed_query.assert_not_awaited()
@pytest.mark.asyncio
async def test_chinese_keyword_fallback_uses_meaningful_subphrases_without_base_stats(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
await repository.import_package(
package={
"source": {"type": "weknora"},
"wiki_pages": [
{
"slug": "wiki/cooling-procedure",
"title": "设备散热操作规程",
"page_type": "article",
"content": "设备过热处理流程要求先断电,再检查风道。",
}
],
},
source_type="weknora",
source_key="kb-cooling",
source_name="设备运维库",
trigger="manual_package_upload",
created_by="admin",
)
repository.get_base = AsyncMock(side_effect=AssertionError("search must not load base statistics"))
results = await repository.search_pages(
base_ids=[base["id"]],
query="设备过热怎么处理",
embedding_client=None,
)
assert len(results) == 1
assert results[0]["retrieval_mode"] == "keyword"
assert "设备过热处理流程" in results[0]["matched_section"]["content"]
repository.get_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_incompatible_vector_fingerprint_skips_query_encoder(
repository: AssistantKnowledgeRepository,
) -> None:
base = await repository.initialize(created_by="admin")
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],
}
],
},
source_type="weknora",
source_key="kb-incompatible-vector",
source_name="向量普通库",
trigger="manual_package_upload",
created_by="admin",
)
embedding = IncompatibleEmbeddingClient()
assert await repository.search_pages(
base_ids=[base["id"]],
query="完全没有字面重合",
embedding_client=embedding,
) == []
embedding.embed_query.assert_not_awaited()
@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"]