1276 lines
44 KiB
Python
1276 lines
44 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
|
|
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"]
|