sync: 鍚屾灞€鍩熺綉鏈€鏂颁唬鐮?(source master @ b69568e4)
This commit is contained in:
gozeros committed
1 parent
4f7f07504a
commit
8bc8b0e082
1651 files changed
+278599
-11416
No files matched your search
@@ -0,0 +1,249 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository
|
||||
from deerflow.persistence.base import Base
|
||||
|
||||
|
||||
class SnapshotEmbeddingClient:
|
||||
dimensions = 3
|
||||
fingerprint = "snapshot-test-model"
|
||||
|
||||
async def embed_query(self, _query: str) -> list[float]:
|
||||
return [1.0, 0.0, 0.0]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def repository(tmp_path) -> AsyncIterator[AssistantKnowledgeRepository]:
|
||||
# A file-backed database gives the atomic-switch test a second connection
|
||||
# that can read the last committed snapshot while the replacement
|
||||
# transaction is deliberately paused before commit.
|
||||
database_path = (tmp_path / "assistant-snapshot.sqlite3").as_posix()
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{database_path}")
|
||||
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()
|
||||
|
||||
|
||||
def _package(*, revision: int, marker: str) -> dict:
|
||||
return {
|
||||
"source": {
|
||||
"provider": "weknora",
|
||||
"type": "weknora",
|
||||
"knowledge_base_id": "stable-kb",
|
||||
"knowledge_base_name": "稳定普通知识库",
|
||||
"index_revision": revision,
|
||||
},
|
||||
"wiki_pages": [
|
||||
{
|
||||
"slug": "wiki/stable",
|
||||
"title": "稳定页面",
|
||||
"page_type": "article",
|
||||
"content": f"# 稳定页面\n\n{marker}",
|
||||
}
|
||||
],
|
||||
"vectors": [
|
||||
{
|
||||
"wiki_slug": "wiki/stable",
|
||||
"section_index": 0,
|
||||
"section_content": marker,
|
||||
"embedding_model": "snapshot-test-model",
|
||||
"embedding_fingerprint": SnapshotEmbeddingClient.fingerprint,
|
||||
"embedding_dimensions": SnapshotEmbeddingClient.dimensions,
|
||||
"vector": [1.0, 0.0, 0.0],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
async def _seed_snapshot(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
) -> tuple[dict, dict, dict]:
|
||||
base = await repository.create_base(created_by="admin", name="助手快照测试库")
|
||||
job = await repository.import_package(
|
||||
package=_package(revision=1, marker="OLD-SNAPSHOT-MARKER-991"),
|
||||
source_type="weknora",
|
||||
source_key="stable-kb",
|
||||
source_name="稳定普通知识库",
|
||||
trigger="initial_import",
|
||||
created_by="admin",
|
||||
mapping_id="stable-kb",
|
||||
base_id=base["id"],
|
||||
mirror_to_global=False,
|
||||
)
|
||||
source = (await repository.list_sources(base_id=base["id"]))[0]
|
||||
assert source["last_import_job_id"] == job["id"]
|
||||
return base, source, job
|
||||
|
||||
|
||||
async def _assert_vector_snapshot(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
*,
|
||||
base_id: str,
|
||||
marker: str,
|
||||
) -> None:
|
||||
results = await repository.search_pages(
|
||||
base_ids=[base_id],
|
||||
query="完全没有字面重合的向量查询",
|
||||
embedding_client=SnapshotEmbeddingClient(),
|
||||
)
|
||||
assert results
|
||||
assert results[0]["matched_section"]["content"] == marker
|
||||
|
||||
|
||||
async def _queue_reimport(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
*,
|
||||
base_id: str,
|
||||
revision: int,
|
||||
) -> dict:
|
||||
return await repository.create_queued_import_job(
|
||||
source_type="weknora",
|
||||
source_key="stable-kb",
|
||||
source_name="稳定普通知识库（待发布）",
|
||||
trigger="source_reimport",
|
||||
created_by="admin",
|
||||
metadata={
|
||||
"provider": "weknora",
|
||||
"knowledge_base_id": "stable-kb",
|
||||
"knowledge_base_name": "稳定普通知识库（待发布）",
|
||||
"index_revision": revision,
|
||||
},
|
||||
mapping_id="stable-kb",
|
||||
base_id=base_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queued_and_running_reimport_keep_last_successful_snapshot(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
) -> None:
|
||||
base, original_source, original_job = await _seed_snapshot(repository)
|
||||
|
||||
queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2)
|
||||
queued_source = await repository.get_source(original_source["id"])
|
||||
|
||||
assert queued_job["status"] == "queued"
|
||||
assert queued_source is not None
|
||||
assert queued_source["last_import_job_id"] == original_job["id"]
|
||||
assert queued_source["metadata"]["index_revision"] == 1
|
||||
assert queued_source["source_name"] == "稳定普通知识库"
|
||||
await _assert_vector_snapshot(
|
||||
repository,
|
||||
base_id=base["id"],
|
||||
marker="OLD-SNAPSHOT-MARKER-991",
|
||||
)
|
||||
|
||||
await repository.update_import_job(
|
||||
queued_job["id"],
|
||||
status="running",
|
||||
phase="exporting",
|
||||
)
|
||||
running_source = await repository.get_source(original_source["id"])
|
||||
assert running_source is not None
|
||||
assert running_source["last_import_job_id"] == original_job["id"]
|
||||
await _assert_vector_snapshot(
|
||||
repository,
|
||||
base_id=base["id"],
|
||||
marker="OLD-SNAPSHOT-MARKER-991",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_reimport_rolls_back_and_keeps_old_vector_snapshot(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
) -> None:
|
||||
base, original_source, original_job = await _seed_snapshot(repository)
|
||||
queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2)
|
||||
invalid_package = _package(revision=2, marker="BROKEN-NEW-SNAPSHOT")
|
||||
invalid_package["vectors"][0]["wiki_slug"] = "wiki/missing"
|
||||
|
||||
with pytest.raises(ValueError, match="向量引用不存在"):
|
||||
await repository.import_package_for_job(
|
||||
queued_job["id"],
|
||||
package=invalid_package,
|
||||
created_by="admin",
|
||||
mirror_to_global=False,
|
||||
)
|
||||
await repository.update_import_job(
|
||||
queued_job["id"],
|
||||
status="failed",
|
||||
phase="failed",
|
||||
error="测试导入失败",
|
||||
completed=True,
|
||||
)
|
||||
|
||||
failed_source = await repository.get_source(original_source["id"])
|
||||
assert failed_source is not None
|
||||
assert failed_source["last_import_job_id"] == original_job["id"]
|
||||
assert failed_source["metadata"]["index_revision"] == 1
|
||||
await _assert_vector_snapshot(
|
||||
repository,
|
||||
base_id=base["id"],
|
||||
marker="OLD-SNAPSHOT-MARKER-991",
|
||||
)
|
||||
pages = await repository.list_pages(base_id=base["id"])
|
||||
assert pages[0]["content_markdown"].endswith("OLD-SNAPSHOT-MARKER-991")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_reimport_switches_snapshot_only_at_commit(
|
||||
repository: AssistantKnowledgeRepository,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
base, original_source, original_job = await _seed_snapshot(repository)
|
||||
queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2)
|
||||
|
||||
snapshot_written = asyncio.Event()
|
||||
allow_commit = asyncio.Event()
|
||||
original_retire_orphans = repository._retire_orphans
|
||||
|
||||
async def pause_before_snapshot_switch(session, base_id: str) -> None:
|
||||
await original_retire_orphans(session, base_id)
|
||||
snapshot_written.set()
|
||||
await allow_commit.wait()
|
||||
|
||||
monkeypatch.setattr(repository, "_retire_orphans", pause_before_snapshot_switch)
|
||||
import_task = asyncio.create_task(
|
||||
repository.import_package_for_job(
|
||||
queued_job["id"],
|
||||
package=_package(revision=2, marker="NEW-SNAPSHOT-MARKER-772"),
|
||||
created_by="admin",
|
||||
mirror_to_global=False,
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.wait_for(snapshot_written.wait(), timeout=5)
|
||||
try:
|
||||
# The replacement rows exist only in the writer's uncommitted
|
||||
# transaction. Other readers must still resolve the old job pointer and
|
||||
# the old vector snapshot until the transaction commits.
|
||||
mid_import_source = await repository.get_source(original_source["id"])
|
||||
assert mid_import_source is not None
|
||||
assert mid_import_source["last_import_job_id"] == original_job["id"]
|
||||
assert mid_import_source["metadata"]["index_revision"] == 1
|
||||
await _assert_vector_snapshot(
|
||||
repository,
|
||||
base_id=base["id"],
|
||||
marker="OLD-SNAPSHOT-MARKER-991",
|
||||
)
|
||||
finally:
|
||||
allow_commit.set()
|
||||
|
||||
completed_job = await asyncio.wait_for(import_task, timeout=5)
|
||||
published_source = await repository.get_source(original_source["id"])
|
||||
assert completed_job["status"] == "completed"
|
||||
assert published_source is not None
|
||||
assert published_source["last_import_job_id"] == queued_job["id"]
|
||||
assert published_source["metadata"]["index_revision"] == 2
|
||||
assert published_source["source_name"] == "稳定普通知识库"
|
||||
await _assert_vector_snapshot(
|
||||
repository,
|
||||
base_id=base["id"],
|
||||
marker="NEW-SNAPSHOT-MARKER-772",
|
||||
)
|
||||
Reference in new issue
Block a user