275 lines
9.4 KiB
Python
275 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page
|
|
from deerflow.integrations.weknora.local_index.vector_codec import decode_vector, encode_vector
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.llmwiki.model import LlmWikiKnowledgeBaseRow
|
|
from deerflow.persistence.llmwiki_index.model import (
|
|
LlmWikiWikiEdgeRow,
|
|
LlmWikiWikiPageRow,
|
|
LlmWikiWikiSyncStateRow,
|
|
LlmWikiWikiVectorRow,
|
|
)
|
|
from deerflow.persistence.llmwiki_index.sql import SqlLlmWikiIndexStore
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_atomically_replaces_vectors_and_marks_missing_deleted(tmp_path) -> None:
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'wiki.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_all(
|
|
[
|
|
LlmWikiKnowledgeBaseRow(
|
|
id="mapping-1",
|
|
weknora_id="remote-1",
|
|
owner_user_id="u1",
|
|
name="Wiki",
|
|
description="",
|
|
kb_type="document",
|
|
publication_status="private",
|
|
created_at=datetime.now(UTC),
|
|
updated_at=datetime.now(UTC),
|
|
),
|
|
LlmWikiKnowledgeBaseRow(
|
|
id="mapping-2",
|
|
weknora_id="remote-2",
|
|
owner_user_id="u2",
|
|
name="Private Wiki",
|
|
description="",
|
|
kb_type="document",
|
|
publication_status="private",
|
|
created_at=datetime.now(UTC),
|
|
updated_at=datetime.now(UTC),
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
store = SqlLlmWikiIndexStore(session_factory, write_batch_size=1)
|
|
assert await store.try_acquire_sync_lease("mapping-1", owner="worker", sync_id="sync-1", lease_seconds=60, fingerprint="fp")
|
|
assert not await store.try_acquire_sync_lease("mapping-1", owner="worker", sync_id="sync-2", lease_seconds=60, fingerprint="fp")
|
|
page = normalize_wiki_page("mapping-1", {"id": "p1", "slug": "people/alice", "title": "Alice", "content": "Alice founded Acme.", "status": "published"})
|
|
blob, dimensions = encode_vector([3.0, 4.0], dimensions=2)
|
|
saved = await store.replace_page_vectors(
|
|
page=page,
|
|
vectors=[
|
|
{
|
|
"section_index": index,
|
|
"heading": "History",
|
|
"section_content": f"Alice founded Acme. Part {index}",
|
|
"content_hash": f"section-{index}",
|
|
"embedding_dimensions": dimensions,
|
|
"vector_blob": blob,
|
|
}
|
|
for index in range(3)
|
|
],
|
|
fingerprint="fp",
|
|
sync_id="sync-1",
|
|
)
|
|
snapshot = await store.load_vector_snapshot("mapping-1", "fp")
|
|
assert saved["slug"] == "people/alice"
|
|
assert store.write_batch_size == 1
|
|
assert len(snapshot["rows"]) == 3
|
|
assert decode_vector(snapshot["rows"][0]["vector_blob"], 2).tolist() == pytest.approx([0.6, 0.8])
|
|
streamed = [row async for row in store.iter_vector_snapshot_rows("mapping-1", "fp", batch_size=2)]
|
|
assert [row["section_index"] for row in streamed] == [0, 1, 2]
|
|
assert streamed[0]["page"]["slug"] == "people/alice"
|
|
page_statuses = await store.list_page_index_status("mapping-1")
|
|
assert len(page_statuses) == 1
|
|
assert page_statuses[0] == {
|
|
**page_statuses[0],
|
|
"id": saved["id"],
|
|
"slug": "people/alice",
|
|
"title": "Alice",
|
|
"page_type": "article",
|
|
"status": "published",
|
|
"index_status": "ready",
|
|
"vector_count": 3,
|
|
"last_error": None,
|
|
"last_seen_sync_id": "sync-1",
|
|
}
|
|
assert page_statuses[0]["indexed_at"]
|
|
assert page_statuses[0]["updated_at"]
|
|
|
|
second_page = normalize_wiki_page(
|
|
"mapping-1",
|
|
{
|
|
"id": "p2",
|
|
"slug": "companies/acme",
|
|
"title": "Acme",
|
|
"content": "Acme is connected to Alice.",
|
|
"status": "published",
|
|
},
|
|
)
|
|
second = await store.replace_page_vectors(
|
|
page=second_page,
|
|
vectors=[
|
|
{
|
|
"section_index": 0,
|
|
"heading": "Acme",
|
|
"section_content": "Acme is connected to Alice.",
|
|
"content_hash": "acme-section",
|
|
"embedding_dimensions": dimensions,
|
|
"vector_blob": blob,
|
|
}
|
|
],
|
|
fingerprint="fp",
|
|
sync_id="sync-1",
|
|
)
|
|
foreign_page = normalize_wiki_page(
|
|
"mapping-2",
|
|
{
|
|
"id": "foreign-p1",
|
|
"slug": "private/foreign",
|
|
"title": "Foreign",
|
|
"content": "This page belongs to another mapping.",
|
|
"status": "published",
|
|
},
|
|
)
|
|
await store.replace_page_vectors(
|
|
page=foreign_page,
|
|
vectors=[
|
|
{
|
|
"section_index": 0,
|
|
"heading": "Foreign",
|
|
"section_content": "This page belongs to another mapping.",
|
|
"content_hash": "foreign-section",
|
|
"embedding_dimensions": dimensions,
|
|
"vector_blob": blob,
|
|
}
|
|
],
|
|
fingerprint="fp",
|
|
sync_id="sync-foreign",
|
|
)
|
|
stale_page = normalize_wiki_page(
|
|
"mapping-1",
|
|
{
|
|
"id": "p1",
|
|
"slug": "people/alice",
|
|
"title": "Stale Alice",
|
|
"content": "A stale worker must not replace this page.",
|
|
"status": "published",
|
|
},
|
|
)
|
|
assert not await store.update_sync_progress(
|
|
"mapping-1",
|
|
owner="worker-stale",
|
|
sync_id="sync-stale",
|
|
remote_page_count=999,
|
|
)
|
|
assert not await store.touch_page(
|
|
str(saved["id"]),
|
|
"sync-stale",
|
|
mapping_id="mapping-1",
|
|
owner="worker-stale",
|
|
)
|
|
assert await store.replace_page_vectors(
|
|
page=stale_page,
|
|
vectors=[],
|
|
fingerprint="fp",
|
|
sync_id="sync-stale",
|
|
owner="worker-stale",
|
|
) is None
|
|
assert not await store.record_page_failure(
|
|
page=stale_page,
|
|
sync_id="sync-stale",
|
|
error="stale failure",
|
|
owner="worker-stale",
|
|
)
|
|
preserved = await store.get_page("mapping-1", "people/alice")
|
|
assert preserved is not None
|
|
assert preserved["title"] == "Alice"
|
|
assert preserved["last_seen_sync_id"] == "sync-1"
|
|
assert preserved["last_error"] is None
|
|
assert await store.replace_graph_edges(
|
|
"mapping-1",
|
|
[],
|
|
owner="worker-stale",
|
|
sync_id="sync-stale",
|
|
) is None
|
|
assert await store.mark_missing_pages_deleted(
|
|
"mapping-1",
|
|
"sync-stale",
|
|
owner="worker-stale",
|
|
) is None
|
|
|
|
graph = await store.replace_graph_edges(
|
|
"mapping-1",
|
|
[
|
|
{
|
|
"source_slug": saved["slug"],
|
|
"target_slug": second["slug"],
|
|
"edge_type": "wiki_link",
|
|
},
|
|
{
|
|
"source_slug": saved["slug"],
|
|
"target_slug": foreign_page["slug"],
|
|
"edge_type": "wiki_link",
|
|
},
|
|
],
|
|
owner="worker",
|
|
sync_id="sync-1",
|
|
)
|
|
assert graph == {"edge_count": 1, "changed": True}
|
|
assert await store.replace_graph_edges(
|
|
"mapping-1",
|
|
[
|
|
{
|
|
"source_slug": saved["slug"],
|
|
"target_slug": second["slug"],
|
|
"edge_type": "wiki_link",
|
|
}
|
|
],
|
|
owner="worker",
|
|
sync_id="sync-1",
|
|
) == {"edge_count": 1, "changed": False}
|
|
outgoing = await store.get_adjacent_pages(["mapping-1"], {"mapping-1": [saved["slug"]]}, direction="outgoing")
|
|
assert len(outgoing) == 1
|
|
assert outgoing[0]["neighbor_slug"] == second["slug"]
|
|
assert outgoing[0]["page"]["knowledge_base_mapping_id"] == "mapping-1"
|
|
incoming = await store.get_adjacent_pages(["mapping-1"], {"mapping-1": [second["slug"]]}, direction="incoming")
|
|
assert len(incoming) == 1
|
|
assert incoming[0]["neighbor_slug"] == saved["slug"]
|
|
|
|
await store.finish_sync(
|
|
"mapping-1",
|
|
owner="worker",
|
|
sync_id="sync-1",
|
|
success=True,
|
|
remote_page_count=2,
|
|
)
|
|
assert await store.try_acquire_sync_lease(
|
|
"mapping-1",
|
|
owner="worker-new",
|
|
sync_id="sync-2",
|
|
lease_seconds=60,
|
|
fingerprint="fp",
|
|
)
|
|
assert await store.mark_missing_pages_deleted(
|
|
"mapping-1",
|
|
"sync-2",
|
|
owner="worker-new",
|
|
) == 2
|
|
assert (await store.load_vector_snapshot("mapping-1", "fp"))["rows"] == []
|
|
await engine.dispose()
|
|
|
|
|
|
def test_vector_codec_rejects_corrupt_and_non_finite_vectors() -> None:
|
|
with pytest.raises(ValueError, match="NaN"):
|
|
encode_vector([1.0, float("nan")], dimensions=2)
|
|
with pytest.raises(ValueError, match="Corrupt"):
|
|
decode_vector(b"123", 2)
|