213 lines
8.2 KiB
Python
213 lines
8.2 KiB
Python
import math
|
||
from collections import Counter
|
||
from uuid import uuid4
|
||
|
||
import pytest
|
||
from sqlalchemy import text
|
||
|
||
from app.services import retrieval
|
||
|
||
|
||
def add_item(client, *, status="verified", **fields):
|
||
response = client.post(
|
||
"/api/knowledge",
|
||
json={
|
||
"title": "检索占位资料",
|
||
"content": "那拉提旅游资源检索测试,不代表真实文旅事实。",
|
||
"source_url": "https://example.com/",
|
||
"collected_at": "2026-10-15T09:00:00+08:00",
|
||
**fields,
|
||
},
|
||
)
|
||
assert response.status_code == 201, response.text
|
||
item = response.json()
|
||
if status != "pending":
|
||
response = client.post(
|
||
f"/api/knowledge/{item['id']}/review",
|
||
json={
|
||
"expected_version": item["version"],
|
||
"verification_status": status,
|
||
"reviewer": "test-reviewer",
|
||
"verification_notes": "仅为自动化审核流程测试",
|
||
},
|
||
)
|
||
assert response.status_code == 200, response.text
|
||
item = response.json()
|
||
return item
|
||
|
||
|
||
def search(client, query="那拉提", **fields):
|
||
response = client.post("/api/retrieval", json={"query": query, **fields})
|
||
assert response.status_code == 200, response.text
|
||
return response.json()
|
||
|
||
|
||
def test_bm25_matches_hand_calculated_scores_and_query_deduplication():
|
||
counters = [Counter({"x": 1}), Counter({"x": 1, "y": 1})]
|
||
scores = retrieval.bm25(counters, ["x"])
|
||
idf = math.log(1 + 0.5 / 2.5)
|
||
assert scores[0] == pytest.approx(idf * 2.2 / (1 + 1.2 * (0.25 + 0.75 / 1.5)))
|
||
assert scores[1] == pytest.approx(idf * 2.2 / (1 + 1.2 * (0.25 + 0.75 * 2 / 1.5)))
|
||
assert scores[0] > scores[1]
|
||
assert retrieval.bm25(counters, ["x", "x"]) == scores
|
||
assert retrieval.bm25([], ["x"]) == []
|
||
assert retrieval.bm25([Counter()], ["x"]) == [0.0]
|
||
|
||
|
||
def test_tokenization_and_exact_original_snippet_offsets():
|
||
terms = retrieval.tokenize("请问那拉提 API api!")
|
||
assert "那拉提" in terms and terms.count("api") == 2
|
||
assert "请问" not in terms
|
||
original = "占位前文。" * 80 + "这里有API检索说明。" + "占位后文。" * 80
|
||
snippet, start, end = retrieval.excerpt(original, ["api"])
|
||
assert start > 0 and "API" in snippet
|
||
assert snippet == original[start:end] and len(snippet) <= 240
|
||
|
||
|
||
def test_only_reviewed_active_knowledge_is_retrieved(client):
|
||
good = add_item(client)
|
||
add_item(client, status="pending")
|
||
add_item(client, status="rejected")
|
||
deleted = add_item(client)
|
||
assert (
|
||
client.delete(
|
||
f"/api/knowledge/{deleted['id']}?expected_version={deleted['version']}"
|
||
).status_code
|
||
== 204
|
||
)
|
||
result = search(client)
|
||
assert result["corpus_size"] == 1 and result["total_hits"] == 1
|
||
assert result["results"][0]["knowledge_id"] == good["id"]
|
||
assert result["results"][0]["knowledge_version"] == good["version"]
|
||
assert result["results"][0]["source_url"] == "https://example.com/"
|
||
assert result["elapsed_ms"] >= 0
|
||
|
||
|
||
def test_tied_results_category_top_k_and_reproducibility(client):
|
||
first = add_item(client, category="交通出行")
|
||
second = add_item(client, category="交通出行")
|
||
third = add_item(client, category="民宿服务")
|
||
result = search(client, top_k=2)
|
||
again = search(client, top_k=2)
|
||
assert result["results"] == again["results"]
|
||
assert result["corpus_fingerprint"] == again["corpus_fingerprint"]
|
||
assert result["config"] == again["config"]
|
||
assert result["run_id"] != again["run_id"]
|
||
assert result["total_hits"] == 3
|
||
assert [hit["knowledge_id"] for hit in result["results"]] == [first["id"], second["id"]]
|
||
filtered = search(client, category="民宿服务")
|
||
assert filtered["corpus_size"] == 1
|
||
assert filtered["results"][0]["knowledge_id"] == third["id"]
|
||
|
||
|
||
def test_title_terms_have_more_weight_than_body_terms(client):
|
||
title_match = add_item(client, title="旅游", content="交通")
|
||
body_match = add_item(client, title="交通", content="旅游")
|
||
result = search(client, query="旅游")
|
||
assert [hit["knowledge_id"] for hit in result["results"]] == [
|
||
title_match["id"],
|
||
body_match["id"],
|
||
]
|
||
assert result["results"][0]["score"] > result["results"][1]["score"]
|
||
|
||
|
||
def test_no_evidence_and_history_are_persisted(client):
|
||
empty = search(client)
|
||
assert empty["status"] == "no_evidence" and empty["results"] == []
|
||
add_item(client)
|
||
unmatched = search(client, query="unmatcheduniquelexeme")
|
||
assert unmatched["status"] == "no_evidence" and unmatched["corpus_size"] == 1
|
||
assert client.get(f"/api/retrieval/runs/{empty['run_id']}").json() == empty
|
||
history = client.get("/api/retrieval/runs?page_size=1").json()
|
||
assert history["total"] == 2 and len(history["items"]) == 1
|
||
assert history["items"][0] == unmatched
|
||
assert client.get(f"/api/retrieval/runs/{uuid4()}").status_code == 404
|
||
assert client.get("/api/retrieval/runs/invalid").status_code == 422
|
||
assert client.get("/api/retrieval/runs?page_size=101").status_code == 422
|
||
|
||
|
||
def test_source_snapshot_and_original_run_survive_invalidation(client):
|
||
src = client.post(
|
||
"/api/sources",
|
||
json={
|
||
"name": "测试来源",
|
||
"source_url": "https://example.com/",
|
||
},
|
||
).json()
|
||
content = "占位前言。" * 80 + "那拉提检索关键词。" + "占位后文。" * 80
|
||
doc = client.post(
|
||
"/api/documents",
|
||
json={
|
||
"source_id": src["id"],
|
||
"title": "测试文档",
|
||
"content": content,
|
||
"collected_at": "2026-10-15T09:00:00+08:00",
|
||
},
|
||
).json()
|
||
item = add_item(client, document_id=doc["id"], content=content)
|
||
original = search(client)
|
||
hit = original["results"][0]
|
||
assert hit["document_id"] == doc["id"] and hit["document_version"] == 1
|
||
assert hit["source_id"] == src["id"] and hit["source_version"] == 1
|
||
assert hit["source_name"] == "测试来源"
|
||
assert "那拉提" in hit["snippet"] and hit["snippet_start"] > 0
|
||
assert hit["snippet"] == content[hit["snippet_start"] : hit["snippet_end"]]
|
||
assert (
|
||
client.patch(
|
||
f"/api/sources/{src['id']}",
|
||
json={
|
||
"expected_version": src["version"],
|
||
"name": "修订来源名称",
|
||
},
|
||
).status_code
|
||
== 200
|
||
)
|
||
revised = search(client)
|
||
assert revised["corpus_size"] == 0 and revised["results"] == []
|
||
assert revised["corpus_fingerprint"] != original["corpus_fingerprint"]
|
||
stored = client.get(f"/api/retrieval/runs/{original['run_id']}").json()
|
||
assert stored == original and stored["results"][0]["knowledge_id"] == item["id"]
|
||
|
||
|
||
def test_inactive_parent_is_excluded_even_if_status_was_not_reset(client, migrated_database):
|
||
src = client.post(
|
||
"/api/sources", json={"name": "测试", "source_url": "https://example.com/"}
|
||
).json()
|
||
add_item(client, source_id=src["id"])
|
||
assert search(client)["corpus_size"] == 1
|
||
with migrated_database.begin() as connection:
|
||
connection.execute(
|
||
text("UPDATE sources SET is_deleted = true WHERE id = :id"), {"id": src["id"]}
|
||
)
|
||
assert search(client)["corpus_size"] == 0
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"payload",
|
||
[
|
||
{"query": " "},
|
||
{"query": "!!!"},
|
||
{"query": "的了"},
|
||
{"query": "那拉提", "top_k": 0},
|
||
{"query": "那拉提", "top_k": 21},
|
||
{"query": "那拉提", "mode": "vector"},
|
||
{"query": "那拉提", "unexpected": True},
|
||
{"query": "a" * 501},
|
||
],
|
||
)
|
||
def test_invalid_query_does_not_create_run(client, payload):
|
||
assert client.post("/api/retrieval", json=payload).status_code == 422
|
||
assert client.get("/api/retrieval/runs").json()["total"] == 0
|
||
|
||
|
||
@pytest.mark.parametrize("setting,value", [("MAX_ITEMS", 1), ("MAX_CHARACTERS", 1)])
|
||
def test_corpus_budget_fails_explicitly_without_silent_truncation(
|
||
client, monkeypatch, setting, value
|
||
):
|
||
add_item(client)
|
||
add_item(client)
|
||
monkeypatch.setattr(retrieval, setting, value)
|
||
response = client.post("/api/retrieval", json={"query": "那拉提"})
|
||
assert response.status_code == 503 and "容量" in response.json()["detail"]
|
||
assert client.get("/api/retrieval/runs").json()["total"] == 0
|