gqt
234 lines
8.5 KiB
Python
234 lines
8.5 KiB
Python
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("*"))
|