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" def __init__(self) -> None: self.queries: list[str] = [] async def embed_query(self, text: str) -> list[float]: self.queries.append(text) 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_title_boost_recalls_and_reranks_exact_title_page() -> None: store = MemoryLlmWikiIndexStore() await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp") title_page = normalize_wiki_page( "kb", {"slug": "mar-a-lago", "title": "海湖庄园", "content": "庄园资料", "status": "published"}, ) broad_page = normalize_wiki_page( "kb", {"slug": "trump", "title": "特朗普", "content": "人物资料", "status": "published"}, ) title_vector, _ = encode_vector([0.8, 0.6], dimensions=2) broad_vector, _ = encode_vector([1.0, 0.0], dimensions=2) await store.replace_page_vectors( page=title_page, vectors=[ { "section_index": 0, "heading": None, "section_content": "庄园资料", "content_hash": "title", "embedding_dimensions": 2, "vector_blob": title_vector, } ], fingerprint="fp", sync_id="s", ) await store.replace_page_vectors( page=broad_page, vectors=[ { "section_index": 0, "heading": None, "section_content": "人物资料提到海湖庄园", "content_hash": "broad", "embedding_dimensions": 2, "vector_blob": broad_vector, } ], fingerprint="fp", sync_id="s", ) mapping = {"id": "kb", "name": "人物资料库"} boosted = WikiVectorSearchService( LocalWikiIndexConfig(min_similarity=0.85, top_k_sections=2, top_k_pages=2), FakeEmbedding(), WikiVectorCache(store, revision_check_seconds=0), ) legacy = WikiVectorSearchService( LocalWikiIndexConfig( min_similarity=0.85, top_k_sections=2, top_k_pages=2, title_boost_enabled=False, ), FakeEmbedding(), WikiVectorCache(store, revision_check_seconds=0), ) zero_weight = WikiVectorSearchService( LocalWikiIndexConfig( min_similarity=0.85, top_k_sections=2, top_k_pages=2, title_boost_enabled=True, title_boost_weight=0, ), FakeEmbedding(), WikiVectorCache(store, revision_check_seconds=0), ) assert [item["wiki_slug"] for item in (await boosted.search("海湖庄园", [mapping]))["results"]] == [ "mar-a-lago", "trump", ] assert [item["wiki_slug"] for item in (await legacy.search("海湖庄园", [mapping]))["results"]] == ["trump"] assert [item["wiki_slug"] for item in (await zero_weight.search("海湖庄园", [mapping]))["results"]] == ["trump"] @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" @pytest.mark.asyncio async def test_overview_query_uses_representative_body_without_similarity_filter_or_embedding() -> None: store = MemoryLlmWikiIndexStore() await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp") directory = normalize_wiki_page( "kb", {"slug": "index", "title": "目录", "page_type": "directory", "content": "目录内容", "status": "published"}, ) article = normalize_wiki_page( "kb", { "slug": "guides/cooling", "title": "设备散热", "page_type": "article", "summary": "摘要不应替代正文", "content": "设备过热时应先断电，再检查风道和散热风扇。", "status": "published", }, ) unrelated, _ = encode_vector([0.0, 1.0], dimensions=2) for page, content in ((directory, "目录内容"), (article, "设备过热时应先断电，再检查风道和散热风扇。")): await store.replace_page_vectors( page=page, vectors=[ { "section_index": 0, "heading": "页面摘要" if page is article else None, "section_content": "摘要不应替代正文" if page is article else content, "content_hash": page["content_hash"], "embedding_dimensions": 2, "vector_blob": unrelated, } ], fingerprint="fp", sync_id="s", ) embedding = FakeEmbedding() service = WikiVectorSearchService( LocalWikiIndexConfig(min_similarity=0.99, top_k_sections=10, top_k_pages=8), embedding, WikiVectorCache(store, revision_check_seconds=0), ) result = await service.search("此知识库都讲了什么？", [{"id": "kb", "name": "设备库"}]) assert result["mode"] == "overview" assert [item["wiki_slug"] for item in result["results"]] == ["guides/cooling"] assert result["results"][0]["retrieval_mode"] == "overview" assert "设备过热时应先断电" in result["results"][0]["content"] assert "摘要不应替代正文" not in result["results"][0]["content"] assert embedding.queries == []