915 lines
34 KiB
Python
Executable File
915 lines
34 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""AI 管家并发压测 + 纠错调教脚本。
|
||
|
||
目标：
|
||
1. 对 /rag/ask 或 /ai-steward/rag/ask 做并发多轮压测。
|
||
2. 覆盖“纠错追问 + AI 引擎调用链”高风险问法，检测模板化误答泄漏。
|
||
3. 可选自动提交反馈（positive / negative + correct_answer）用于线上调教数据积累。
|
||
4. 输出 JSON 报告并按阈值给出 pass/fail 退出码。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import random
|
||
import statistics
|
||
import sys
|
||
import time
|
||
from dataclasses import asdict, dataclass, field
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Iterable
|
||
|
||
import httpx
|
||
|
||
|
||
DEFAULT_TEMPLATE_LEAK_PATTERNS = [
|
||
"当前系统采用前后端分层架构",
|
||
"fastapi + sqlalchemy",
|
||
"postgresql (pgvector-ready)",
|
||
"后端模块规模",
|
||
]
|
||
|
||
ENGINE_QUERY_HINTS = [
|
||
"ai引擎",
|
||
"openclaw",
|
||
"调用链",
|
||
"模型网关",
|
||
"路由",
|
||
]
|
||
|
||
CORRECTION_HINTS = [
|
||
"不是",
|
||
"不对",
|
||
"错了",
|
||
"我说的是",
|
||
"纠正",
|
||
]
|
||
|
||
ARCHITECTURE_HINTS = [
|
||
"系统架构",
|
||
"技术架构",
|
||
"前后端",
|
||
"整体架构",
|
||
"架构",
|
||
]
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TurnCase:
|
||
question: str
|
||
thinking_mode: str = "normal"
|
||
expect_deep: bool = False
|
||
engine_sensitive: bool = False
|
||
expect_requirement_sync: bool = False
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Scenario:
|
||
name: str
|
||
turns: list[TurnCase]
|
||
|
||
|
||
@dataclass
|
||
class TurnSample:
|
||
worker_id: int
|
||
round_id: int
|
||
scenario: str
|
||
turn_index: int
|
||
question: str
|
||
requested_thinking_mode: str
|
||
expected_deep: bool
|
||
endpoint_latency_ms: int
|
||
backend_response_time_ms: int
|
||
status_code: int
|
||
ok: bool
|
||
error: str | None
|
||
session_id: int | None
|
||
qa_id: int | None
|
||
query_type: str
|
||
intent: str
|
||
returned_thinking_mode: str
|
||
confidence: float
|
||
evidence_guard_hit: bool
|
||
template_leak: bool
|
||
direct_reply_required: bool
|
||
requirement_sync_expected: bool
|
||
requirement_sync_hit: bool
|
||
answer_preview: str
|
||
feedback_sent: str | None = None
|
||
|
||
|
||
@dataclass
|
||
class Metrics:
|
||
samples: list[TurnSample] = field(default_factory=list)
|
||
session_mismatch: int = 0
|
||
feedback_positive: int = 0
|
||
feedback_negative: int = 0
|
||
feedback_failed: int = 0
|
||
|
||
def add(self, sample: TurnSample) -> None:
|
||
self.samples.append(sample)
|
||
|
||
@property
|
||
def total(self) -> int:
|
||
return len(self.samples)
|
||
|
||
@property
|
||
def ok_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.ok)
|
||
|
||
@property
|
||
def fail_count(self) -> int:
|
||
return self.total - self.ok_count
|
||
|
||
@property
|
||
def success_rate(self) -> float:
|
||
if self.total == 0:
|
||
return 0.0
|
||
return self.ok_count / self.total
|
||
|
||
@property
|
||
def template_leak_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.template_leak)
|
||
|
||
@property
|
||
def template_leak_rate(self) -> float:
|
||
if self.total == 0:
|
||
return 0.0
|
||
return self.template_leak_count / self.total
|
||
|
||
@property
|
||
def deep_expected_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.expected_deep)
|
||
|
||
@property
|
||
def deep_hit_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.expected_deep and x.returned_thinking_mode == "deep")
|
||
|
||
@property
|
||
def deep_hit_rate(self) -> float:
|
||
if self.deep_expected_count == 0:
|
||
return 1.0
|
||
return self.deep_hit_count / self.deep_expected_count
|
||
|
||
@property
|
||
def requirement_sync_expected_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.requirement_sync_expected)
|
||
|
||
@property
|
||
def requirement_sync_hit_count(self) -> int:
|
||
return sum(1 for x in self.samples if x.requirement_sync_expected and x.requirement_sync_hit)
|
||
|
||
@property
|
||
def requirement_sync_hit_rate(self) -> float:
|
||
if self.requirement_sync_expected_count == 0:
|
||
return 0.0
|
||
return self.requirement_sync_hit_count / self.requirement_sync_expected_count
|
||
|
||
@property
|
||
def evidence_guard_hits(self) -> int:
|
||
return sum(1 for x in self.samples if x.evidence_guard_hit)
|
||
|
||
@property
|
||
def endpoint_latencies(self) -> list[int]:
|
||
return [x.endpoint_latency_ms for x in self.samples if x.endpoint_latency_ms > 0]
|
||
|
||
@property
|
||
def backend_latencies(self) -> list[int]:
|
||
return [x.backend_response_time_ms for x in self.samples if x.backend_response_time_ms > 0]
|
||
|
||
|
||
def _percentile(values: Iterable[int], q: float) -> int:
|
||
data = sorted(int(v) for v in values if v >= 0)
|
||
if not data:
|
||
return 0
|
||
if len(data) == 1:
|
||
return data[0]
|
||
pos = max(0.0, min(1.0, q)) * (len(data) - 1)
|
||
lo = int(pos)
|
||
hi = min(lo + 1, len(data) - 1)
|
||
frac = pos - lo
|
||
return int(round(data[lo] * (1.0 - frac) + data[hi] * frac))
|
||
|
||
|
||
def _build_rag_scenarios() -> list[Scenario]:
|
||
return [
|
||
Scenario(
|
||
name="engine_correction_chain",
|
||
turns=[
|
||
TurnCase(
|
||
question="给我说下当前系统架构",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=False,
|
||
),
|
||
TurnCase(
|
||
question="不是啊，我说的是ai引擎驱动的方式",
|
||
thinking_mode="normal",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
),
|
||
TurnCase(
|
||
question="内置的ai管家第一层也是去调用openclaw吗？",
|
||
thinking_mode="normal",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
),
|
||
],
|
||
),
|
||
Scenario(
|
||
name="manual_deep_engine_chain",
|
||
turns=[
|
||
TurnCase(
|
||
question="请深入分析本项目AI问答引擎调用链，按入口->路由->模型网关列证据。",
|
||
thinking_mode="deep",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
),
|
||
TurnCase(
|
||
question="如果证据不足请直接说不确定，不要泛化成系统架构模板。",
|
||
thinking_mode="deep",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
),
|
||
],
|
||
),
|
||
Scenario(
|
||
name="quick_correction_probe",
|
||
turns=[
|
||
TurnCase(
|
||
question="给我说下当前系统的ai架构",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=True,
|
||
),
|
||
TurnCase(
|
||
question="不是啊。",
|
||
thinking_mode="normal",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
),
|
||
TurnCase(
|
||
question="草泥马（这是鲁棒性测试，应该礼貌回应并要求明确目标）",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=False,
|
||
),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
def _build_feishu_capture_scenarios() -> list[Scenario]:
|
||
return [
|
||
Scenario(
|
||
name="question_requirement_sync",
|
||
turns=[
|
||
TurnCase(
|
||
question="这个需求要不要支持飞书同步？",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=False,
|
||
expect_requirement_sync=True,
|
||
),
|
||
TurnCase(
|
||
question="再加一个导出按钮可以吗？",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=False,
|
||
expect_requirement_sync=True,
|
||
),
|
||
TurnCase(
|
||
question="好的，按这个需求推进。",
|
||
thinking_mode="normal",
|
||
expect_deep=False,
|
||
engine_sensitive=False,
|
||
expect_requirement_sync=False,
|
||
),
|
||
],
|
||
),
|
||
Scenario(
|
||
name="feishu_explicit_deep_chain",
|
||
turns=[
|
||
TurnCase(
|
||
question="请深入分析内置AI管家到飞书需求同步链路，按入口->路由->写库给证据。",
|
||
thinking_mode="normal",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
expect_requirement_sync=False,
|
||
),
|
||
TurnCase(
|
||
question="不是系统架构，我要你继续深挖飞书链路里具体是哪个服务落库。",
|
||
thinking_mode="normal",
|
||
expect_deep=True,
|
||
engine_sensitive=True,
|
||
expect_requirement_sync=False,
|
||
),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
def _build_scenarios(mode: str) -> list[Scenario]:
|
||
if str(mode or "").strip().lower() == "feishu_capture":
|
||
return _build_feishu_capture_scenarios()
|
||
return _build_rag_scenarios()
|
||
|
||
|
||
def _is_engine_or_correction_question(question: str) -> bool:
|
||
q = str(question or "").strip().lower()
|
||
if not q:
|
||
return False
|
||
return any(k in q for k in ENGINE_QUERY_HINTS) or any(k in q for k in CORRECTION_HINTS)
|
||
|
||
|
||
def _is_plain_architecture_question(question: str) -> bool:
|
||
q = str(question or "").strip().lower()
|
||
if not q:
|
||
return False
|
||
has_arch = any(k in q for k in ARCHITECTURE_HINTS)
|
||
has_engine = any(k in q for k in ENGINE_QUERY_HINTS)
|
||
has_correction = any(k in q for k in CORRECTION_HINTS)
|
||
if not has_arch:
|
||
return False
|
||
# “纠错 + 架构词”不算纯架构问法，比如：不是系统架构，我要AI引擎路由。
|
||
return (not has_engine) and (not has_correction)
|
||
|
||
|
||
def _contains_template_leak(answer: str) -> bool:
|
||
text = str(answer or "").strip().lower()
|
||
if not text:
|
||
return False
|
||
return any(pattern.lower() in text for pattern in DEFAULT_TEMPLATE_LEAK_PATTERNS)
|
||
|
||
|
||
def _is_requirement_like_question(question: str) -> bool:
|
||
q = str(question or "").strip().lower()
|
||
if not q:
|
||
return False
|
||
has_question = ("?" in q) or ("？" in q) or any(k in q for k in ("吗", "是否", "要不要", "能不能", "可不可以"))
|
||
if not has_question:
|
||
return False
|
||
verb_hits = ("新增", "增加", "加一个", "加个", "支持", "接入", "实现", "优化", "修复", "改进", "同步")
|
||
object_hits = ("功能", "能力", "接口", "页面", "按钮", "流程", "模块", "需求", "项目", "队列", "信号")
|
||
return any(v in q for v in verb_hits) and any(o in q for o in object_hits)
|
||
|
||
|
||
def _pick_negative_reason(question: str) -> tuple[list[str], str]:
|
||
q = str(question or "").strip()
|
||
if _is_engine_or_correction_question(q):
|
||
return (
|
||
["答非所问", "答案错误", "来源错误"],
|
||
"请基于项目AI引擎调用链回答：入口、路由、模型网关分别是什么；证据不足就明确写“不确定”。",
|
||
)
|
||
return (["答案错误"], "请结合证据回答，并避免使用通用模板。")
|
||
|
||
|
||
async def _login_token(
|
||
*,
|
||
base_url: str,
|
||
username: str,
|
||
password: str,
|
||
timeout_seconds: float,
|
||
) -> str:
|
||
async with httpx.AsyncClient(timeout=timeout_seconds, trust_env=False) as client:
|
||
resp = await client.post(
|
||
f"{base_url}/auth/login",
|
||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||
data={"username": username, "password": password},
|
||
)
|
||
resp.raise_for_status()
|
||
payload = resp.json()
|
||
token = str(payload.get("access_token") or "").strip()
|
||
if not token:
|
||
raise RuntimeError("login success but access_token missing")
|
||
return token
|
||
|
||
|
||
async def _send_feedback(
|
||
*,
|
||
client: httpx.AsyncClient,
|
||
base_url: str,
|
||
token: str,
|
||
qa_id: int,
|
||
payload: dict[str, Any],
|
||
) -> tuple[bool, int]:
|
||
try:
|
||
resp = await client.post(
|
||
f"{base_url}/rag/qa-records/{qa_id}/feedback",
|
||
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
||
json=payload,
|
||
)
|
||
return (resp.status_code < 300, int(resp.status_code))
|
||
except Exception:
|
||
return (False, 0)
|
||
|
||
|
||
async def _ask_once(
|
||
*,
|
||
client: httpx.AsyncClient,
|
||
base_url: str,
|
||
endpoint: str,
|
||
token: str,
|
||
mode: str,
|
||
question: str,
|
||
thinking_mode: str,
|
||
session_id: int | None,
|
||
project_id: int | None,
|
||
group_chat_id: str,
|
||
sender: str,
|
||
feishu_message_id: str,
|
||
) -> tuple[int, int, dict[str, Any] | None, str | None]:
|
||
mode_key = str(mode or "").strip().lower()
|
||
if mode_key == "feishu_capture":
|
||
payload = {
|
||
"group_chat_id": group_chat_id,
|
||
"message_content": question,
|
||
"sender": sender,
|
||
"feishu_message_id": feishu_message_id,
|
||
}
|
||
else:
|
||
payload = {
|
||
"question": question,
|
||
"asked_via": "pmflow",
|
||
"thinking_mode": "deep" if thinking_mode == "deep" else "normal",
|
||
}
|
||
if session_id is not None:
|
||
payload["session_id"] = int(session_id)
|
||
if project_id is not None:
|
||
payload["project_id"] = int(project_id)
|
||
|
||
started = time.perf_counter()
|
||
try:
|
||
resp = await client.post(
|
||
f"{base_url}{endpoint}",
|
||
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
||
json=payload,
|
||
)
|
||
elapsed = int((time.perf_counter() - started) * 1000)
|
||
if resp.status_code >= 400:
|
||
text = resp.text[:300]
|
||
return (resp.status_code, elapsed, None, text)
|
||
body = resp.json()
|
||
data = body.get("data") if isinstance(body, dict) and isinstance(body.get("data"), dict) else body
|
||
if not isinstance(data, dict):
|
||
return (resp.status_code, elapsed, None, "response_data_not_dict")
|
||
return (resp.status_code, elapsed, data, None)
|
||
except Exception as exc:
|
||
elapsed = int((time.perf_counter() - started) * 1000)
|
||
return (0, elapsed, None, str(exc))
|
||
|
||
|
||
async def _run_worker(
|
||
*,
|
||
worker_id: int,
|
||
rounds: int,
|
||
base_url: str,
|
||
endpoint: str,
|
||
token: str,
|
||
mode: str,
|
||
timeout_seconds: float,
|
||
project_id: int | None,
|
||
scenarios: list[Scenario],
|
||
rng_seed: int,
|
||
auto_feedback: bool,
|
||
sleep_ms: int,
|
||
chat_id_prefix: str,
|
||
sender_prefix: str,
|
||
) -> Metrics:
|
||
mode_key = str(mode or "").strip().lower()
|
||
rng = random.Random(rng_seed + worker_id)
|
||
metrics = Metrics()
|
||
async with httpx.AsyncClient(timeout=timeout_seconds, trust_env=False) as client:
|
||
for round_id in range(1, rounds + 1):
|
||
scenario = scenarios[rng.randrange(0, len(scenarios))]
|
||
active_session_id: int | None = None
|
||
previous_session_id: int | None = None
|
||
group_chat_id = f"{chat_id_prefix}-{worker_id}-{round_id}"
|
||
sender = f"{sender_prefix}{worker_id:02d}"
|
||
|
||
for turn_idx, turn in enumerate(scenario.turns, start=1):
|
||
feishu_message_id = f"pc-{worker_id}-{round_id}-{turn_idx}-{rng.randrange(100000, 999999)}"
|
||
status_code, endpoint_latency, data, err = await _ask_once(
|
||
client=client,
|
||
base_url=base_url,
|
||
endpoint=endpoint,
|
||
token=token,
|
||
mode=mode_key,
|
||
question=turn.question,
|
||
thinking_mode=turn.thinking_mode,
|
||
session_id=active_session_id,
|
||
project_id=project_id,
|
||
group_chat_id=group_chat_id,
|
||
sender=sender,
|
||
feishu_message_id=feishu_message_id,
|
||
)
|
||
|
||
payload = data if isinstance(data, dict) else {}
|
||
ok = bool(status_code and status_code < 300 and isinstance(payload, dict))
|
||
if mode_key == "feishu_capture":
|
||
if str(payload.get("status") or "").strip().lower() in {"error", "failed"}:
|
||
ok = False
|
||
rag_payload = payload.get("rag_result") if isinstance(payload.get("rag_result"), dict) else {}
|
||
answer = str(payload.get("direct_reply_text") or rag_payload.get("answer") or "")
|
||
returned_mode = str(rag_payload.get("thinking_mode") or "normal").strip().lower()
|
||
query_type = str(rag_payload.get("query_type") or "")
|
||
qa_id_raw = rag_payload.get("qa_id")
|
||
sess_raw = rag_payload.get("session_id")
|
||
confidence_raw = rag_payload.get("confidence")
|
||
backend_response_time_raw = rag_payload.get("response_time_ms")
|
||
sources = rag_payload.get("sources") if isinstance(rag_payload.get("sources"), list) else []
|
||
intent = str(payload.get("intent") or "").strip().lower()
|
||
direct_reply_required = bool(payload.get("direct_reply_required"))
|
||
requirement_sync_expected = bool(
|
||
turn.expect_requirement_sync or _is_requirement_like_question(turn.question)
|
||
)
|
||
has_sync_artifact = bool(payload.get("results")) or (payload.get("signal_id") is not None)
|
||
requirement_sync_hit = bool(
|
||
requirement_sync_expected
|
||
and intent == "question_requirement"
|
||
and str(payload.get("status") or "").strip().lower() == "captured"
|
||
and has_sync_artifact
|
||
)
|
||
else:
|
||
answer = str(payload.get("answer") or "")
|
||
returned_mode = str(payload.get("thinking_mode") or "normal").strip().lower()
|
||
query_type = str(payload.get("query_type") or "")
|
||
qa_id_raw = payload.get("qa_id")
|
||
sess_raw = payload.get("session_id")
|
||
confidence_raw = payload.get("confidence")
|
||
backend_response_time_raw = payload.get("response_time_ms")
|
||
sources = payload.get("sources") if isinstance(payload.get("sources"), list) else []
|
||
intent = ""
|
||
direct_reply_required = False
|
||
requirement_sync_expected = bool(turn.expect_requirement_sync)
|
||
requirement_sync_hit = False
|
||
|
||
try:
|
||
qa_id = int(qa_id_raw) if qa_id_raw is not None else None
|
||
except Exception:
|
||
qa_id = None
|
||
try:
|
||
current_session_id = int(sess_raw) if sess_raw is not None else None
|
||
except Exception:
|
||
current_session_id = None
|
||
try:
|
||
confidence = float(confidence_raw) if confidence_raw is not None else 0.0
|
||
except Exception:
|
||
confidence = 0.0
|
||
try:
|
||
backend_response_time_ms = int(backend_response_time_raw) if backend_response_time_raw is not None else 0
|
||
except Exception:
|
||
backend_response_time_ms = 0
|
||
|
||
evidence_guard_hit = False
|
||
for src in sources:
|
||
if not isinstance(src, dict):
|
||
continue
|
||
if str(src.get("label") or "").strip().lower() == "evidence guard":
|
||
evidence_guard_hit = True
|
||
break
|
||
|
||
template_leak = False
|
||
if (turn.engine_sensitive or _is_engine_or_correction_question(turn.question)) and (
|
||
not _is_plain_architecture_question(turn.question)
|
||
):
|
||
template_leak = _contains_template_leak(answer)
|
||
|
||
if (
|
||
mode_key != "feishu_capture"
|
||
and previous_session_id is not None
|
||
and current_session_id is not None
|
||
and previous_session_id != current_session_id
|
||
):
|
||
metrics.session_mismatch += 1
|
||
if current_session_id is not None:
|
||
active_session_id = current_session_id
|
||
previous_session_id = current_session_id
|
||
|
||
sample = TurnSample(
|
||
worker_id=worker_id,
|
||
round_id=round_id,
|
||
scenario=scenario.name,
|
||
turn_index=turn_idx,
|
||
question=turn.question,
|
||
requested_thinking_mode=turn.thinking_mode,
|
||
expected_deep=turn.expect_deep,
|
||
endpoint_latency_ms=endpoint_latency,
|
||
backend_response_time_ms=backend_response_time_ms,
|
||
status_code=status_code,
|
||
ok=ok,
|
||
error=err,
|
||
session_id=current_session_id,
|
||
qa_id=qa_id,
|
||
query_type=query_type,
|
||
intent=intent,
|
||
returned_thinking_mode=returned_mode if returned_mode in {"normal", "deep"} else "normal",
|
||
confidence=round(max(0.0, min(1.0, confidence)), 4),
|
||
evidence_guard_hit=evidence_guard_hit,
|
||
template_leak=template_leak,
|
||
direct_reply_required=direct_reply_required,
|
||
requirement_sync_expected=requirement_sync_expected,
|
||
requirement_sync_hit=requirement_sync_hit,
|
||
answer_preview=answer.replace("\n", " ")[:180],
|
||
)
|
||
|
||
if auto_feedback and ok and qa_id is not None:
|
||
if template_leak:
|
||
error_types, correct_answer = _pick_negative_reason(turn.question)
|
||
fb_payload = {
|
||
"rating": "negative",
|
||
"error_types": error_types,
|
||
"correct_answer": correct_answer,
|
||
}
|
||
fb_ok, _ = await _send_feedback(
|
||
client=client,
|
||
base_url=base_url,
|
||
token=token,
|
||
qa_id=qa_id,
|
||
payload=fb_payload,
|
||
)
|
||
if fb_ok:
|
||
metrics.feedback_negative += 1
|
||
sample.feedback_sent = "negative"
|
||
else:
|
||
metrics.feedback_failed += 1
|
||
else:
|
||
fb_payload = {
|
||
"rating": "positive",
|
||
"helpfulness": "fully",
|
||
}
|
||
fb_ok, _ = await _send_feedback(
|
||
client=client,
|
||
base_url=base_url,
|
||
token=token,
|
||
qa_id=qa_id,
|
||
payload=fb_payload,
|
||
)
|
||
if fb_ok:
|
||
metrics.feedback_positive += 1
|
||
sample.feedback_sent = "positive"
|
||
else:
|
||
metrics.feedback_failed += 1
|
||
|
||
metrics.add(sample)
|
||
if sleep_ms > 0:
|
||
await asyncio.sleep(max(0.0, sleep_ms / 1000.0))
|
||
return metrics
|
||
|
||
|
||
def _merge_metrics(items: list[Metrics]) -> Metrics:
|
||
merged = Metrics()
|
||
for m in items:
|
||
merged.samples.extend(m.samples)
|
||
merged.session_mismatch += m.session_mismatch
|
||
merged.feedback_positive += m.feedback_positive
|
||
merged.feedback_negative += m.feedback_negative
|
||
merged.feedback_failed += m.feedback_failed
|
||
return merged
|
||
|
||
|
||
def _build_report(
|
||
*,
|
||
started_at: str,
|
||
finished_at: str,
|
||
args: argparse.Namespace,
|
||
metrics: Metrics,
|
||
) -> dict[str, Any]:
|
||
endpoint_lat = metrics.endpoint_latencies
|
||
backend_lat = metrics.backend_latencies
|
||
errors = [x for x in metrics.samples if not x.ok]
|
||
template_hits = [x for x in metrics.samples if x.template_leak]
|
||
deep_miss = [x for x in metrics.samples if x.expected_deep and x.returned_thinking_mode != "deep"]
|
||
|
||
summary = {
|
||
"started_at": started_at,
|
||
"finished_at": finished_at,
|
||
"mode": args.mode,
|
||
"base_url": args.base_url,
|
||
"endpoint": args.endpoint,
|
||
"workers": args.workers,
|
||
"rounds": args.rounds,
|
||
"total_turns": metrics.total,
|
||
"ok_turns": metrics.ok_count,
|
||
"failed_turns": metrics.fail_count,
|
||
"success_rate": round(metrics.success_rate, 4),
|
||
"template_leak_count": metrics.template_leak_count,
|
||
"template_leak_rate": round(metrics.template_leak_rate, 4),
|
||
"deep_expected_count": metrics.deep_expected_count,
|
||
"deep_hit_count": metrics.deep_hit_count,
|
||
"deep_hit_rate": round(metrics.deep_hit_rate, 4),
|
||
"question_requirement_expected_count": metrics.requirement_sync_expected_count,
|
||
"question_requirement_hit_count": metrics.requirement_sync_hit_count,
|
||
"question_requirement_hit_rate": round(metrics.requirement_sync_hit_rate, 4),
|
||
"evidence_guard_hits": metrics.evidence_guard_hits,
|
||
"session_mismatch": metrics.session_mismatch,
|
||
"feedback_positive": metrics.feedback_positive,
|
||
"feedback_negative": metrics.feedback_negative,
|
||
"feedback_failed": metrics.feedback_failed,
|
||
"endpoint_latency_ms": {
|
||
"p50": _percentile(endpoint_lat, 0.50),
|
||
"p95": _percentile(endpoint_lat, 0.95),
|
||
"p99": _percentile(endpoint_lat, 0.99),
|
||
"avg": round(statistics.mean(endpoint_lat), 2) if endpoint_lat else 0.0,
|
||
},
|
||
"backend_response_time_ms": {
|
||
"p50": _percentile(backend_lat, 0.50),
|
||
"p95": _percentile(backend_lat, 0.95),
|
||
"p99": _percentile(backend_lat, 0.99),
|
||
"avg": round(statistics.mean(backend_lat), 2) if backend_lat else 0.0,
|
||
},
|
||
}
|
||
|
||
return {
|
||
"summary": summary,
|
||
"thresholds": {
|
||
"min_success_rate": args.min_success_rate,
|
||
"max_template_leak_rate": args.max_template_leak_rate,
|
||
"min_deep_hit_rate": args.min_deep_hit_rate,
|
||
"min_requirement_sync_hit_rate": args.min_requirement_sync_hit_rate,
|
||
"max_session_mismatch": args.max_session_mismatch,
|
||
},
|
||
"samples_head": [asdict(x) for x in metrics.samples[:30]],
|
||
"error_samples": [asdict(x) for x in errors[:20]],
|
||
"template_leak_samples": [asdict(x) for x in template_hits[:20]],
|
||
"deep_miss_samples": [asdict(x) for x in deep_miss[:20]],
|
||
}
|
||
|
||
|
||
def _print_console_summary(report: dict[str, Any]) -> None:
|
||
s = report.get("summary") or {}
|
||
print("==== AI Steward Pressure Coach Summary ====")
|
||
print(f"mode={s.get('mode') or ''} base_url={s.get('base_url') or ''} endpoint={s.get('endpoint') or ''}")
|
||
print(
|
||
f"turns={s.get('total_turns')} ok={s.get('ok_turns')} fail={s.get('failed_turns')} "
|
||
f"success_rate={s.get('success_rate')}"
|
||
)
|
||
print(
|
||
f"template_leak={s.get('template_leak_count')}({s.get('template_leak_rate')}) "
|
||
f"deep_hit={s.get('deep_hit_count')}/{s.get('deep_expected_count')}({s.get('deep_hit_rate')}) "
|
||
f"question_requirement_hit={s.get('question_requirement_hit_count')}/{s.get('question_requirement_expected_count')}({s.get('question_requirement_hit_rate')}) "
|
||
f"session_mismatch={s.get('session_mismatch')}"
|
||
)
|
||
el = s.get("endpoint_latency_ms") or {}
|
||
bl = s.get("backend_response_time_ms") or {}
|
||
print(
|
||
f"latency endpoint(p50/p95/p99/avg)={el.get('p50')}/{el.get('p95')}/{el.get('p99')}/{el.get('avg')} ms"
|
||
)
|
||
print(
|
||
f"latency backend(p50/p95/p99/avg)={bl.get('p50')}/{bl.get('p95')}/{bl.get('p99')}/{bl.get('avg')} ms"
|
||
)
|
||
print(
|
||
f"feedback +={s.get('feedback_positive')} -={s.get('feedback_negative')} failed={s.get('feedback_failed')}"
|
||
)
|
||
|
||
|
||
def _parse_args() -> argparse.Namespace:
|
||
now_tag = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
parser = argparse.ArgumentParser(description="AI 管家并发压测 + 调教脚本")
|
||
parser.add_argument("--mode", choices=["rag", "feishu_capture"], default=os.getenv("STRESS_MODE", "rag"))
|
||
parser.add_argument("--base-url", default=os.getenv("STRESS_BASE_URL", "http://127.0.0.1:8000"))
|
||
parser.add_argument("--endpoint", default=os.getenv("STRESS_ENDPOINT", ""))
|
||
parser.add_argument("--user", default=os.getenv("STRESS_USER", "boss"))
|
||
parser.add_argument("--password", default=os.getenv("STRESS_PASSWORD", "boss123"))
|
||
parser.add_argument("--workers", type=int, default=int(os.getenv("STRESS_WORKERS", "8")))
|
||
parser.add_argument("--rounds", type=int, default=int(os.getenv("STRESS_ROUNDS", "6")))
|
||
parser.add_argument("--project-id", type=int, default=None)
|
||
parser.add_argument("--timeout-seconds", type=float, default=float(os.getenv("STRESS_TIMEOUT", "60")))
|
||
parser.add_argument("--sleep-ms", type=int, default=int(os.getenv("STRESS_SLEEP_MS", "80")))
|
||
parser.add_argument("--seed", type=int, default=int(os.getenv("STRESS_SEED", "20260306")))
|
||
parser.add_argument("--auto-feedback", action="store_true", default=False)
|
||
parser.add_argument("--report-file", default=f"logs/ai_steward_pressure_{now_tag}.json")
|
||
parser.add_argument("--min-success-rate", type=float, default=0.95)
|
||
parser.add_argument("--max-template-leak-rate", type=float, default=0.08)
|
||
parser.add_argument("--min-deep-hit-rate", type=float, default=0.90)
|
||
parser.add_argument("--min-requirement-sync-hit-rate", type=float, default=0.90)
|
||
parser.add_argument("--max-session-mismatch", type=int, default=0)
|
||
parser.add_argument("--chat-id-prefix", default=os.getenv("STRESS_CHAT_PREFIX", "oc_pressure_chat"))
|
||
parser.add_argument("--sender-prefix", default=os.getenv("STRESS_SENDER_PREFIX", "压测用户"))
|
||
return parser.parse_args()
|
||
|
||
|
||
def _normalize_base_url(raw: str) -> str:
|
||
base = str(raw or "").strip().rstrip("/")
|
||
if not base:
|
||
return "http://127.0.0.1:8000"
|
||
if not (base.startswith("http://") or base.startswith("https://")):
|
||
return f"http://{base}"
|
||
return base
|
||
|
||
|
||
def _normalize_endpoint(raw: str) -> str:
|
||
ep = str(raw or "").strip()
|
||
if not ep:
|
||
return "/rag/ask"
|
||
if not ep.startswith("/"):
|
||
ep = f"/{ep}"
|
||
return ep
|
||
|
||
|
||
async def _amain() -> int:
|
||
args = _parse_args()
|
||
args.mode = str(args.mode or "rag").strip().lower()
|
||
args.base_url = _normalize_base_url(args.base_url)
|
||
if not str(args.endpoint or "").strip():
|
||
args.endpoint = "/ai-steward/signal/capture-feishu" if args.mode == "feishu_capture" else "/rag/ask"
|
||
args.endpoint = _normalize_endpoint(args.endpoint)
|
||
args.workers = max(1, min(args.workers, 256))
|
||
args.rounds = max(1, min(args.rounds, 10000))
|
||
args.sleep_ms = max(0, min(args.sleep_ms, 10000))
|
||
args.timeout_seconds = max(5.0, min(args.timeout_seconds, 300.0))
|
||
|
||
started_at = datetime.now(timezone.utc).isoformat()
|
||
scenarios = _build_scenarios(args.mode)
|
||
token = await _login_token(
|
||
base_url=args.base_url,
|
||
username=args.user,
|
||
password=args.password,
|
||
timeout_seconds=args.timeout_seconds,
|
||
)
|
||
|
||
tasks = [
|
||
_run_worker(
|
||
worker_id=i + 1,
|
||
rounds=args.rounds,
|
||
base_url=args.base_url,
|
||
endpoint=args.endpoint,
|
||
token=token,
|
||
mode=args.mode,
|
||
timeout_seconds=args.timeout_seconds,
|
||
project_id=args.project_id,
|
||
scenarios=scenarios,
|
||
rng_seed=args.seed,
|
||
auto_feedback=bool(args.auto_feedback),
|
||
sleep_ms=args.sleep_ms,
|
||
chat_id_prefix=args.chat_id_prefix,
|
||
sender_prefix=args.sender_prefix,
|
||
)
|
||
for i in range(args.workers)
|
||
]
|
||
results = await asyncio.gather(*tasks)
|
||
metrics = _merge_metrics(results)
|
||
finished_at = datetime.now(timezone.utc).isoformat()
|
||
|
||
report = _build_report(
|
||
started_at=started_at,
|
||
finished_at=finished_at,
|
||
args=args,
|
||
metrics=metrics,
|
||
)
|
||
_print_console_summary(report)
|
||
|
||
report_path = Path(args.report_file)
|
||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(f"report_file={report_path}")
|
||
|
||
summary = report.get("summary") or {}
|
||
success_rate = float(summary.get("success_rate") or 0.0)
|
||
template_leak_rate = float(summary.get("template_leak_rate") or 0.0)
|
||
deep_hit_rate = float(summary.get("deep_hit_rate") or 0.0)
|
||
requirement_sync_expected_count = int(summary.get("question_requirement_expected_count") or 0)
|
||
requirement_sync_hit_rate = float(summary.get("question_requirement_hit_rate") or 0.0)
|
||
session_mismatch = int(summary.get("session_mismatch") or 0)
|
||
|
||
failed_checks: list[str] = []
|
||
if success_rate < float(args.min_success_rate):
|
||
failed_checks.append(f"success_rate<{args.min_success_rate}")
|
||
if template_leak_rate > float(args.max_template_leak_rate):
|
||
failed_checks.append(f"template_leak_rate>{args.max_template_leak_rate}")
|
||
if deep_hit_rate < float(args.min_deep_hit_rate):
|
||
failed_checks.append(f"deep_hit_rate<{args.min_deep_hit_rate}")
|
||
if (
|
||
args.mode == "feishu_capture"
|
||
and requirement_sync_expected_count > 0
|
||
and requirement_sync_hit_rate < float(args.min_requirement_sync_hit_rate)
|
||
):
|
||
failed_checks.append(f"question_requirement_hit_rate<{args.min_requirement_sync_hit_rate}")
|
||
if session_mismatch > int(args.max_session_mismatch):
|
||
failed_checks.append(f"session_mismatch>{args.max_session_mismatch}")
|
||
|
||
if failed_checks:
|
||
print(f"RESULT=FAIL checks={','.join(failed_checks)}")
|
||
return 2
|
||
print("RESULT=PASS")
|
||
return 0
|
||
|
||
|
||
def main() -> int:
|
||
try:
|
||
return asyncio.run(_amain())
|
||
except KeyboardInterrupt:
|
||
print("interrupted")
|
||
return 130
|
||
except Exception as exc:
|
||
print(f"fatal_error={exc}")
|
||
return 2
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|