Files
border-red-post-nalati/backend/tests/test_retrieval.py
T

213 lines
8.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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