235 lines
10 KiB
Python
235 lines
10 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"
|
|
|
|
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 == []
|