45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
|
# ============================================================
|
||
|
|
# 混合嵌入降级黄金测试
|
||
|
|
# ============================================================
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from server.knowledge.embedding import EmbeddingProvider, EmbeddingStore, _cosine
|
||
|
|
from server.knowledge.retrieval import hybrid_search
|
||
|
|
|
||
|
|
|
||
|
|
def test_embedding_provider_falls_back_to_none(monkeypatch):
|
||
|
|
monkeypatch.setenv("EMBEDDING_PROVIDER", "off")
|
||
|
|
monkeypatch.delenv("EMBEDDING_API_KEY", raising=False)
|
||
|
|
monkeypatch.delenv("LLM_API_KEY", raising=False)
|
||
|
|
p = EmbeddingProvider()
|
||
|
|
assert p.backend == "none"
|
||
|
|
assert p.enabled is False
|
||
|
|
assert p.embed(["hello"]) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_cosine_and_store_search(tmp_path):
|
||
|
|
store = EmbeddingStore(path=str(tmp_path / "emb.json"))
|
||
|
|
store.upsert_many({
|
||
|
|
"a:1": [1.0, 0.0, 0.0],
|
||
|
|
"b:1": [0.0, 1.0, 0.0],
|
||
|
|
"c:1": [0.9, 0.1, 0.0],
|
||
|
|
})
|
||
|
|
hits = store.search([1.0, 0.0, 0.0], top_k=2)
|
||
|
|
assert hits[0][0] == "a:1"
|
||
|
|
assert hits[0][1] == 1.0
|
||
|
|
assert _cosine([1, 0], [0, 1]) == 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def test_hybrid_search_works_without_vectors():
|
||
|
|
units = [{
|
||
|
|
"assetId": "x1", "title": "换线标准SOP", "kind": "sop", "version": "v1",
|
||
|
|
"tags": ["换线"], "content": "产线换型必须遵守夜班不换线规定。",
|
||
|
|
"chunkId": "c1", "heading": None, "page": None,
|
||
|
|
}]
|
||
|
|
hits = hybrid_search(units, "换线夜班", top_k=3)
|
||
|
|
assert hits
|
||
|
|
assert hits[0]["assetId"] == "x1"
|
||
|
|
assert hits[0]["version"] == "v1"
|
||
|
|
miss = hybrid_search(units, "quantum blockchain xyzzy", top_k=3)
|
||
|
|
assert miss == []
|