63 lines
3.5 KiB
Python
63 lines
3.5 KiB
Python
"""Opt-in acceptance against an actual OpenAI-compatible embedding endpoint."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.llmwiki_config import LocalWikiEmbeddingConfig, LocalWikiIndexConfig
|
|
from deerflow.integrations.weknora.local_index.embedding import StrictWikiEmbeddingClient
|
|
from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page
|
|
from deerflow.integrations.weknora.local_index.search import 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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skipif(not os.getenv("KNOWLEDGE_TEST_EMBEDDING_KEY"), reason="Live embedding credentials not configured")
|
|
async def test_real_embedding_search_preserves_library_scope_and_excludes_drafts():
|
|
client = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(
|
|
provider="openai_compatible",
|
|
base_url=os.getenv("KNOWLEDGE_TEST_EMBEDDING_URL", "https://api.siliconflow.cn/v1"),
|
|
model=os.getenv("KNOWLEDGE_TEST_EMBEDDING_MODEL", "BAAI/bge-m3"),
|
|
api_key=os.environ["KNOWLEDGE_TEST_EMBEDDING_KEY"],
|
|
dimensions=1024,
|
|
)
|
|
)
|
|
content = "设备过热时应先断电，再检查风道和散热风扇。"
|
|
try:
|
|
vectors = await client.embed_texts([content, "读书交流活动通过分享心得促进阅读。"])
|
|
assert [len(vector) for vector in vectors] == [1024, 1024]
|
|
store = MemoryLlmWikiIndexStore()
|
|
for base_id in ("selected", "other"):
|
|
await store.try_acquire_sync_lease(base_id, owner="test", sync_id="sync", lease_seconds=60, fingerprint=client.fingerprint)
|
|
for slug, status, index in (("cooling", "published", 0), ("books", "published", 1), ("draft", "draft", 0)):
|
|
page = normalize_wiki_page(base_id, {"slug": slug, "title": slug, "content": content if index == 0 else "读书交流活动", "status": status})
|
|
blob, _ = encode_vector(vectors[index], dimensions=1024)
|
|
await store.replace_page_vectors(
|
|
page=page,
|
|
vectors=[
|
|
{
|
|
"section_index": 0,
|
|
"heading": slug,
|
|
"section_content": content if index == 0 else "读书交流活动",
|
|
"content_hash": slug,
|
|
"embedding_dimensions": 1024,
|
|
"vector_blob": blob,
|
|
}
|
|
],
|
|
fingerprint=client.fingerprint,
|
|
sync_id="sync",
|
|
)
|
|
service = WikiVectorSearchService(LocalWikiIndexConfig(min_similarity=0, top_k_sections=10, top_k_pages=8), client, WikiVectorCache(store, revision_check_seconds=0))
|
|
result = await service.search("设备太热应该如何处理？", [{"id": "selected", "name": "验收设备库"}])
|
|
assert result["results"][0]["wiki_slug"] == "cooling"
|
|
assert {row["wiki_slug"] for row in result["results"]} == {"cooling", "books"}
|
|
assert all("selected" in row["url"] for row in result["results"])
|
|
overview = await service.search("此知识库都讲了什么？", [{"id": "selected", "name": "验收设备库"}])
|
|
assert overview["mode"] == "overview"
|
|
assert overview["results"]
|
|
finally:
|
|
await client.close()
|