PMFlow/scripts/rag_followup_smoke.py

118 lines
3.9 KiB
Python
Executable File

#!/usr/bin/env python3
"""RAG 追问链烟测:防止“分别是什么”回退成 count。"""
from __future__ import annotations
import os
import sys
import httpx
def main() -> int:
default_base = f"http://localhost:{os.getenv('BACKEND_PORT', '8000')}"
base = os.getenv("SMOKE_BASE_URL", default_base).rstrip("/")
username = os.getenv("SMOKE_USER", "boss")
password = os.getenv("SMOKE_PASSWORD", "boss123")
with httpx.Client(timeout=30.0) as client:
login = client.post(
f"{base}/auth/login",
headers={"Content-Type": "application/x-www-form-urlencoded"},
data={"username": username, "password": password},
)
if login.status_code != 200:
print(f"[FAIL] login status={login.status_code} body={login.text[:240]}")
return 2
token = login.json().get("access_token")
if not token:
print("[FAIL] login token missing")
return 2
headers = {"Authorization": f"Bearer {token}"}
q0 = client.post(
f"{base}/rag/ask",
headers=headers,
json={"question": "健身食谱项目里面有多少个信号源", "asked_via": "pmflow"},
)
if q0.status_code != 200:
print(f"[FAIL] q0 status={q0.status_code} body={q0.text[:240]}")
return 2
q1 = client.post(
f"{base}/rag/ask",
headers=headers,
json={"question": "目前系统里面有多少个项目", "asked_via": "pmflow"},
)
if q1.status_code != 200:
print(f"[FAIL] q1 status={q1.status_code} body={q1.text[:240]}")
return 2
q2 = client.post(
f"{base}/rag/ask",
headers=headers,
json={"question": "分别是什么", "asked_via": "pmflow"},
)
if q2.status_code != 200:
print(f"[FAIL] q2 status={q2.status_code} body={q2.text[:240]}")
return 2
d0 = (q0.json() or {}).get("data") or {}
d1 = (q1.json() or {}).get("data") or {}
d2 = (q2.json() or {}).get("data") or {}
qt0 = str(d0.get("query_type") or "")
qt1 = str(d1.get("query_type") or "")
qt2 = str(d2.get("query_type") or "")
a0 = str(d0.get("answer") or "")
a2 = str(d2.get("answer") or "")
q3 = client.post(
f"{base}/rag/ask",
headers=headers,
json={"question": "健身食谱项目里面有多少个信号源", "asked_via": "pmflow"},
)
if q3.status_code != 200:
print(f"[FAIL] q3 status={q3.status_code} body={q3.text[:240]}")
return 2
q4 = client.post(
f"{base}/rag/ask",
headers=headers,
json={"question": "分别什么信号源", "asked_via": "pmflow"},
)
if q4.status_code != 200:
print(f"[FAIL] q4 status={q4.status_code} body={q4.text[:240]}")
return 2
d3 = (q3.json() or {}).get("data") or {}
d4 = (q4.json() or {}).get("data") or {}
qt3 = str(d3.get("query_type") or "")
qt4 = str(d4.get("query_type") or "")
a4 = str(d4.get("answer") or "")
ok = (
(qt0 == "count")
and ("信号" in a0)
and (qt1 == "count")
and ("项目总数" in str(d1.get("answer") or ""))
and (qt2 == "list")
and ("项目列表" in a2 or "#" in a2)
and (qt3 == "count")
and (qt4 == "list")
and ("信号清单" in a4 or "信号" in a4)
)
if not ok:
print(
f"[FAIL] qt0={qt0} qt1={qt1} qt2={qt2} qt3={qt3} qt4={qt4} "
f"a0={a0[:120]} a1={str(d1.get('answer') or '')[:120]} a2={a2[:120]} a4={a4[:120]}"
)
return 2
print("[OK] followup smoke passed")
return 0
if __name__ == "__main__":
raise SystemExit(main())