444 lines
15 KiB
Python
444 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import Counter
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.config.llmwiki_config import LocalWikiIndexConfig
|
|
from deerflow.integrations.weknora.local_index.sync import WikiSyncService
|
|
from deerflow.integrations.weknora.local_index.vector_cache import WikiVectorCache
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.llmwiki.model import LlmWikiKnowledgeBaseRow
|
|
from deerflow.persistence.llmwiki_index.base import LlmWikiIndexStore
|
|
from deerflow.persistence.llmwiki_index.memory import MemoryLlmWikiIndexStore
|
|
from deerflow.persistence.llmwiki_index.model import (
|
|
LlmWikiWikiEdgeRow,
|
|
LlmWikiWikiPageRow,
|
|
LlmWikiWikiSyncStateRow,
|
|
LlmWikiWikiVectorRow,
|
|
)
|
|
from deerflow.persistence.llmwiki_index.sql import SqlLlmWikiIndexStore
|
|
|
|
|
|
class RecordingEmbedding:
|
|
dimensions = 2
|
|
fingerprint = "incremental-test-fingerprint"
|
|
|
|
def __init__(self) -> None:
|
|
self.batches: list[list[str]] = []
|
|
self.fail = False
|
|
|
|
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
|
self.batches.append(list(texts))
|
|
if self.fail:
|
|
raise RuntimeError("temporary embedding outage")
|
|
return [[1.0, 0.0] for _ in texts]
|
|
|
|
|
|
class VersionedWikiClient:
|
|
def __init__(self) -> None:
|
|
self.pages: dict[str, dict[str, Any]] = {
|
|
"one": self.page("one", version=1, marker="ONE-V1"),
|
|
"two": self.page("two", version=1, marker="TWO-V1"),
|
|
}
|
|
self.detail_calls: list[str] = []
|
|
|
|
@staticmethod
|
|
def page(
|
|
slug: str,
|
|
*,
|
|
version: int,
|
|
marker: str,
|
|
category_path: list[str] | None = None,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"id": f"remote-{slug}",
|
|
"slug": slug,
|
|
"title": slug.title(),
|
|
"summary": f"Summary for {marker} with enough text to index.",
|
|
"content": f"# {slug.title()}\nContent marker {marker} with enough text to index.",
|
|
"status": "published",
|
|
"version": version,
|
|
"updated_at": f"2026-09-{version + 1:02d}T00:00:00Z",
|
|
"category_path": category_path or ["root"],
|
|
}
|
|
|
|
async def get_knowledge_base(self, knowledge_base_id: str) -> dict[str, Any]:
|
|
return {"id": knowledge_base_id}
|
|
|
|
async def list_wiki_pages(
|
|
self,
|
|
knowledge_base_id: str,
|
|
*,
|
|
page: int,
|
|
page_size: int,
|
|
) -> dict[str, Any]:
|
|
del knowledge_base_id, page, page_size
|
|
summaries = [
|
|
{
|
|
key: value
|
|
for key, value in item.items()
|
|
if key
|
|
in {
|
|
"id",
|
|
"slug",
|
|
"title",
|
|
"status",
|
|
"version",
|
|
"updated_at",
|
|
}
|
|
}
|
|
for item in self.pages.values()
|
|
]
|
|
return {
|
|
"pages": summaries,
|
|
"total": len(summaries),
|
|
"total_pages": 1,
|
|
}
|
|
|
|
async def get_wiki_page(
|
|
self,
|
|
knowledge_base_id: str,
|
|
slug: str,
|
|
) -> dict[str, Any]:
|
|
del knowledge_base_id
|
|
self.detail_calls.append(slug)
|
|
return dict(self.pages[slug])
|
|
|
|
|
|
class BodyInSummaryWikiClient(VersionedWikiClient):
|
|
async def list_wiki_pages(
|
|
self,
|
|
knowledge_base_id: str,
|
|
*,
|
|
page: int,
|
|
page_size: int,
|
|
) -> dict[str, Any]:
|
|
del knowledge_base_id, page, page_size
|
|
return {
|
|
"pages": [dict(item) for item in self.pages.values()],
|
|
"total": len(self.pages),
|
|
"total_pages": 1,
|
|
}
|
|
|
|
async def get_wiki_page(
|
|
self,
|
|
knowledge_base_id: str,
|
|
slug: str,
|
|
) -> dict[str, Any]:
|
|
del knowledge_base_id, slug
|
|
raise AssertionError("a list item with a complete body must not need a detail request")
|
|
|
|
|
|
class ProcessingTransitionWikiClient(VersionedWikiClient):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.processing_sequence: list[bool] = []
|
|
|
|
async def get_knowledge_base(self, knowledge_base_id: str) -> dict[str, Any]:
|
|
processing = (
|
|
self.processing_sequence.pop(0)
|
|
if self.processing_sequence
|
|
else False
|
|
)
|
|
return {
|
|
"id": knowledge_base_id,
|
|
"processing_count": 1 if processing else 0,
|
|
}
|
|
|
|
|
|
def _service(
|
|
store: LlmWikiIndexStore,
|
|
embedding: RecordingEmbedding,
|
|
client: VersionedWikiClient,
|
|
) -> WikiSyncService:
|
|
config = LocalWikiIndexConfig(
|
|
chunk_min_chars=10,
|
|
min_similarity=0,
|
|
sync_lease_seconds=60,
|
|
)
|
|
return WikiSyncService(
|
|
config,
|
|
store,
|
|
embedding,
|
|
WikiVectorCache(store),
|
|
client,
|
|
)
|
|
|
|
|
|
_MAPPING = {"id": "kb", "weknora_id": "remote-kb", "name": "Wiki"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_summary_body_avoids_detail_requests_for_new_and_modified_pages() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = BodyInSummaryWikiClient()
|
|
service = _service(store, embedding, client)
|
|
|
|
initial = await service.sync_mapping(_MAPPING)
|
|
assert initial["changed_pages"] == 2
|
|
assert len(embedding.batches) == 2
|
|
|
|
embedding.batches.clear()
|
|
client.pages["one"] = client.page("one", version=2, marker="ONE-V2")
|
|
modified = await service.sync_mapping(_MAPPING)
|
|
assert modified["changed_pages"] == 1
|
|
assert len(embedding.batches) == 1
|
|
assert "ONE-V2" in (await store.get_page("kb", "one"))["content_md"]
|
|
|
|
embedding.batches.clear()
|
|
client.pages["three"] = client.page("three", version=1, marker="THREE-V1")
|
|
added = await service.sync_mapping(_MAPPING)
|
|
assert added["changed_pages"] == 1
|
|
assert len(embedding.batches) == 1
|
|
assert "THREE-V1" in (await store.get_page("kb", "three"))["content_md"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_sync_fetches_and_embeds_only_new_or_changed_pages() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = VersionedWikiClient()
|
|
service = _service(store, embedding, client)
|
|
|
|
first = await service.sync_mapping(_MAPPING)
|
|
assert first["changed_pages"] == 2
|
|
assert Counter(client.detail_calls) == Counter({"one": 1, "two": 1})
|
|
assert len(embedding.batches) == 2
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
unchanged = await service.sync_mapping(_MAPPING)
|
|
assert unchanged["changed_pages"] == 0
|
|
assert client.detail_calls == []
|
|
assert embedding.batches == []
|
|
|
|
client.pages["three"] = client.page("three", version=1, marker="THREE-V1")
|
|
added = await service.sync_mapping(_MAPPING)
|
|
assert added["changed_pages"] == 1
|
|
assert client.detail_calls == ["three"]
|
|
assert len(embedding.batches) == 1
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
client.pages["two"] = client.page("two", version=2, marker="TWO-V2")
|
|
modified = await service.sync_mapping(_MAPPING)
|
|
assert modified["changed_pages"] == 1
|
|
assert client.detail_calls == ["two"]
|
|
assert len(embedding.batches) == 1
|
|
local_two = await store.get_page("kb", "two")
|
|
assert local_two is not None
|
|
assert "TWO-V2" in local_two["content_md"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_incremental_sync_compares_serialized_timestamps_without_refetching(
|
|
tmp_path,
|
|
) -> None:
|
|
engine = create_async_engine(
|
|
f"sqlite+aiosqlite:///{tmp_path / 'incremental-sync.db'}"
|
|
)
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(
|
|
lambda sync_connection: Base.metadata.create_all(
|
|
sync_connection,
|
|
tables=[
|
|
LlmWikiKnowledgeBaseRow.__table__,
|
|
LlmWikiWikiPageRow.__table__,
|
|
LlmWikiWikiVectorRow.__table__,
|
|
LlmWikiWikiEdgeRow.__table__,
|
|
LlmWikiWikiSyncStateRow.__table__,
|
|
],
|
|
)
|
|
)
|
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
async with session_factory() as session:
|
|
session.add(
|
|
LlmWikiKnowledgeBaseRow(
|
|
id="kb",
|
|
weknora_id="remote-kb",
|
|
owner_user_id="owner",
|
|
name="Wiki",
|
|
description="",
|
|
kb_type="document",
|
|
publication_status="private",
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
store = SqlLlmWikiIndexStore(session_factory)
|
|
embedding = RecordingEmbedding()
|
|
client = VersionedWikiClient()
|
|
service = _service(store, embedding, client)
|
|
await service.sync_mapping(_MAPPING)
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
result = await service.sync_mapping(_MAPPING)
|
|
|
|
assert result["changed_pages"] == 0
|
|
assert client.detail_calls == []
|
|
assert embedding.batches == []
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_sync_updates_metadata_without_reembedding_body() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = VersionedWikiClient()
|
|
service = _service(store, embedding, client)
|
|
await service.sync_mapping(_MAPPING)
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
current = dict(client.pages["one"])
|
|
current.update(
|
|
{
|
|
"version": 2,
|
|
"updated_at": "2026-09-03T00:00:00Z",
|
|
"category_path": ["root", "moved"],
|
|
}
|
|
)
|
|
client.pages["one"] = current
|
|
|
|
result = await service.sync_mapping(_MAPPING)
|
|
|
|
assert client.detail_calls == ["one"]
|
|
assert embedding.batches == []
|
|
assert result["changed_pages"] == 0
|
|
local = await store.get_page("kb", "one")
|
|
assert local is not None
|
|
assert local["page_version"] == 2
|
|
assert local["category_path"] == ["root", "moved"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_sync_deletes_only_missing_pages_without_refetching_survivors() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = VersionedWikiClient()
|
|
service = _service(store, embedding, client)
|
|
await service.sync_mapping(_MAPPING)
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
del client.pages["one"]
|
|
result = await service.sync_mapping(_MAPPING)
|
|
|
|
assert result["deleted_pages"] == 1
|
|
assert client.detail_calls == []
|
|
assert embedding.batches == []
|
|
deleted = next(page for page in store.pages.values() if page["slug"] == "one")
|
|
assert deleted["is_remote_deleted"] is True
|
|
assert deleted["id"] not in store.vectors
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_processing_transition_preserves_missing_pages_until_a_quiet_scan() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = ProcessingTransitionWikiClient()
|
|
service = _service(store, embedding, client)
|
|
await service.sync_mapping(_MAPPING)
|
|
|
|
before = await store.get_page("kb", "one")
|
|
assert before is not None
|
|
old_vectors = list(store.vectors[before["id"]])
|
|
del client.pages["one"]
|
|
client.processing_sequence = [False, True]
|
|
|
|
deferred = await service.sync_mapping(_MAPPING)
|
|
|
|
assert deferred["status"] == "partial"
|
|
assert deferred["deleted_pages"] == 0
|
|
assert deferred["deferred_deletion"] is True
|
|
assert deferred["deferred_reason"] == "upstream_processing"
|
|
retained = await store.get_page("kb", "one")
|
|
assert retained is not None
|
|
assert retained["is_remote_deleted"] is False
|
|
assert store.vectors[retained["id"]] == old_vectors
|
|
assert store.states["kb"]["last_sync_result"]["sync_id"] == deferred["sync_id"]
|
|
|
|
client.processing_sequence = [False, False]
|
|
completed = await service.sync_mapping(_MAPPING)
|
|
assert completed["status"] == "completed"
|
|
assert completed["deleted_pages"] == 1
|
|
retired = next(page for page in store.pages.values() if page["slug"] == "one")
|
|
assert retired["is_remote_deleted"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_changed_page_is_retried_and_old_generation_stays_searchable() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = RecordingEmbedding()
|
|
client = VersionedWikiClient()
|
|
service = _service(store, embedding, client)
|
|
await service.sync_mapping(_MAPPING)
|
|
before = await store.get_page("kb", "one")
|
|
assert before is not None
|
|
old_hash = before["content_hash"]
|
|
old_vectors = list(store.vectors[before["id"]])
|
|
|
|
client.pages["one"] = client.page("one", version=2, marker="ONE-V2")
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
embedding.fail = True
|
|
partial = await service.sync_mapping(_MAPPING)
|
|
assert partial["status"] == "partial"
|
|
retained = await store.get_page("kb", "one")
|
|
assert retained is not None
|
|
assert retained["content_hash"] == old_hash
|
|
assert retained["index_status"] == "ready"
|
|
assert store.vectors[retained["id"]] == old_vectors
|
|
|
|
client.detail_calls.clear()
|
|
embedding.batches.clear()
|
|
embedding.fail = False
|
|
recovered = await service.sync_mapping(_MAPPING)
|
|
assert recovered["status"] == "completed"
|
|
assert recovered["changed_pages"] == 1
|
|
assert client.detail_calls == ["one"]
|
|
assert len(embedding.batches) == 1
|
|
updated = await store.get_page("kb", "one")
|
|
assert updated is not None
|
|
assert updated["content_hash"] != old_hash
|
|
assert "ONE-V2" in updated["content_md"]
|
|
|
|
|
|
class BlockingWikiClient(VersionedWikiClient):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.started = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
|
|
async def get_wiki_page(
|
|
self,
|
|
knowledge_base_id: str,
|
|
slug: str,
|
|
) -> dict[str, Any]:
|
|
self.started.set()
|
|
await self.release.wait()
|
|
return await super().get_wiki_page(knowledge_base_id, slug)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_incremental_sync_is_deduplicated_by_mapping_lease() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = BlockingWikiClient()
|
|
first_service = _service(store, RecordingEmbedding(), client)
|
|
second_service = _service(store, RecordingEmbedding(), client)
|
|
|
|
first_task = asyncio.create_task(first_service.sync_mapping(_MAPPING))
|
|
await asyncio.wait_for(client.started.wait(), timeout=2)
|
|
duplicate = await second_service.sync_mapping(_MAPPING)
|
|
assert duplicate == {"mapping_id": "kb", "status": "already_running"}
|
|
|
|
client.release.set()
|
|
first = await asyncio.wait_for(first_task, timeout=2)
|
|
assert first["status"] == "completed"
|