116 lines
5.3 KiB
Python
116 lines
5.3 KiB
Python
"""Explicit opt-in: actual pinned model -> HTTP -> PostgreSQL -> HTTP process restart."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from test_runtime import server
|
|
|
|
from app.services.embedding import SMALL_ZH
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("RUN_LOCAL_EMBEDDING_TEST") != "1",
|
|
reason="requires optional CPU runtime and offline pinned model cache",
|
|
)
|
|
def test_real_local_model_vectors_retrieval_and_http_restart(migrated_database):
|
|
options = {
|
|
"EMBEDDING_ENABLED": "true",
|
|
"EMBEDDING_MODEL": "small",
|
|
"EMBEDDING_CACHE_DIR": os.environ["EMBEDDING_CACHE_DIR"],
|
|
"HF_HUB_OFFLINE": "1",
|
|
"HF_HUB_DISABLE_TELEMETRY": "1",
|
|
"HF_HUB_DISABLE_IMPLICIT_TOKEN": "1",
|
|
"TOKENIZERS_PARALLELISM": "false",
|
|
}
|
|
query = {"query": "合成测试中交通安排应查阅哪份资料？", "top_k": 3}
|
|
with server(migrated_database, environment_overrides=options) as client:
|
|
ids = []
|
|
for topic, body in [
|
|
("交通", "【合成测试】占位交通资料介绍测试公交线路与班次的查询入口。"),
|
|
("住宿", "【合成测试】占位住宿资料介绍测试房型与入住须知。"),
|
|
("审核", "【合成测试】占位审核资料说明测试出处核验流程。"),
|
|
]:
|
|
item = client.post(
|
|
"/api/knowledge",
|
|
json={
|
|
"title": f"{topic}合成测试",
|
|
"content": body,
|
|
"source_url": "https://example.com/",
|
|
"collected_at": "2026-10-16T00:00:00Z",
|
|
},
|
|
).json()
|
|
response = client.post(
|
|
f"/api/knowledge/{item['id']}/review",
|
|
json={
|
|
"expected_version": item["version"],
|
|
"verification_status": "verified",
|
|
"reviewer": "test-reviewer",
|
|
"verification_notes": "仅为模型/HTTP/数据库持久化测试，不是事实资料",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
ids.append(item["id"])
|
|
text_ = client.post("/api/indexes", json={"chunk_size": 200, "overlap": 20}).json()
|
|
assert (
|
|
client.post(
|
|
f"/api/indexes/{text_['id']}/publish", json={"expected_generation": 0}
|
|
).status_code
|
|
== 200
|
|
)
|
|
built = client.post(
|
|
"/api/vector-indexes", json={"text_index_version_id": text_["id"]}, timeout=60
|
|
)
|
|
assert built.status_code == 201, built.text
|
|
version = built.json()
|
|
assert version["space_id"] == SMALL_ZH.space_id and version["dimension"] == 512
|
|
published = client.post(
|
|
f"/api/vector-indexes/{version['id']}/publish", json={"expected_generation": 0}
|
|
)
|
|
assert published.status_code == 200
|
|
active = published.json()
|
|
response = client.post("/api/retrieval/vector", json=query, timeout=60)
|
|
assert response.status_code == 200, response.text
|
|
first = response.json()
|
|
assert {h["knowledge_id"] for h in first["results"]} == set(ids)
|
|
assert all(-1 <= h["score"] <= 1 for h in first["results"])
|
|
assert first["config"]["space"]["spec"]["revision"] == SMALL_ZH.revision
|
|
response = client.post(
|
|
"/api/retrieval/hybrid", json={**query, "candidate_k": 3}, timeout=60
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
hybrid = response.json()
|
|
assert hybrid["vector_version_id"] == version["id"]
|
|
assert hybrid["text_index_version_id"] == text_["id"]
|
|
assert hybrid["config"]["space"]["spec"]["revision"] == SMALL_ZH.revision
|
|
assert {h["knowledge_id"] for h in hybrid["results"]} == set(ids)
|
|
assert hybrid["vector_candidates"]
|
|
keyword = client.post("/api/retrieval/chunks", json=query).json()
|
|
assert keyword["mode"] == "keyword_chunks" and keyword["results"]
|
|
with migrated_database.connect() as c:
|
|
assert (
|
|
c.scalar(
|
|
text("SELECT count(*) FROM chunk_vectors WHERE public.vector_dims(embedding)=512")
|
|
)
|
|
== 3
|
|
)
|
|
assert c.scalar(text("SELECT count(*) FROM vector_publications")) == 1
|
|
with server(migrated_database, environment_overrides=options) as client:
|
|
assert client.get("/api/vector-indexes/active").json() == active
|
|
assert client.get(f"/api/retrieval/runs/{first['run_id']}").json() == first
|
|
response = client.post("/api/retrieval/vector", json=query, timeout=60)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["results"] == first["results"]
|
|
for result, endpoint in [(hybrid, "hybrid"), (keyword, "chunks")]:
|
|
assert client.get(f"/api/retrieval/runs/{result['run_id']}").json() == result
|
|
response = client.post(
|
|
f"/api/retrieval/{endpoint}",
|
|
json={**query, **({"candidate_k": 3} if endpoint == "hybrid" else {})},
|
|
timeout=60,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["results"] == result["results"]
|
|
if endpoint == "hybrid":
|
|
assert response.json()["lexical_candidates"] == result["lexical_candidates"]
|
|
assert response.json()["vector_candidates"] == result["vector_candidates"]
|