97 lines
4.9 KiB
Python
97 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from deerflow.config.llmwiki_config import LocalWikiIndexConfig
|
|
from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page
|
|
from deerflow.integrations.weknora.local_index.search import WikiIndexError, WikiVectorSearchService
|
|
from deerflow.integrations.weknora.local_index.vector_cache import WikiVectorCache
|
|
from deerflow.integrations.weknora.local_index.vector_codec import encode_vector
|
|
from deerflow.persistence.llmwiki_index.memory import MemoryLlmWikiIndexStore
|
|
|
|
|
|
class FakeEmbedding:
|
|
dimensions = 2
|
|
fingerprint = "fp"
|
|
|
|
async def embed_query(self, text: str) -> list[float]:
|
|
return [1.0, 0.0]
|
|
|
|
|
|
def test_enabled_index_requires_supported_enabled_embedding() -> None:
|
|
with pytest.raises(ValidationError, match="embedding must be enabled"):
|
|
LocalWikiIndexConfig(enabled=True, embedding={"enabled": False})
|
|
with pytest.raises(ValidationError, match="provider must be openai"):
|
|
LocalWikiIndexConfig(embedding={"provider": "custom"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exact_search_groups_sections_by_wiki_page_and_filters_drafts() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp")
|
|
published = normalize_wiki_page("kb", {"slug": "alice", "title": "Alice", "summary": "Founder", "content": "Alpha", "status": "published"})
|
|
draft = normalize_wiki_page("kb", {"slug": "secret", "title": "Secret", "content": "Draft", "status": "draft"})
|
|
a, _ = encode_vector([1.0, 0.0], dimensions=2)
|
|
b, _ = encode_vector([0.9, 0.1], dimensions=2)
|
|
await store.replace_page_vectors(
|
|
page=published,
|
|
vectors=[
|
|
{"section_index": 0, "heading": "A", "section_content": "Alpha", "content_hash": "a", "embedding_dimensions": 2, "vector_blob": a},
|
|
{"section_index": 1, "heading": "B", "section_content": "Beta", "content_hash": "b", "embedding_dimensions": 2, "vector_blob": b},
|
|
],
|
|
fingerprint="fp",
|
|
sync_id="s",
|
|
)
|
|
await store.replace_page_vectors(page=draft, vectors=[{"section_index": 0, "heading": None, "section_content": "Draft", "content_hash": "d", "embedding_dimensions": 2, "vector_blob": a}], fingerprint="fp", sync_id="s")
|
|
config = LocalWikiIndexConfig(min_similarity=0, top_k_sections=10, top_k_pages=8, max_sections_per_page=2)
|
|
service = WikiVectorSearchService(config, FakeEmbedding(), WikiVectorCache(store, revision_check_seconds=0))
|
|
mapping = {"id": "kb", "name": "People"}
|
|
|
|
result = await service.search("alice", [mapping])
|
|
assert [item["wiki_slug"] for item in result["results"]] == ["alice"]
|
|
assert len(result["results"][0]["matched_sections"]) == 2
|
|
assert "wiki_slug=alice" in result["results"][0]["url"]
|
|
|
|
with_drafts = await service.search("alice", [mapping], include_drafts=True)
|
|
assert {item["wiki_slug"] for item in with_drafts["results"]} == {"alice", "secret"}
|
|
|
|
section_result = await service.search_sections("alice", [mapping], top_k_sections=2, external=True)
|
|
assert [item["content"] for item in section_result["results"]] == ["Alpha", "Beta"]
|
|
assert all(item["wiki_slug"] == "alice" for item in section_result["results"])
|
|
assert all(item["kind"] == "wiki_section" for item in section_result["results"])
|
|
assert section_result["results"][0]["wiki_page_id"]
|
|
assert "page_content" in section_result["results"][0]
|
|
|
|
narrow_service = WikiVectorSearchService(
|
|
LocalWikiIndexConfig(min_similarity=0, top_k_sections=1, top_k_pages=1),
|
|
FakeEmbedding(),
|
|
WikiVectorCache(store, revision_check_seconds=0),
|
|
)
|
|
assert [item["wiki_slug"] for item in (await narrow_service.search("alice", [mapping]))["results"]] == ["alice"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_corrupt_blob_marks_index_for_rebuild() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp")
|
|
page = normalize_wiki_page("kb", {"slug": "broken", "title": "Broken", "content": "Broken", "status": "published"})
|
|
blob, _ = encode_vector([1.0, 0.0], dimensions=2)
|
|
saved = await store.replace_page_vectors(
|
|
page=page,
|
|
vectors=[{"section_index": 0, "heading": None, "section_content": "Broken", "content_hash": "x", "embedding_dimensions": 2, "vector_blob": blob}],
|
|
fingerprint="fp",
|
|
sync_id="s",
|
|
)
|
|
store.vectors[saved["id"]][0]["vector_blob"] = b"broken"
|
|
service = WikiVectorSearchService(
|
|
LocalWikiIndexConfig(min_similarity=0, top_k_sections=10, top_k_pages=8),
|
|
FakeEmbedding(),
|
|
WikiVectorCache(store, revision_check_seconds=0),
|
|
)
|
|
|
|
with pytest.raises(WikiIndexError) as exc_info:
|
|
await service.search("broken", [{"id": "kb", "name": "Broken"}])
|
|
assert exc_info.value.code == "WIKI_INDEX_REBUILDING"
|
|
assert store.states["kb"]["state"] == "rebuild_required"
|