test: verify real kbqa workflow
This commit is contained in:
gqt
@@ -0,0 +1,233 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pymilvus import MilvusClient
|
||||
|
||||
|
||||
pytestmark = pytest.mark.live
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
FIXTURES = ROOT / "backend" / "tests" / "fixtures"
|
||||
FIXTURE_PATHS = tuple(
|
||||
FIXTURES / name for name in ("knowledge.md", "knowledge.txt", "knowledge.pdf")
|
||||
)
|
||||
|
||||
|
||||
def _port() -> int:
|
||||
with socket.socket() as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return int(sock.getsockname()[1])
|
||||
|
||||
|
||||
def _start_server(port: int, env: dict[str, str]) -> subprocess.Popen[str]:
|
||||
process = subprocess.Popen(
|
||||
[
|
||||
str(ROOT / "backend" / ".venv" / "bin" / "python"),
|
||||
"-m",
|
||||
"uvicorn",
|
||||
"kbqa.main:app",
|
||||
"--app-dir",
|
||||
"backend/src",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
],
|
||||
cwd=ROOT,
|
||||
env=env,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
for _ in range(120):
|
||||
if process.poll() is not None:
|
||||
raise RuntimeError("live Uvicorn exited during startup")
|
||||
try:
|
||||
if httpx.get(f"{base_url}/health", timeout=1).status_code == 200:
|
||||
return process
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
time.sleep(0.25)
|
||||
process.terminate()
|
||||
raise TimeoutError("live Uvicorn did not become healthy")
|
||||
|
||||
|
||||
def _stop_server(process: subprocess.Popen[str]) -> None:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=15)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=5)
|
||||
|
||||
|
||||
def _stream(client: httpx.Client, session_id: str, query: str) -> list[tuple[str, dict]]:
|
||||
events: list[tuple[str, dict]] = []
|
||||
with client.stream(
|
||||
"POST", "/api/v1/chat/stream", json={"session_id": session_id, "query": query}, timeout=180
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
event_name = ""
|
||||
for line in response.iter_lines():
|
||||
if line.startswith("event: "):
|
||||
event_name = line.removeprefix("event: ")
|
||||
elif line.startswith("data: "):
|
||||
events.append((event_name, json.loads(line.removeprefix("data: "))))
|
||||
return events
|
||||
|
||||
|
||||
def _answer(events: list[tuple[str, dict]]) -> str:
|
||||
return "".join(data["content"] for name, data in events if name == "token")
|
||||
|
||||
|
||||
def _assert_flow(events: list[tuple[str, dict]]) -> None:
|
||||
stages = [
|
||||
f"status:{data['phase']}" if name == "status" else name
|
||||
for name, data in events
|
||||
if name != "ping"
|
||||
]
|
||||
expected = ["status:researching", "sources", "status:answering", "token", "done"]
|
||||
positions = [stages.index(stage) for stage in expected]
|
||||
assert positions == sorted(positions), stages
|
||||
|
||||
|
||||
def test_real_model_pipeline_survives_restart_and_cleans_up(tmp_path) -> None:
|
||||
if os.environ.get("KBQA_RUN_LIVE_TESTS") != "1":
|
||||
pytest.skip("set KBQA_RUN_LIVE_TESTS=1 to call real models")
|
||||
|
||||
database = tmp_path / "kbqa.sqlite3"
|
||||
milvus = tmp_path / "milvus.db"
|
||||
raw = tmp_path / "raw"
|
||||
env = os.environ.copy()
|
||||
env.update(
|
||||
{
|
||||
"KBQA_DATABASE_URL": f"sqlite+aiosqlite:///{database}",
|
||||
"KBQA_MILVUS_URI": str(milvus),
|
||||
"KBQA_RAW_DATA_DIR": str(raw),
|
||||
"KBQA_RUN_LIVE_TESTS": "1",
|
||||
}
|
||||
)
|
||||
subprocess.run(
|
||||
[
|
||||
str(ROOT / "backend" / ".venv" / "bin" / "alembic"),
|
||||
"-c",
|
||||
"backend/alembic.ini",
|
||||
"upgrade",
|
||||
"head",
|
||||
],
|
||||
cwd=ROOT,
|
||||
env=env,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
port = _port()
|
||||
process = _start_server(port, env)
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
document_ids: list[str] = []
|
||||
try:
|
||||
with httpx.Client(base_url=base_url, timeout=30) as client:
|
||||
for path in FIXTURE_PATHS:
|
||||
with path.open("rb") as file_handle:
|
||||
response = client.post(
|
||||
"/api/v1/documents", files={"file": (path.name, file_handle)}
|
||||
)
|
||||
response.raise_for_status()
|
||||
document_ids.append(response.json()["id"])
|
||||
|
||||
deadline = time.monotonic() + 120
|
||||
while time.monotonic() < deadline:
|
||||
documents = client.get("/api/v1/documents", params={"limit": 100}).json()["items"]
|
||||
if all(item["status"] == "indexed" for item in documents):
|
||||
break
|
||||
assert not any(item["status"] == "failed" for item in documents), documents
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
pytest.fail("documents did not finish indexing")
|
||||
assert len(documents) == 3
|
||||
assert all(item["chunk_count"] > 0 for item in documents)
|
||||
|
||||
session = client.post("/api/v1/sessions", json={"title": "真实模型测试"}).json()
|
||||
session_id = session["id"]
|
||||
first = _stream(client, session_id, "星槐项目的内部代号和负责人分别是什么？")
|
||||
_assert_flow(first)
|
||||
first_answer = _answer(first)
|
||||
assert "ORCHID-731" in first_answer
|
||||
assert "林岚" in first_answer
|
||||
first_sources = [
|
||||
source for name, data in first if name == "sources" for source in data["sources"]
|
||||
]
|
||||
assert first_sources
|
||||
assert {source["filename"] for source in first_sources} <= {
|
||||
path.name for path in FIXTURE_PATHS
|
||||
}
|
||||
chunks_by_document: dict[str, set[str]] = {}
|
||||
for source in first_sources:
|
||||
document_id = source["document_id"]
|
||||
if document_id not in chunks_by_document:
|
||||
page = client.get(
|
||||
f"/api/v1/documents/{document_id}/chunks", params={"limit": 100}
|
||||
).json()
|
||||
chunks_by_document[document_id] = {item["id"] for item in page["items"]}
|
||||
assert source["chunk_id"] in chunks_by_document[document_id]
|
||||
|
||||
follow_up = _stream(client, session_id, "她负责的项目首次评审日期是什么？")
|
||||
_assert_flow(follow_up)
|
||||
follow_up_answer = _answer(follow_up)
|
||||
assert "2026" in follow_up_answer and "18" in follow_up_answer
|
||||
absent = _stream(client, session_id, "星槐项目部署在哪个月球基地？")
|
||||
_assert_flow(absent)
|
||||
assert any(marker in _answer(absent) for marker in ("不足", "没有", "未检索"))
|
||||
history = client.get(
|
||||
f"/api/v1/sessions/{session_id}/messages", params={"limit": 100}
|
||||
).json()
|
||||
assert history["total"] == 6
|
||||
assert all(item["status"] == "completed" for item in history["items"])
|
||||
assert history["items"][1]["sources"]
|
||||
finally:
|
||||
_stop_server(process)
|
||||
|
||||
process = _start_server(port, env)
|
||||
try:
|
||||
with httpx.Client(base_url=base_url, timeout=30) as client:
|
||||
assert client.get("/api/v1/documents").json()["total"] == 3
|
||||
history = client.get(
|
||||
f"/api/v1/sessions/{session_id}/messages", params={"limit": 100}
|
||||
).json()
|
||||
assert history["total"] == 6
|
||||
after_restart = _stream(client, session_id, "青屿实验的校准周期是多久？")
|
||||
_assert_flow(after_restart)
|
||||
assert "14" in _answer(after_restart)
|
||||
for document_id in document_ids:
|
||||
assert client.delete(f"/api/v1/documents/{document_id}").status_code == 204
|
||||
assert client.delete(f"/api/v1/sessions/{session_id}").status_code == 204
|
||||
assert client.get("/api/v1/documents").json()["total"] == 0
|
||||
assert client.get(f"/api/v1/sessions/{session_id}/messages").status_code == 404
|
||||
finally:
|
||||
_stop_server(process)
|
||||
|
||||
with sqlite3.connect(database) as connection:
|
||||
for table in (
|
||||
"documents",
|
||||
"index_jobs",
|
||||
"chunks",
|
||||
"chat_sessions",
|
||||
"messages",
|
||||
"message_sources",
|
||||
):
|
||||
assert connection.execute(f"SELECT count(*) FROM {table}").fetchone()[0] == 0
|
||||
vector_client = MilvusClient(str(milvus))
|
||||
try:
|
||||
vector_client.load_collection("kbqa_chunks")
|
||||
assert int(vector_client.get_collection_stats("kbqa_chunks")["row_count"]) == 0
|
||||
finally:
|
||||
vector_client.close()
|
||||
assert not list(raw.glob("*"))
|
||||
Reference in New Issue
Block a user