mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 02:21:58 +00:00
合入期 0 核查工具箱：SQL 语义（规范化、方面、差异说明）、只读核查 agent、核查每日花费账、核查设置组、数据源查询工具的强制只读
This commit is contained in:
ThearchyHelios
+172
-7
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
@@ -14,9 +15,11 @@ from app.core.config import settings as app_settings
|
||||
from app.core.crypto import encrypt, mask
|
||||
from app.db.base import get_session
|
||||
from app.db.models import Provider, Setting
|
||||
from app.engine.judge import JUDGE_DEFAULTS, SETTING_GROUP as JUDGE_GROUP, SPEND_KEY as JUDGE_SPEND
|
||||
from app.engine.judge import daily_spend, judge_settings
|
||||
from app.engine.judge import JUDGE_DEFAULTS, MAX_TOKENS as JUDGE_MAX_TOKENS
|
||||
from app.engine.judge import SETTING_GROUP as JUDGE_GROUP, SPEND_KEY as JUDGE_SPEND
|
||||
from app.engine.judge import daily_spend, judge_model_spec, judge_settings
|
||||
from app.engine.judge import sanitize_settings as judge_values, settings_problem as judge_problem
|
||||
from app.engine.spend import SPEND_KEY as VERIFY_SPEND, daily_spend as verify_daily_spend
|
||||
from app.providers import catalog
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, _resolve_key, build_chat_model
|
||||
|
||||
@@ -412,6 +415,153 @@ async def probe_provider_models(
|
||||
return {"ok": True, "models": models, "url": url}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 核查：双盲复算、反方审稿、审批质询共用的一组设置
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
#: 设置里 verify 这一组。三项都是让一个只读 agent 查库核对（engine/verify_agent.py），共用核查模型和每日金额上限；
|
||||
#: 各自的开关和上限分开放。上限写 null 表示不限（和裁判一组的约定一样）
|
||||
VERIFY_GROUP = "verify"
|
||||
VERIFY_DEFAULTS: dict[str, Any] = {
|
||||
"provider": None, # 核查模型：空着沿用证据裁判模型的取值链（verify_model_spec）
|
||||
"model": None,
|
||||
# 双盲复算：tolerance 是允许的相对误差，两边的值差在这个比例以内算一致
|
||||
"dual_check": {"enabled": True, "tolerance": 1e-6, "max_queries": 6, "timeout_s": 120},
|
||||
# 反方审稿：max_claims 是每份报告最多挑几句关键结论来审，查询次数、时长按每句算
|
||||
"challenge": {"enabled": True, "max_claims": 5, "max_queries": 4, "timeout_s": 90},
|
||||
# 审批质询：查询次数、时长按每次提问算
|
||||
"inquiry": {"enabled": True, "max_queries": 6, "timeout_s": 120},
|
||||
"daily_max_usd": 2.0, # 全局每日核查金额上限（按本地日期），和裁判的每日上限分开记账（engine/spend.py）
|
||||
}
|
||||
#: 三项各自那一组的键
|
||||
_VERIFY_PARTS = ("dual_check", "challenge", "inquiry")
|
||||
_VERIFY_COUNTS = frozenset({"max_queries", "max_claims"})
|
||||
|
||||
#: 设置页「核查」一节各项的叫法：报错时写它，不写键名。嵌套的项写成「组.项」
|
||||
VERIFY_SETTING_LABEL = {
|
||||
"provider": "核查模型 · 接入", "model": "核查模型 · 模型",
|
||||
"daily_max_usd": "核查每日金额上限（美元）",
|
||||
"dual_check": "双盲复算", "dual_check.enabled": "双盲复算 · 启用",
|
||||
"dual_check.tolerance": "双盲复算 · 允许的相对误差", "dual_check.max_queries": "双盲复算 · 每项最多查询次数",
|
||||
"dual_check.timeout_s": "双盲复算 · 每项时长上限（秒）",
|
||||
"challenge": "反方审稿", "challenge.enabled": "反方审稿 · 启用",
|
||||
"challenge.max_claims": "反方审稿 · 每份报告最多审几句", "challenge.max_queries": "反方审稿 · 每句最多查询次数",
|
||||
"challenge.timeout_s": "反方审稿 · 每句时长上限（秒）",
|
||||
"inquiry": "审批质询", "inquiry.enabled": "审批质询 · 启用",
|
||||
"inquiry.max_queries": "审批质询 · 每次最多查询次数", "inquiry.timeout_s": "审批质询 · 每次时长上限（秒）",
|
||||
}
|
||||
|
||||
|
||||
def _positive(value: Any, *, whole: bool) -> bool:
|
||||
"""上限：null（不限），或大于 0 的有限数；whole 时还得是整数（次数、句数）。True / False 不算数。"""
|
||||
if value is None:
|
||||
return True
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0:
|
||||
return False
|
||||
return not (whole and int(value) != value)
|
||||
|
||||
|
||||
def _tolerance_ok(value: Any) -> bool:
|
||||
"""允许的相对误差：不小于 0 的有限数（0 表示必须完全一致）。不能是 null：没有「不限」的说法。"""
|
||||
return not isinstance(value, bool) and isinstance(value, (int, float)) and math.isfinite(value) and value >= 0
|
||||
|
||||
|
||||
def _verify_item_ok(key: str, value: Any) -> bool:
|
||||
if key == "enabled":
|
||||
return isinstance(value, bool)
|
||||
if key == "tolerance":
|
||||
return _tolerance_ok(value)
|
||||
return _positive(value, whole=key in _VERIFY_COUNTS)
|
||||
|
||||
|
||||
def _verify_item_problem(part: str, key: str, value: Any) -> str | None:
|
||||
label = VERIFY_SETTING_LABEL[f"{part}.{key}"]
|
||||
if _verify_item_ok(key, value):
|
||||
return None
|
||||
if key == "enabled":
|
||||
return f"「{label}」需要选择开启或关闭，当前为「{value}」"
|
||||
if key == "tolerance":
|
||||
return f"「{label}」需要填写不小于 0 的数（例如 0.000001 表示百万分之一），当前为「{value}」"
|
||||
words = "正整数（次数）" if key == "max_queries" else "正整数（句数）" if key == "max_claims" else "大于 0 的秒数"
|
||||
return f"「{label}」需要填写{words}，留空表示不限，当前为「{value}」"
|
||||
|
||||
|
||||
def verify_problem(value: Any) -> str | None:
|
||||
"""设置页写进来的 verify 这一组有什么问题，没问题返回 None。说明写设置页上的叫法。"""
|
||||
if not isinstance(value, dict):
|
||||
return f"「核查」设置格式有误，需要是一个对象，当前为「{value}」"
|
||||
for key, item in value.items():
|
||||
if key not in VERIFY_DEFAULTS:
|
||||
return (f"「核查」设置中不支持「{key}」，可设置的项："
|
||||
+ "、".join(f"「{VERIFY_SETTING_LABEL[k]}」" for k in VERIFY_DEFAULTS))
|
||||
label = VERIFY_SETTING_LABEL[key]
|
||||
if key in ("provider", "model"):
|
||||
if item is not None and not isinstance(item, str):
|
||||
return f"「{label}」需要填写模型接入的名称或模型 ID，留空表示沿用证据裁判模型，当前为「{item}」"
|
||||
elif key == "daily_max_usd":
|
||||
if not _positive(item, whole=False):
|
||||
return f"「{label}」需要填写大于 0 的金额（美元），留空表示不限，当前为「{item}」"
|
||||
elif not isinstance(item, dict):
|
||||
return f"「{label}」设置格式有误，需要是一个对象，当前为「{item}」"
|
||||
else:
|
||||
allowed = VERIFY_DEFAULTS[key]
|
||||
for sub, sub_value in item.items():
|
||||
if sub not in allowed:
|
||||
return (f"「{label}」设置中不支持「{sub}」，可设置的项："
|
||||
+ "、".join(f"「{VERIFY_SETTING_LABEL[f'{key}.{k}']}」" for k in allowed))
|
||||
if problem := _verify_item_problem(key, sub, sub_value):
|
||||
return problem
|
||||
return None
|
||||
|
||||
|
||||
def verify_values(stored: Any) -> dict[str, Any]:
|
||||
"""库里存的 verify 一组 → 补齐默认的完整设置。存坏了的项（旧版本、手改）按默认，存了 null 的上限就是不限；
|
||||
三项各自那一组逐项补，只存了一项的其余照默认。"""
|
||||
stored = stored if isinstance(stored, dict) else {}
|
||||
out: dict[str, Any] = {k: dict(v) if isinstance(v, dict) else v for k, v in VERIFY_DEFAULTS.items()}
|
||||
for key in ("provider", "model"):
|
||||
value = stored.get(key)
|
||||
out[key] = value.strip() or None if isinstance(value, str) else None
|
||||
if "daily_max_usd" in stored and _positive(stored["daily_max_usd"], whole=False):
|
||||
out["daily_max_usd"] = stored["daily_max_usd"]
|
||||
for part in _VERIFY_PARTS:
|
||||
saved = stored.get(part)
|
||||
if not isinstance(saved, dict):
|
||||
continue
|
||||
for key in VERIFY_DEFAULTS[part]:
|
||||
if key in saved and _verify_item_ok(key, saved[key]):
|
||||
out[part][key] = saved[key]
|
||||
return out
|
||||
|
||||
|
||||
async def verify_settings(session: AsyncSession) -> dict[str, Any]:
|
||||
"""设置里的 verify 一组，没存过的键补默认（读法同 judge_settings）。节点、质询接口都从这里读。"""
|
||||
row = await session.get(Setting, VERIFY_GROUP)
|
||||
return verify_values(row.value if row else None)
|
||||
|
||||
|
||||
def _model_text(value: Any) -> str | None:
|
||||
return value.strip() or None if isinstance(value, str) else None
|
||||
|
||||
|
||||
async def verify_model_spec(session: AsyncSession, node_config: dict[str, Any] | None = None, *,
|
||||
settings: dict[str, Any] | None = None) -> ModelSpec:
|
||||
"""核查用哪个模型：节点配置里的 provider / model > 设置里的核查模型 > 证据裁判模型的取值链（裁判设置 >
|
||||
助手模型 > 默认接入，见 judge.judge_model_spec）。哪一级写了就整组用哪一级，不跨级拼接入和模型。
|
||||
|
||||
和裁判一样关掉 thinking、max_tokens 4096：核查最后要按结构交结果（强制工具调用），和思考模式不兼容。
|
||||
建议和被核查的模型不同，否则等于自己核自己；比较、提醒由调用方做（它知道被核查的是哪个模型）。
|
||||
"""
|
||||
node = node_config if isinstance(node_config, dict) else {}
|
||||
provider, model = _model_text(node.get("provider")), _model_text(node.get("model"))
|
||||
if not (provider or model):
|
||||
chosen = settings if settings is not None else await verify_settings(session)
|
||||
provider, model = _model_text(chosen.get("provider")), _model_text(chosen.get("model"))
|
||||
if not (provider or model):
|
||||
return await judge_model_spec(session, None)
|
||||
return ModelSpec(provider=provider, model=model, max_tokens=JUDGE_MAX_TOKENS, thinking="off")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 用户设置
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -440,6 +590,8 @@ DEFAULT_SETTINGS: dict[str, Any] = {
|
||||
# 结论句裁判（engine/judge.py）：证据裁判模型和各项上限，上限存 null 表示不限。不放进 limits：
|
||||
# 那一组是环境变量定的、只读，这一组是用户自己定的
|
||||
JUDGE_GROUP: dict(JUDGE_DEFAULTS),
|
||||
# 核查（双盲复算、反方审稿、审批质询）：核查模型、三项的开关和上限、每日金额上限。和裁判一组一样由用户定
|
||||
VERIFY_GROUP: verify_values(None),
|
||||
}
|
||||
|
||||
|
||||
@@ -488,11 +640,13 @@ async def get_settings(session: AsyncSession = Depends(get_session)) -> dict[str
|
||||
# limits 是服务端配置（环境变量），只读：库里可能存着旧版设置页写进去的一份
|
||||
#（见过 max_agent_steps: 25），它从来不生效，却会盖住真实的上限，界面照着错的数许诺
|
||||
stored.pop("limits", None)
|
||||
# 裁判的每日计数是引擎记的账，不是用户设置：设置页整组回写时带着一份旧的，会把今天的花费清零
|
||||
# 裁判、核查的每日计数是引擎记的账，不是用户设置：设置页整组回写时带着一份旧的，会把今天的花费清零
|
||||
stored.pop(JUDGE_SPEND, None)
|
||||
stored.pop(VERIFY_SPEND, None)
|
||||
merged = {k: {**v, **(stored.get(k) or {})} for k, v in DEFAULT_SETTINGS.items()}
|
||||
# 裁判设置按引擎实际生效的样子给：存坏了的项显示默认，存了 null 的上限显示不限
|
||||
# 裁判、核查设置按引擎实际生效的样子给：存坏了的项显示默认，存了 null 的上限显示不限
|
||||
merged[JUDGE_GROUP] = judge_values(stored.get(JUDGE_GROUP))
|
||||
merged[VERIFY_GROUP] = verify_values(stored.get(VERIFY_GROUP))
|
||||
for key, value in stored.items():
|
||||
merged.setdefault(key, value)
|
||||
return merged
|
||||
@@ -508,14 +662,16 @@ async def put_settings(
|
||||
) -> dict[str, Any]:
|
||||
from app.tools.trust import SETTING_KEY as TRUST_KEY
|
||||
|
||||
# 先查后写：一组里有一项不对就整个请求不落库，免得存下半组
|
||||
if JUDGE_GROUP in payload.values and (problem := judge_problem(payload.values[JUDGE_GROUP])):
|
||||
# 先查后写：一组里有一项不对就整个请求不落库，免得存下半组
|
||||
raise HTTPException(422, problem)
|
||||
if VERIFY_GROUP in payload.values and (problem := verify_problem(payload.values[VERIFY_GROUP])):
|
||||
raise HTTPException(422, problem)
|
||||
for key, value in payload.values.items():
|
||||
if key in (TRUST_KEY, "limits", JUDGE_SPEND):
|
||||
if key in (TRUST_KEY, "limits", JUDGE_SPEND, VERIFY_SPEND):
|
||||
# 信任三档只走 PUT /api/tools/trust：设置页整组回写时带着一份旧的，
|
||||
# 会把刚在工具页改的档位冲掉。limits 是环境变量定的，写进库也不生效。
|
||||
# 裁判的每日计数是引擎记的账，只有裁判自己写
|
||||
# 裁判、核查的每日计数是引擎记的账，只有引擎自己写
|
||||
continue
|
||||
row = await session.get(Setting, key)
|
||||
if row:
|
||||
@@ -536,6 +692,15 @@ async def get_judge_spend(session: AsyncSession = Depends(get_session)) -> dict[
|
||||
return {**await daily_spend(), "daily_max_usd": (await judge_settings(session))["daily_max_usd"]}
|
||||
|
||||
|
||||
@router.get("/settings/verify/spend")
|
||||
async def get_verify_spend(session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""今天（本地日期）核查（双盲复算、反方审稿、审批质询）花了多少，设置页放在核查的每日上限旁边。只读。
|
||||
|
||||
形状同 /settings/judge/spend：{date, usd, calls, unpriced_calls, daily_max_usd}，daily_max_usd 为 null 表示不限。
|
||||
"""
|
||||
return {**await verify_daily_spend(), "daily_max_usd": (await verify_settings(session))["daily_max_usd"]}
|
||||
|
||||
|
||||
@router.get("/system")
|
||||
async def system_info() -> dict[str, Any]:
|
||||
"""给设置页顶部展示的运行环境概览。"""
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,98 @@
|
||||
"""核查花费账：双盲复算、反方审稿、审批质询里的模型调用，按本地日期记一本每日账。
|
||||
|
||||
照结论句裁判的每日计数做（judge.py 的 SPEND_KEY、_charge、_settle）：调用前按目录价估算、先预留，加上这笔会超过
|
||||
每日上限就不记；调用后按实际用量找补差额。账落在设置表里，按本地日期滚动。
|
||||
|
||||
**为什么和裁判分开记。** 核查是一个 agent 查库、来回好几轮，一次花的钱比裁判判一份报告多得多。记在同一本账上，
|
||||
核查一多，当天的裁判就全因为「每日上限」判不了——报告的出具会因此降档，而降档的原因和这份报告毫无关系。两本账
|
||||
各管各的上限，互不挤占。
|
||||
|
||||
目录里没有价格的模型估不出金额：只计次数（unpriced_calls），金额上限对它不起作用，由调用方照实说明。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import weakref
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Setting
|
||||
|
||||
#: 每日计数存在设置表的这一行：{date: 本地日期, usd, calls, unpriced_calls}。不从设置接口进出（api/settings.py）
|
||||
SPEND_KEY = "verify_spend"
|
||||
|
||||
_locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]" = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
def _lock() -> asyncio.Lock:
|
||||
"""读改写这一行时持的锁，按事件循环各一把（asyncio.Lock 不能跨循环用）。
|
||||
|
||||
几个核查并行时各自预留：不串行的话两边读到同一个余额，会一起越过上限。
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
lock = _locks.get(loop)
|
||||
if lock is None:
|
||||
lock = _locks[loop] = asyncio.Lock()
|
||||
return lock
|
||||
|
||||
|
||||
def _today() -> str:
|
||||
"""本地日期：每日上限按它滚动。"""
|
||||
return datetime.now().astimezone().date().isoformat()
|
||||
|
||||
|
||||
def _spend_of(row: Any) -> dict[str, Any]:
|
||||
value = row.value if row is not None and isinstance(row.value, dict) else {}
|
||||
today = _today()
|
||||
if value.get("date") != today:
|
||||
return {"date": today, "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
return {"date": today, "usd": float(value.get("usd") or 0.0), "calls": int(value.get("calls") or 0),
|
||||
"unpriced_calls": int(value.get("unpriced_calls") or 0)}
|
||||
|
||||
|
||||
async def daily_spend() -> dict[str, Any]:
|
||||
"""今天（本地日期）核查已经花了多少：{date, usd, calls, unpriced_calls}。估不出金额的调用只计次数。"""
|
||||
async with SessionLocal() as session:
|
||||
return _spend_of(await session.get(Setting, SPEND_KEY))
|
||||
|
||||
|
||||
async def charge(amount: float, *, limit: float | None, priced: bool) -> str | None:
|
||||
"""记一次调用、先按估算预留金额。加上这笔会超过每日上限就不记，返回 None；记下了返回记在哪天（找补时要用）。
|
||||
|
||||
limit 为 None 表示不限。估不出金额的调用（priced=False）照记次数，金额记 0，调用方应传 limit=None：
|
||||
拿 0 去比上限，等于说它不花钱。
|
||||
"""
|
||||
async with _lock():
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, SPEND_KEY)
|
||||
spend = _spend_of(row)
|
||||
if limit is not None and spend["usd"] + amount > limit + 1e-12:
|
||||
return None
|
||||
spend["usd"] = round(spend["usd"] + amount, 8)
|
||||
spend["calls"] += 1
|
||||
spend["unpriced_calls"] += 0 if priced else 1
|
||||
if row is None:
|
||||
session.add(Setting(key=SPEND_KEY, value=spend))
|
||||
else:
|
||||
row.value = spend
|
||||
await session.commit()
|
||||
return spend["date"]
|
||||
|
||||
|
||||
async def settle(date: str, delta: float) -> None:
|
||||
"""调用结束后按实际用量找补预留的差额。跨了天就不找补：预留记在前一天，今天从零算。"""
|
||||
if not delta:
|
||||
return
|
||||
async with _lock():
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, SPEND_KEY)
|
||||
spend = _spend_of(row)
|
||||
if row is None or spend["date"] != date:
|
||||
return
|
||||
spend["usd"] = max(0.0, round(spend["usd"] + delta, 8))
|
||||
row.value = spend
|
||||
await session.commit()
|
||||
|
||||
|
||||
__all__ = ["SPEND_KEY", "charge", "daily_spend", "settle"]
|
||||
@@ -0,0 +1,982 @@
|
||||
"""核查用的只读 agent：查库，交回有出处的结果。双盲复算、反方审稿、审批质询共用这一个。
|
||||
|
||||
三处要的是同一件事：另起一个模型，只给它几个数据源的只读查询工具，让它自己写 SQL 去核对一个数、一句结论、一个
|
||||
提问，最后交回结构化的结果——而且结果里的数要能核对。
|
||||
|
||||
**只读。** 工具只建给定数据源的 db_schema / db_query，建的时候 force_readonly：可写源上的写语句也被守卫拒绝
|
||||
（tools/datasource.py）。这一路没有审批关卡，不能指望人在半路拦下一条 DELETE。
|
||||
|
||||
**数从快照取。** 每次成功的查询是一份查询快照（工具自己落的 put_json），按这个 agent 自己的顺序编 Q1、Q2…。
|
||||
最后交结果时，调用方结构里的数值字段一律改成格子引用 {alias, row, column}，由代码按引用从快照里取值：模型自己
|
||||
写的数字不采信，取不到就是 None，不兜底成 0。数值字段的各种写法都认（type 列表、anyOf 里的数、$ref 引用进来的、
|
||||
allOf 包着的）；认不出的结构下模型写的数一律丢掉、照记下来，宁可少一个数，不让它绕过格子引用。理由文字里的数字
|
||||
要写成 [[v:Qn.rX.列]]，用本 agent 自己的台账（evidence.build_catalog）组装成文档（evidence.compose_doc），裸数字
|
||||
会被标出来。
|
||||
|
||||
**不进写作目录。** 节点里调用时（给了 emit）每次工具调用发 tool.start / tool.end，查询的 tool.end 带
|
||||
query_artifact 和 role: verify——随运行封存、证据面板能下钻。但这里绝不往 state["evidence"] 写 kind: query：
|
||||
报告的写作目录按台账里的 query 编号，核查的查询混进去，写作者就会把它们当成可以引用的证据。调用方只拿到结果，
|
||||
台账条目由它按自己的种类记（dual_check / challenge / inquiry）。
|
||||
|
||||
**上限。** 查询次数、时长、金额（这一次、全局每日）任何一项到顶就停止查询，按已经查到的结果收尾：收尾那次
|
||||
调用总是要做的（时长到顶时再给 SETTLE_GRACE_S 的宽限），所以每一轮都按收尾最贵的样子给它留出钱——收尾时要把
|
||||
查到的结果交给模型，结果有多长事先不知道，按 _digest 的上限算。留过钱的收尾不再被上限拦下：某一轮的实际用量
|
||||
超出了估算、或者并行的核查先用掉了当天的余额时，宁可超出上限一次收尾调用的钱，也不让已经查到的结果白查。
|
||||
每日上限不传就按设置里的「核查每日金额上限」，明确传 None 才是不限。模型不在价格目录里时估不出金额，金额上限
|
||||
对它无效，照裁判的说法记一条说明。每次调用都记进核查的每日账（engine/spend.py）。
|
||||
|
||||
永远不抛：模型没配好、调用失败、格式不对、触顶、库写不进去，都落在返回值里（error / limit）。模型调用都经
|
||||
get_chat_model，测试换成剧本。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
from app.core import artifact_store
|
||||
from app.core.errors import describe_exception, not_configured
|
||||
from app.core.events import EventType
|
||||
from app.db.base import SessionLocal
|
||||
from app.engine import spend
|
||||
from app.engine.evidence import build_catalog, cell_eid, compose_doc, query_entry_fields, schema_ledger_entry
|
||||
from app.engine.expressions import CellError, column_kind, locate_cell, numeric_text
|
||||
from app.engine.judge import _json_in, _tokens, _usd
|
||||
from app.engine.nodes.llm import _DIGEST_BUDGET, _digest, _split_structured, _typed, explain_model_error
|
||||
from app.engine.state import message_text
|
||||
from app.engine.toolcalls import ToolTimeout, limit_fields, limit_of, run_bounded
|
||||
from app.providers import catalog as pricing
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, bind_tools_safely, get_chat_model
|
||||
from app.tools.datasource import QUERY_PREFIX, SCHEMA_PREFIX, build_datasource_tools
|
||||
from app.tools.registry import ToolArgsError, ToolContext, Versions, args_model_of, prepare_args
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: 时长到顶之后，收尾那次调用（交结构化结果）最多再等这么久。不给宽限的话，到点时手里的查询全白查了
|
||||
SETTLE_GRACE_S = 30.0
|
||||
#: 估算金额时每一轮、收尾那次的输出 token 数（回复里带了用量就按实际找补）
|
||||
_OUT_STEP = 512
|
||||
_OUT_FINAL = 1024
|
||||
#: 工具结果交给模型时最多这么多字：完整结果在查询快照里，收尾时另按 _digest 列出
|
||||
_TOOL_CHARS = 12_000
|
||||
#: 收尾时带上的模型最后那段说明最多这么多字。收尾那次调用的长度有了上限，循环里才能按它把钱留够
|
||||
_SAID_CHARS = 2_000
|
||||
#: 收尾那次调用里查询结果最多占多少字：_digest 的上限，加上它截断时末尾那句说明
|
||||
_DIGEST_ROOM = _DIGEST_BUDGET + 64
|
||||
#: 留钱时给「已达上限（…）」那句话留的字数
|
||||
_LIMIT_ROOM = 40
|
||||
#: 本地 $ref 最多展开这么多层
|
||||
_REF_DEPTH = 8
|
||||
#: 结果里自动加上的理由字段（调用方的结构里已经有同名的字符串字段就沿用它）
|
||||
RATIONALE = "rationale"
|
||||
#: 触顶原因
|
||||
LIMITS = ("max_queries", "timeout_s", "max_cost_usd", "daily_max_usd", "max_steps")
|
||||
|
||||
VERIFY_SYSTEM = (
|
||||
"你是数据核查员：只读地查库，用查询结果完成交给你的任务。\n"
|
||||
"- 先用 db_schema 工具看表结构，确认表名、字段名，再用 db_query 工具写只读 SQL。不能改数据，写语句会被拒绝\n"
|
||||
"- 每次成功的查询前面都标着编号【证据 Qn】，最后交结果时用它说明每个数取自哪次查询的哪一格\n"
|
||||
"- {limit}先想清楚要查什么，不要重复同样的查询\n"
|
||||
"- 不要编造数字，查不到就照实说"
|
||||
)
|
||||
VERIFY_FINAL_PROMPT = (
|
||||
"现在交结果，不能再查询。按给定结构填写：\n"
|
||||
"- 数值字段一律写格子引用，形如 {\"alias\": \"Q1\", \"row\": 0, \"column\": \"列名\"}：alias 是【证据 Qn】里的编号，"
|
||||
"row 是行号（从 0 数），column 是列名。系统按引用从查询结果里取值，你自己写的数字不会被采用；查不到就写 null，"
|
||||
"不要估算\n"
|
||||
"- rationale 写理由，简短说清依据。理由里出现的每个数字都写成 [[v:Q1.r0.列名]] 这样的引用，系统会换成查询结果"
|
||||
"里的值；不要直接写数字，提到某次查询时写 [[see:Q1]]\n"
|
||||
"- 其余字段按各自的说明填写"
|
||||
)
|
||||
_CELL_REF: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"description": "格子引用：alias 是【证据 Qn】里的编号，row 是行号（从 0 数），column 是列名",
|
||||
"properties": {"alias": {"type": "string"}, "row": {"type": "integer"}, "column": {"type": "string"}},
|
||||
"required": ["alias", "row", "column"],
|
||||
}
|
||||
_RATIONALE_FIELD: dict[str, Any] = {
|
||||
"type": "string",
|
||||
"description": "理由：简短说清依据，里面的每个数字都写成 [[v:Q1.r0.列名]] 这样的引用",
|
||||
}
|
||||
|
||||
Emit = Callable[..., Any]
|
||||
|
||||
|
||||
class Default(Enum):
|
||||
"""run_verify_agent 的 daily_max_usd 缺省值：按设置里的「核查每日金额上限」（api/settings.py 的
|
||||
verify_settings）。节点、质询接口不必各自记得去读；明确传 None 才是不限。"""
|
||||
|
||||
SETTINGS = "settings"
|
||||
|
||||
|
||||
@dataclass
|
||||
class VerifyOutcome:
|
||||
"""一次核查的结果。
|
||||
|
||||
result：调用方结构的形状，数值字段是从快照取到的值（取不到为 None），另有 rationale（理由渲染后的正文）；
|
||||
没交成结果时为 None，error 说明为什么。cells：每个数值字段的出处 {path, status, ref, alias, row, column,
|
||||
artifact, eid, value}，status 是 resolved（value 是那一格的值，那一格本身是 NULL 时为 None）/ missing（模型
|
||||
交了 null：查不到）/ unresolved（reason 写原因：没交这个字段、引用指不到、指到的不是数；模型直接写了数的记
|
||||
model_value，非数值字段里写的数也照此记下）。queries：成功的查询 {alias, artifact, sql, row_count, truncated,
|
||||
source, tool, call_id}。
|
||||
rationale：理由组装成的报告文档（evidence.compose_doc 的形状，violations 里有裸数字）。usage：{input_tokens,
|
||||
output_tokens, total_tokens, cost_usd, calls, unpriced_calls}。limit：触顶原因（LIMITS 之一），没触顶为 None。
|
||||
answered：模型是不是自己停下来的（没被上限掐断）。
|
||||
"""
|
||||
|
||||
result: dict[str, Any] | None
|
||||
cells: list[dict[str, Any]]
|
||||
queries: list[dict[str, Any]]
|
||||
rationale: dict[str, Any] | None
|
||||
usage: dict[str, Any]
|
||||
limit: str | None = None
|
||||
error: str | None = None
|
||||
model: str | None = None
|
||||
answered: bool = False
|
||||
notes: list[str] = field(default_factory=list)
|
||||
duration_ms: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 结构：数值字段换成格子引用
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _types(schema: Any) -> set[str]:
|
||||
declared = schema.get("type") if isinstance(schema, dict) else None
|
||||
return {declared} if isinstance(declared, str) else set(declared) if isinstance(declared, list) else set()
|
||||
|
||||
|
||||
def _is_object(schema: Any) -> bool:
|
||||
return isinstance(schema, dict) and ("object" in _types(schema) or "properties" in schema)
|
||||
|
||||
|
||||
def _is_array(schema: Any) -> bool:
|
||||
return isinstance(schema, dict) and ("array" in _types(schema) or "items" in schema)
|
||||
|
||||
|
||||
def _numeric_types(schema: Any) -> set[str] | None:
|
||||
"""数值字段允许的类型（number / integer，连同写在一起的 null、string）；不是数值字段返回 None。
|
||||
|
||||
type 写成 "number"、["number", "null"]，或者 anyOf / oneOf 里有一支是数（可为 null 的数一般这么写，
|
||||
Pydantic 的 Optional[float] 就是）都算：有一支能放数，整个字段就得是格子引用，不然数从那一支漏过去。
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return None
|
||||
types = set(_types(schema))
|
||||
for key in ("anyOf", "oneOf"):
|
||||
for branch in schema.get(key) or ():
|
||||
if isinstance(branch, dict) and not _is_object(branch) and not _is_array(branch):
|
||||
types |= _numeric_types(branch) or _types(branch)
|
||||
return types if types & {"number", "integer"} else None
|
||||
|
||||
|
||||
def _has_numbers(schema: Any, depth: int = 0) -> bool:
|
||||
"""结构里有没有数值字段（嵌套的也算）。"""
|
||||
if not isinstance(schema, dict) or depth > 32:
|
||||
return False
|
||||
if _numeric_types(schema) is not None:
|
||||
return True
|
||||
subs = [*(schema.get("properties") or {}).values(), schema.get("items"), schema.get("additionalProperties"),
|
||||
*(schema.get("anyOf") or ()), *(schema.get("oneOf") or ())]
|
||||
return any(_has_numbers(sub, depth + 1) for sub in subs if isinstance(sub, dict))
|
||||
|
||||
|
||||
#: 值是「名字 → 结构」的键、值是一个结构的键、值是一串结构的键：展开引用时往里走；其余的键（enum、description…）原样
|
||||
_SCHEMA_MAPS = ("properties", "patternProperties")
|
||||
_SCHEMA_ONE = ("items", "additionalProperties", "contains", "not")
|
||||
_SCHEMA_LIST = ("anyOf", "oneOf", "prefixItems")
|
||||
|
||||
|
||||
def _merge(a: dict[str, Any], b: dict[str, Any]) -> dict[str, Any]:
|
||||
"""allOf 的两支并成一个节点：属性合在一起，必填合在一起，类型取并集（只用来认是不是数、是不是对象）。"""
|
||||
out = dict(a)
|
||||
for key, value in b.items():
|
||||
if key == "properties" and isinstance(value, dict):
|
||||
out["properties"] = {**(out.get("properties") or {}), **value}
|
||||
elif key == "required" and isinstance(value, list):
|
||||
out["required"] = list(dict.fromkeys([*(out.get("required") or []), *value]))
|
||||
elif key == "type":
|
||||
merged = _types(out) | _types(b)
|
||||
out["type"] = next(iter(merged)) if len(merged) == 1 else sorted(merged)
|
||||
else:
|
||||
out.setdefault(key, value)
|
||||
return out
|
||||
|
||||
|
||||
def _inline(schema: Any) -> dict[str, Any]:
|
||||
"""展开调用方结构里的本地引用（$ref: #/$defs/…、#/definitions/…），allOf 并成一个节点，去掉 $defs。
|
||||
|
||||
后面换格子引用、取值只认 type / properties / items / anyOf 这几种写法：引用不展开，引用进来的数值字段就绕过了
|
||||
格子引用这一关（Pydantic 生成的结构里，嵌套的对象全是 $ref）。引用不到的、绕回自己的（树形结构第二次遇到同一个
|
||||
引用）就地断开，当作认不出的结构：模型在那里写的数照样不采信。
|
||||
"""
|
||||
root = schema if isinstance(schema, dict) else {}
|
||||
|
||||
def target(ref: Any) -> dict[str, Any] | None:
|
||||
if not isinstance(ref, str) or not ref.startswith("#/"):
|
||||
return None
|
||||
node: Any = root
|
||||
for part in ref[2:].split("/"):
|
||||
node = node.get(part.replace("~1", "/").replace("~0", "~")) if isinstance(node, dict) else None
|
||||
return node if isinstance(node, dict) else None
|
||||
|
||||
def expand(node: Any, trail: tuple[str, ...]) -> Any:
|
||||
while isinstance(node, dict) and "$ref" in node:
|
||||
ref, rest = node["$ref"], {k: v for k, v in node.items() if k != "$ref"}
|
||||
found = target(ref)
|
||||
if found is None or ref in trail or len(trail) >= _REF_DEPTH:
|
||||
node = rest
|
||||
break
|
||||
node, trail = {**found, **rest}, (*trail, ref)
|
||||
if not isinstance(node, dict):
|
||||
return node
|
||||
out: dict[str, Any] = {}
|
||||
for key, value in node.items():
|
||||
if key in ("$defs", "definitions", "allOf"):
|
||||
continue
|
||||
if key in _SCHEMA_MAPS and isinstance(value, dict):
|
||||
out[key] = {k: expand(v, trail) for k, v in value.items()}
|
||||
elif key in _SCHEMA_ONE and isinstance(value, dict):
|
||||
out[key] = expand(value, trail)
|
||||
elif key in _SCHEMA_LIST and isinstance(value, list):
|
||||
out[key] = [expand(v, trail) for v in value]
|
||||
else:
|
||||
out[key] = value
|
||||
for part in node.get("allOf") or ():
|
||||
if isinstance(part, dict):
|
||||
out = _merge(out, expand(part, trail))
|
||||
return out
|
||||
|
||||
return expand(root, ())
|
||||
|
||||
|
||||
def result_schema(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""交给模型的结构：调用方结构里的数值字段（见 _numeric_types；嵌套的对象、数组、anyOf 的分支里的也算，引用先
|
||||
展开）一律换成可为 null 的格子引用，另加一个理由字段（已有同名的字符串字段就沿用）。其余字段原样。"""
|
||||
|
||||
def rewrite(node: Any) -> Any:
|
||||
if not isinstance(node, dict):
|
||||
return node
|
||||
if _numeric_types(node) is not None:
|
||||
out: dict[str, Any] = {"anyOf": [_CELL_REF, {"type": "null"}]}
|
||||
if node.get("description"):
|
||||
out["description"] = f"{node['description']}（写格子引用，查不到写 null）"
|
||||
return out
|
||||
out = dict(node)
|
||||
if isinstance(node.get("properties"), dict):
|
||||
out["properties"] = {k: rewrite(v) for k, v in node["properties"].items()}
|
||||
out.setdefault("type", "object")
|
||||
for key in ("items", "additionalProperties"):
|
||||
if isinstance(node.get(key), dict):
|
||||
out[key] = rewrite(node[key])
|
||||
for key in ("anyOf", "oneOf"):
|
||||
if isinstance(node.get(key), list):
|
||||
out[key] = [rewrite(b) for b in node[key]]
|
||||
return out
|
||||
|
||||
inlined = _inline(schema)
|
||||
root = rewrite(inlined if _is_object(inlined) else {"type": "object", "properties": {}})
|
||||
props = dict(root.get("properties") or {})
|
||||
props.setdefault(RATIONALE, _RATIONALE_FIELD)
|
||||
required = list(dict.fromkeys([*(root.get("required") or []), RATIONALE]))
|
||||
# 走工具调用的结构化输出要求函数名只含字母、数字、下划线
|
||||
return {**root, "title": "verify_result", "type": "object", "properties": props, "required": required}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Gap:
|
||||
"""模型交的结果里这个位置没有值。reason 说为什么：没交这个字段、上一层的格式不对。"""
|
||||
|
||||
reason: str
|
||||
|
||||
|
||||
_LEFT_OUT = _Gap("模型没有交这个字段")
|
||||
|
||||
|
||||
def _fits(branch: dict[str, Any], value: Any) -> bool:
|
||||
"""anyOf / oneOf 的这一支收不收这个值（按形状认）。"""
|
||||
if isinstance(value, _Gap):
|
||||
return _has_numbers(branch)
|
||||
if isinstance(value, dict):
|
||||
return _is_object(branch)
|
||||
if isinstance(value, list):
|
||||
return _is_array(branch)
|
||||
if value is None:
|
||||
return "null" in _types(branch)
|
||||
if isinstance(value, bool):
|
||||
return "boolean" in _types(branch)
|
||||
if isinstance(value, str):
|
||||
return "string" in _types(branch) or "enum" in branch or "const" in branch
|
||||
# 数：能放数的分支已经整个换成了格子引用（_numeric_types），这里只剩码值里的数。True == 1，真假不算数
|
||||
choices = [*(branch.get("enum") or ()), *([branch["const"]] if "const" in branch else [])]
|
||||
return any(c == value and not isinstance(c, bool) for c in choices)
|
||||
|
||||
|
||||
def _decoded(value: Any, kind: type) -> Any:
|
||||
"""有的模型把嵌套的对象、数组写成一段 JSON 文本（裁判见过同样的事，judge._items_of）：解得开、形状对就用解开的。"""
|
||||
if isinstance(value, str) and isinstance(found := _json_in(value), kind):
|
||||
return found
|
||||
return value
|
||||
|
||||
|
||||
class _Cells:
|
||||
"""按格子引用从查询快照里取值：快照经工件库读（复验哈希），同一份只读一次。"""
|
||||
|
||||
def __init__(self, queries: list[dict[str, Any]]) -> None:
|
||||
self.by_alias = {q["alias"]: q for q in queries}
|
||||
self.memo: dict[str, tuple[Any, str | None]] = {}
|
||||
self.cells: list[dict[str, Any]] = []
|
||||
|
||||
def snapshot(self, artifact: str) -> tuple[Any, str | None]:
|
||||
if artifact not in self.memo:
|
||||
try:
|
||||
content = artifact_store.load(artifact)
|
||||
self.memo[artifact] = (content, None if content is not None else "查询快照无法读取")
|
||||
except ValueError:
|
||||
self.memo[artifact] = (None, "查询快照与哈希不一致，可能已被修改")
|
||||
return self.memo[artifact]
|
||||
|
||||
def take(self, value: Any, path: str, types: set[str]) -> Any:
|
||||
"""一个数值字段：按引用取值，记下出处，返回值（取不到为 None）。types 是字段允许的类型：不许文本的，
|
||||
引用的格子不是数（日期、名称、真假）就不采信——拿它和原流程的数比，比较本身就会出错。"""
|
||||
if isinstance(value, _Gap):
|
||||
self.cells.append({"path": path, "status": "unresolved", "reason": value.reason})
|
||||
return None
|
||||
if value is None:
|
||||
self.cells.append({"path": path, "status": "missing"})
|
||||
return None
|
||||
value = _decoded(value, dict)
|
||||
if not isinstance(value, dict):
|
||||
# 模型直接写了数：不采信。照记下来，界面上能看到它想填什么
|
||||
number = isinstance(value, (int, float)) or numeric_text(value, loose=True) is not None
|
||||
self.cells.append({"path": path, "status": "unresolved", "model_value": value,
|
||||
"reason": "模型直接写了数值，没有给出格子引用，不采信" if number
|
||||
else "格子引用的格式不对（应写成 {alias, row, column}），不采信"})
|
||||
return None
|
||||
alias = str(value.get("alias") or value.get("call") or "").strip()
|
||||
row, column = value.get("row"), value.get("column")
|
||||
if isinstance(row, str) and row.strip().isdigit():
|
||||
row = int(row) # 有的模型把整数写成 "0"：意思没有歧义
|
||||
elif isinstance(row, float) and row.is_integer():
|
||||
row = int(row)
|
||||
query = self.by_alias.get(alias)
|
||||
if query is None:
|
||||
self.cells.append({"path": path, "status": "unresolved", "alias": alias or None,
|
||||
"reason": f"没有编号为 {alias or '（空）'} 的查询，可用的查询："
|
||||
+ ("、".join(self.by_alias) or "（无）")})
|
||||
return None
|
||||
content, why = self.snapshot(query["artifact"])
|
||||
if why:
|
||||
self.cells.append({"path": path, "status": "unresolved", "alias": alias, "artifact": query["artifact"],
|
||||
"reason": f"{alias} 的{why}"})
|
||||
return None
|
||||
try:
|
||||
raw, name = locate_cell(content, row, column)
|
||||
except CellError as e:
|
||||
self.cells.append({"path": path, "status": "unresolved", "alias": alias, "artifact": query["artifact"],
|
||||
"reason": f"{alias} {e}"})
|
||||
return None
|
||||
truth = _typed(raw, {"type": sorted(types)}, column_kind(content, name))
|
||||
cited = {"ref": f"{alias}.r{row}.{name}", "alias": alias, "row": row, "column": name,
|
||||
"artifact": query["artifact"], "eid": cell_eid(query["artifact"], row, name), "value": truth}
|
||||
if truth is not None and "string" not in types and (isinstance(truth, bool)
|
||||
or not isinstance(truth, (int, float))):
|
||||
self.cells.append({"path": path, "status": "unresolved", **cited,
|
||||
"reason": f"{alias} 第 {row} 行「{name}」的值是「{truth}」，不是数值，不采信"})
|
||||
return None
|
||||
self.cells.append({"path": path, "status": "resolved", **cited})
|
||||
return truth
|
||||
|
||||
def walk(self, schema: Any, value: Any, path: str = "") -> Any:
|
||||
"""按调用方的结构（_inline 展开过的）走一遍模型交的结果：数值字段换成从快照取到的值，其余原样。"""
|
||||
types = _numeric_types(schema)
|
||||
if types is not None:
|
||||
return self.take(value, path, types)
|
||||
if not isinstance(schema, dict):
|
||||
return self.scrub({}, value, path)
|
||||
branches = [b for key in ("anyOf", "oneOf") for b in schema.get(key) or () if isinstance(b, dict)]
|
||||
if branches:
|
||||
rest = {k: v for k, v in schema.items() if k not in ("anyOf", "oneOf")}
|
||||
if any(_is_object(b) for b in branches):
|
||||
value = _decoded(value, dict)
|
||||
if any(_is_array(b) for b in branches):
|
||||
value = _decoded(value, list)
|
||||
if value is None and any("null" in _types(b) for b in branches):
|
||||
return None
|
||||
chosen = next((b for b in branches if _fits(b, value)), None)
|
||||
return self.walk({**rest, **chosen} if chosen is not None else rest, value, path)
|
||||
if _is_object(schema):
|
||||
return self._object(schema, value, path)
|
||||
if _is_array(schema):
|
||||
return self._array(schema, value, path)
|
||||
return self.scrub(schema, value, path)
|
||||
|
||||
def _object(self, schema: dict[str, Any], value: Any, path: str) -> dict[str, Any]:
|
||||
given = _decoded(value, dict)
|
||||
if isinstance(given, dict):
|
||||
def child(key: str) -> Any:
|
||||
return given[key] if key in given else _LEFT_OUT
|
||||
else:
|
||||
gap = given if given is None or isinstance(given, _Gap) else _Gap("模型交的格式不对：这里应是一个对象")
|
||||
|
||||
def child(key: str) -> Any:
|
||||
return gap
|
||||
props = schema.get("properties") or {}
|
||||
out = {k: self.walk(sub, child(k), f"{path}.{k}" if path else k) for k, sub in props.items()}
|
||||
extra = schema.get("additionalProperties")
|
||||
for k, v in (given.items() if isinstance(given, dict) else ()):
|
||||
if k in props:
|
||||
continue
|
||||
if isinstance(extra, dict):
|
||||
out[k] = self.walk(extra, v, f"{path}.{k}" if path else k)
|
||||
elif extra is True or (not props and extra is not False):
|
||||
# 没写属性的对象（任意键）照收，里面的数照样不采信；写了属性的，没声明的键不要
|
||||
out[k] = self.scrub({}, v, f"{path}.{k}" if path else k)
|
||||
return out
|
||||
|
||||
def _array(self, schema: dict[str, Any], value: Any, path: str) -> list[Any] | None:
|
||||
items = _decoded(value, list)
|
||||
if isinstance(items, list):
|
||||
sub = schema.get("items") if isinstance(schema.get("items"), dict) else {}
|
||||
return [self.walk(sub, v, f"{path}[{i}]") for i, v in enumerate(items)]
|
||||
if _has_numbers(schema):
|
||||
if items is None:
|
||||
self.cells.append({"path": path, "status": "missing"})
|
||||
elif isinstance(items, _Gap):
|
||||
self.cells.append({"path": path, "status": "unresolved", "reason": items.reason})
|
||||
else:
|
||||
self.cells.append({"path": path, "status": "unresolved", "model_value": items,
|
||||
"reason": "模型交的格式不对：这里应是一个数组"})
|
||||
elif items is not None and not isinstance(items, _Gap):
|
||||
self.scrub({}, items, path)
|
||||
return None
|
||||
|
||||
def scrub(self, schema: Any, value: Any, path: str) -> Any:
|
||||
"""认不出的结构（没写类型、文本字段、没写 items 的数组……）下的值：文本、真假、null 原样，数一律不采信、
|
||||
照记下来（码值里的数是调用方给的选项，不是查出来的数，照收），对象、数组逐项照此处理。"""
|
||||
if isinstance(value, _Gap):
|
||||
return None
|
||||
if value is None or isinstance(value, (bool, str)):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
if isinstance(schema, dict) and _fits(schema, value):
|
||||
return value
|
||||
self.cells.append({"path": path, "status": "unresolved", "model_value": value,
|
||||
"reason": "这里不是数值字段，模型却写了数值，不采信"})
|
||||
return None
|
||||
if isinstance(value, dict):
|
||||
return {k: self.scrub({}, v, f"{path}.{k}" if path else k) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [self.scrub({}, v, f"{path}[{i}]") for i, v in enumerate(value)]
|
||||
return None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 花费
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _Budget:
|
||||
"""金额上限（这一次、全局每日）和核查的每日账。调用前按估算预留，调用后按实际用量找补。"""
|
||||
|
||||
def __init__(self, model_id: str, max_cost_usd: float | None, daily_max_usd: float | None) -> None:
|
||||
self.model_id = model_id
|
||||
self.priced = pricing.find_pricing(model_id) is not None
|
||||
self.max_cost_usd = max_cost_usd
|
||||
self.daily_max_usd = daily_max_usd
|
||||
self.spent = 0.0
|
||||
|
||||
def estimate(self, text: str, out: int, *, room: int = 0) -> float:
|
||||
"""按文本估一次调用的金额。room 是还不知道内容、按最坏情况（一个字一个 token）算的输入字数。"""
|
||||
return pricing.estimate_cost(self.model_id, _tokens(text) + room, out) if self.priced else 0.0
|
||||
|
||||
async def reserve(self, estimate: float, *, keep: float = 0.0, kept: bool = False
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""(记在哪天, None) 或 (None, 拦下的上限)。
|
||||
|
||||
keep：要额外留出的钱（收尾那次调用最贵的样子），只用来判断、不预留。kept：这就是循环里一直给它留着钱的
|
||||
那次收尾，不再按上限拦，照记账：每一轮都留够了，还超出只能是某一轮实际用得比估算多，或者并行的核查先用掉了
|
||||
当天的余额（每日余额只是比了一下、没有真的预留）。这时宁可超出上限一次收尾调用的钱，也不让查到的结果白查。
|
||||
"""
|
||||
cap = self.daily_max_usd if self.priced and not kept else None
|
||||
if not kept:
|
||||
if self.priced and self.max_cost_usd is not None and \
|
||||
self.spent + estimate + keep > self.max_cost_usd + 1e-12:
|
||||
return None, "max_cost_usd"
|
||||
if cap is not None and keep and (await spend.daily_spend())["usd"] + estimate + keep > cap + 1e-12:
|
||||
return None, "daily_max_usd"
|
||||
day = await spend.charge(estimate, limit=cap, priced=self.priced)
|
||||
return (day, None) if day is not None else (None, "daily_max_usd")
|
||||
|
||||
async def settle(self, day: str, estimate: float, inp: int, out: int) -> float:
|
||||
"""按实际用量找补，返回这一次的金额。回复没带用量（有的网关不给、调用超时）就按估算记：宁可多记。"""
|
||||
actual = (pricing.estimate_cost(self.model_id, inp, out) if inp or out else estimate) if self.priced else 0.0
|
||||
await spend.settle(day, actual - estimate)
|
||||
self.spent += actual
|
||||
return actual
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 主循环
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _announce(emit: Emit | None, kind: EventType, **data: Any) -> None:
|
||||
"""经 emit 发一条事件（同步、协程函数都认）；没给 emit 就不发。发不出去不影响核查本身。"""
|
||||
if emit is None:
|
||||
return
|
||||
try:
|
||||
out = emit(kind, **data)
|
||||
if inspect.isawaitable(out):
|
||||
await out
|
||||
except Exception: # noqa: BLE001 - 事件只是留痕
|
||||
logger.warning("核查的事件没发出去：%s", kind, exc_info=True)
|
||||
|
||||
|
||||
async def _bounded(job: Any, seconds: float | None) -> Any:
|
||||
if seconds is None:
|
||||
return await job
|
||||
return await asyncio.wait_for(job, timeout=max(seconds, 0.001))
|
||||
|
||||
|
||||
def _usage(message: Any) -> tuple[int, int]:
|
||||
meta = getattr(message, "usage_metadata", None) or {}
|
||||
return int(meta.get("input_tokens") or 0), int(meta.get("output_tokens") or 0)
|
||||
|
||||
|
||||
def _limit_words(reason: str, *, max_queries: int | None, timeout_s: float | None, max_cost_usd: float | None,
|
||||
daily_max_usd: float | None, steps: int) -> str:
|
||||
if reason == "max_queries":
|
||||
return f"最多查询 {max_queries} 次"
|
||||
if reason == "timeout_s":
|
||||
return f"时长上限 {timeout_s:g} 秒"
|
||||
if reason == "max_cost_usd":
|
||||
return f"这次核查的金额上限 {_usd(max_cost_usd)}"
|
||||
if reason == "daily_max_usd":
|
||||
return f"全局每日核查金额上限 {_usd(daily_max_usd)}"
|
||||
return f"最多 {steps} 轮"
|
||||
|
||||
|
||||
def _text_of(messages: list[BaseMessage]) -> str:
|
||||
return "\n".join(str(m.content) if isinstance(m.content, str) else json.dumps(m.content, ensure_ascii=False)
|
||||
for m in messages)
|
||||
|
||||
|
||||
def _final_messages(system: str, task: str, queries: list[dict[str, Any]], said: str, limit_words: str | None
|
||||
) -> list[BaseMessage]:
|
||||
"""收尾那次调用的消息：任务、查到的结果（每次查询列出行号）、模型最后说的话，以及交结果的规则。
|
||||
limit_words：触顶时是哪个上限（「最多查询 6 次」），告诉模型为什么不能再查。
|
||||
|
||||
长度有上限（任务和规则之外，查询结果按 _digest 的上限、说明按 _SAID_CHARS）：循环里按这个上限给收尾留钱。
|
||||
"""
|
||||
parts = [f"任务：\n{task.strip()}"]
|
||||
if limit_words:
|
||||
parts.append(f"已达上限（{limit_words}），不能再查询，按已经查到的结果交结果。")
|
||||
parts.append(f"你查到的结果：\n{_digest(queries) if queries else '（没有成功的查询）'}")
|
||||
parts.append(f"你最后的说明（只作参考，数以查询结果为准）：\n{said.strip()[:_SAID_CHARS] or '（没有）'}")
|
||||
parts.append(VERIFY_FINAL_PROMPT)
|
||||
return [SystemMessage(content=system), HumanMessage(content="\n\n".join(parts))]
|
||||
|
||||
|
||||
def _plain_messages(final: list[BaseMessage], schema: dict[str, Any]) -> list[BaseMessage]:
|
||||
"""普通调用（不走结构化输出）的消息：结构不在请求里了，要写进消息，不然模型不知道要交哪些字段、字段叫什么，
|
||||
交回来的字段对不上，全成了没交。"""
|
||||
*head, last = final
|
||||
text = (f"{last.content}\n\n只交一个 JSON 对象，不要写别的内容。对象的结构（JSON Schema）：\n"
|
||||
f"```json\n{json.dumps(schema, ensure_ascii=False)}\n```")
|
||||
return [*head, HumanMessage(content=text)]
|
||||
|
||||
|
||||
def _parsed(value: Any) -> dict[str, Any] | None:
|
||||
"""结构化结果：对象；有的模型把整个对象写成一段 JSON 文本，也认。"""
|
||||
if isinstance(value, str):
|
||||
value = _json_in(value)
|
||||
return value if isinstance(value, dict) else None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _State:
|
||||
"""一次核查过程中攒下的东西：出意外时也照这些交回，用了多少、查了什么不丢。"""
|
||||
|
||||
started: float
|
||||
usage: dict[str, Any] = field(default_factory=lambda: {
|
||||
"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cost_usd": 0.0, "calls": 0, "unpriced_calls": 0})
|
||||
notes: list[str] = field(default_factory=list)
|
||||
queries: list[dict[str, Any]] = field(default_factory=list)
|
||||
ledger: list[dict[str, Any]] = field(default_factory=list)
|
||||
model: str | None = None
|
||||
|
||||
def outcome(self, **fields: Any) -> VerifyOutcome:
|
||||
usage = self.usage
|
||||
usage["total_tokens"] = usage["input_tokens"] + usage["output_tokens"]
|
||||
usage["cost_usd"] = round(usage["cost_usd"], 6)
|
||||
shown = [{k: q[k] for k in ("alias", "artifact", "sql", "row_count", "truncated", "source", "tool", "call_id")}
|
||||
for q in self.queries]
|
||||
base: dict[str, Any] = {"result": None, "cells": [], "queries": shown, "rationale": None, "usage": usage,
|
||||
"notes": self.notes, "model": self.model,
|
||||
"duration_ms": int((time.monotonic() - self.started) * 1000)}
|
||||
return VerifyOutcome(**{**base, **fields})
|
||||
|
||||
|
||||
async def run_verify_agent(
|
||||
*,
|
||||
spec: ModelSpec,
|
||||
system: str,
|
||||
task: str,
|
||||
sources: Iterable[str],
|
||||
output_schema: dict[str, Any],
|
||||
max_queries: int | None = 6,
|
||||
timeout_s: float | None = 120,
|
||||
max_cost_usd: float | None = None,
|
||||
daily_max_usd: float | None | Default = Default.SETTINGS,
|
||||
max_steps: int | None = None,
|
||||
emit: Emit | None = None,
|
||||
run_id: str = "",
|
||||
node_id: str = "",
|
||||
data_versions: dict[str, Any] | Versions | None = Versions.UNSET,
|
||||
) -> VerifyOutcome:
|
||||
"""让一个只读 agent 查库，按 output_schema 交回有出处的结果。永远不抛。
|
||||
|
||||
- spec：用哪个模型（一般是 api/settings.py 的 verify_model_spec）。system / task：交给模型的角色说明和任务。
|
||||
- sources：可以查的数据源名；不存在、停用了的照实记进 notes。
|
||||
- output_schema：调用方要的结果结构（JSON Schema 对象，可以带 $defs / $ref）。数值字段会换成格子引用交给模型、
|
||||
由代码取值。
|
||||
- 上限：max_queries（db_query 的执行次数，被拒的语句也算）、timeout_s（整个循环的时长，收尾另给宽限）、
|
||||
max_cost_usd（这一次）、daily_max_usd（全局每日，和裁判分开记账；缺省按设置里的「核查每日金额上限」）；
|
||||
None 表示不限。max_steps 是模型最多想几轮，缺省按查询次数推。
|
||||
- emit：节点里调用时传 NodeContext.emit，每次模型调用发 llm.start / llm.end、每次工具调用发 tool.start /
|
||||
tool.end（role: verify，查询带 query_artifact）；不在节点里（质询接口）就不传，什么都不发。
|
||||
- run_id / node_id：查询快照、表结构快照记在谁名下。data_versions：上传表格用哪一版（运行固定的那份
|
||||
data_versions；不在运行里传 Versions.CURRENT；缺省 UNSET 时上传源一律报错，不拿当前版本顶上）。
|
||||
"""
|
||||
state = _State(started=time.monotonic())
|
||||
try:
|
||||
return await _verify(state, spec=spec, system=system, task=task, sources=sources,
|
||||
output_schema=output_schema, max_queries=max_queries, timeout_s=timeout_s,
|
||||
max_cost_usd=max_cost_usd, daily_max_usd=daily_max_usd, max_steps=max_steps, emit=emit,
|
||||
run_id=run_id, node_id=node_id, data_versions=data_versions)
|
||||
except Exception as e: # noqa: BLE001 - 永远不抛：库写不进去、建工具出错这类意外也落在返回值里
|
||||
logger.warning("核查出错", exc_info=True)
|
||||
return state.outcome(error=f"核查出错（{describe_exception(e)}）")
|
||||
|
||||
|
||||
async def _verify(state: _State, *, spec: ModelSpec, system: str, task: str, sources: Iterable[str],
|
||||
output_schema: dict[str, Any], max_queries: int | None, timeout_s: float | None,
|
||||
max_cost_usd: float | None, daily_max_usd: float | None | Default, max_steps: int | None,
|
||||
emit: Emit | None, run_id: str, node_id: str,
|
||||
data_versions: dict[str, Any] | Versions | None) -> VerifyOutcome:
|
||||
deadline = state.started + timeout_s if timeout_s is not None else None
|
||||
usage, notes, queries, ledger = state.usage, state.notes, state.queries, state.ledger
|
||||
outcome = state.outcome
|
||||
steps = max_steps if max_steps is not None else max(4, (max_queries if max_queries is not None else 6) * 2 + 2)
|
||||
owner = node_id or "verify"
|
||||
|
||||
def left() -> float | None:
|
||||
return None if deadline is None else deadline - time.monotonic()
|
||||
|
||||
# ---- 模型
|
||||
try:
|
||||
async with SessionLocal() as session:
|
||||
base_model, model_id = await get_chat_model(session, spec)
|
||||
except Exception as e: # noqa: BLE001 - 核查模型没配好：这回核不了，照实交回
|
||||
why = not_configured(e) if isinstance(e, ProviderNotConfigured) else describe_exception(e)
|
||||
return outcome(error=f"核查模型不可用（{why}）")
|
||||
state.model = model_id
|
||||
if daily_max_usd is Default.SETTINGS:
|
||||
from app.api.settings import verify_settings # 在这里导入：api.settings 导入了 engine 里的模块
|
||||
|
||||
async with SessionLocal() as session:
|
||||
daily_max_usd = (await verify_settings(session))["daily_max_usd"]
|
||||
budget = _Budget(model_id, max_cost_usd, daily_max_usd)
|
||||
if not budget.priced and (max_cost_usd is not None or daily_max_usd is not None):
|
||||
note = (f"模型「{model_id}」不在价格目录中，无法按 token 数估算金额：金额上限（这次核查、每日）对其无效，"
|
||||
"费用只受查询次数和时长上限约束")
|
||||
notes.append(note)
|
||||
await _announce(emit, EventType.LOG, level="warn", code="verify_unpriced", message=note)
|
||||
|
||||
# ---- 工具：只读的 db_schema / db_query
|
||||
wanted = list(dict.fromkeys(s for s in sources if isinstance(s, str) and s))
|
||||
ctx = ToolContext(run_id=run_id, node_id=node_id, data_versions=data_versions, emit=emit)
|
||||
async with SessionLocal() as session:
|
||||
tools = await build_datasource_tools([f"{p}{s}" for s in wanted for p in (SCHEMA_PREFIX, QUERY_PREFIX)], ctx,
|
||||
session, force_readonly=True)
|
||||
tool_map = {t.name: t for t in tools}
|
||||
if lost := [s for s in wanted if f"{QUERY_PREFIX}{s}" not in tool_map]:
|
||||
notes.append("数据源" + "、".join(f"「{s}」" for s in lost) + "不存在或已停用，核查没有查它")
|
||||
model = bind_tools_safely(base_model, tools, parallel=False)
|
||||
limit_note = f"查询次数有上限，本次最多 {max_queries} 次，用完就不能再查。" if max_queries is not None else ""
|
||||
head = "\n\n".join(p for p in (system.strip(), VERIFY_SYSTEM.format(limit=limit_note)) if p)
|
||||
messages: list[BaseMessage] = [SystemMessage(content=head), HumanMessage(content=task)]
|
||||
tool_text = "\n".join(f"{t.name}：{t.description}" for t in tools)
|
||||
resolved = _inline(output_schema)
|
||||
schema = result_schema(resolved)
|
||||
# 结构也算输入：结构化输出时它是工具定义，退回普通调用时写在消息里
|
||||
schema_text = json.dumps(schema, ensure_ascii=False)
|
||||
# 收尾那次最贵的样子：查询结果、最后的说明都按上限、按一个字一个 token 算。查询结果要到查完才知道多长，
|
||||
# 按眼下的消息估，等最后一条查询的结果进来就不够了（大段中文的结果尤其如此）
|
||||
keep = budget.estimate(_text_of(_final_messages(head, task, [], "", "上" * _LIMIT_ROOM)) + schema_text,
|
||||
_OUT_FINAL, room=_DIGEST_ROOM + _SAID_CHARS)
|
||||
|
||||
async def ask(payload: list[BaseMessage], runner: Any, estimate: float, day: str, purpose: str,
|
||||
seconds: float | None) -> tuple[Any, BaseMessage | None, BaseException | None]:
|
||||
"""调一次模型并记账。返回 (结果, 原始回复, 异常)；超时的异常是 asyncio.TimeoutError。"""
|
||||
await _announce(emit, EventType.LLM_START, model=model_id, purpose=purpose, role="verify",
|
||||
message_count=len(payload))
|
||||
began = time.perf_counter()
|
||||
got: Any = None
|
||||
raw: BaseMessage | None = None
|
||||
failure: BaseException | None = None
|
||||
try:
|
||||
got = await _bounded(runner(payload), seconds)
|
||||
raw = got if isinstance(got, BaseMessage) else _split_structured(got)[0]
|
||||
except (asyncio.TimeoutError, Exception) as e: # noqa: BLE001 - 调用失败落在返回值里
|
||||
failure = e
|
||||
inp, out = _usage(raw)
|
||||
cost = await budget.settle(day, estimate, inp, out)
|
||||
usage["input_tokens"] += inp
|
||||
usage["output_tokens"] += out
|
||||
usage["cost_usd"] += cost
|
||||
usage["calls"] += 1
|
||||
usage["unpriced_calls"] += 0 if budget.priced else 1
|
||||
end = {"model": model_id, "purpose": purpose, "role": "verify",
|
||||
"duration_ms": int((time.perf_counter() - began) * 1000), "input_tokens": inp, "output_tokens": out,
|
||||
"total_tokens": inp + out, "cost_usd": round(cost, 6), "calls": 1}
|
||||
if failure is not None:
|
||||
end["error"] = "超过核查的时长上限，已放弃这次调用" if isinstance(failure, asyncio.TimeoutError) \
|
||||
else explain_model_error(failure, model_id)
|
||||
await _announce(emit, EventType.LLM_END, **end)
|
||||
return got, raw, failure
|
||||
|
||||
# ---- 查库的循环
|
||||
stopped: str | None = None
|
||||
answered = False
|
||||
kept = False # 有一轮按收尾最贵的样子留过钱：收尾不再按金额上限拦
|
||||
said = ""
|
||||
attempts = 0
|
||||
done: dict[str, str] = {} # 成功过的查询（工具名 + 参数）→ 编号：一模一样的不再执行
|
||||
for _ in range(steps):
|
||||
remaining = left()
|
||||
if remaining is not None and remaining <= 0:
|
||||
stopped = "timeout_s"
|
||||
break
|
||||
estimate = budget.estimate(_text_of(messages) + tool_text, _OUT_STEP)
|
||||
day, refused = await budget.reserve(estimate, keep=keep)
|
||||
if refused:
|
||||
stopped = refused
|
||||
break
|
||||
kept = True
|
||||
response, _, failure = await ask(messages, model.ainvoke, estimate, day or "", "verify", remaining)
|
||||
if isinstance(failure, asyncio.TimeoutError):
|
||||
stopped = "timeout_s"
|
||||
break
|
||||
if failure is not None or not isinstance(response, BaseMessage):
|
||||
return outcome(error=explain_model_error(failure, model_id) if failure else "核查模型没有给出回复")
|
||||
messages.append(response)
|
||||
calls = getattr(response, "tool_calls", None) or []
|
||||
if not calls:
|
||||
answered, said = True, message_text(response)
|
||||
break
|
||||
said = message_text(response) or said
|
||||
for call in calls:
|
||||
name, args = str(call.get("name") or ""), dict(call.get("args") or {})
|
||||
call_id = str(call.get("id") or f"verify-{len(messages)}-{name}")
|
||||
if stopped:
|
||||
messages.append(ToolMessage(content=f"未执行：已达上限（{stopped}），不能再调用工具。", tool_call_id=call_id))
|
||||
continue
|
||||
tool = tool_map.get(name)
|
||||
if tool is None:
|
||||
messages.append(ToolMessage(content=f"错误：没有名为 {name} 的工具，可用的工具："
|
||||
+ ("、".join(tool_map) or "（无）"), tool_call_id=call_id))
|
||||
continue
|
||||
args_schema = args_model_of(tool)
|
||||
if args_schema is not None:
|
||||
try:
|
||||
args, _ = prepare_args(args_schema, args)
|
||||
except ToolArgsError as e:
|
||||
await _announce(emit, EventType.TOOL_ERROR, tool=name, call_id=call_id, error=str(e), role="verify")
|
||||
messages.append(ToolMessage(content=str(e), tool_call_id=call_id))
|
||||
continue
|
||||
is_query = name.startswith(QUERY_PREFIX)
|
||||
fingerprint = f"{name}\n{json.dumps(args, ensure_ascii=False, sort_keys=True, default=str)}"
|
||||
if is_query and fingerprint in done:
|
||||
alias = done[fingerprint]
|
||||
messages.append(ToolMessage(content=f"这条查询和 {alias} 完全相同，结果见 {alias}，没有再执行。",
|
||||
tool_call_id=call_id))
|
||||
continue
|
||||
if is_query and max_queries is not None and attempts >= max_queries:
|
||||
stopped = "max_queries"
|
||||
messages.append(ToolMessage(content=f"未执行：已达查询次数上限（{max_queries} 次），不能再查询。",
|
||||
tool_call_id=call_id))
|
||||
continue
|
||||
remaining = left()
|
||||
if remaining is not None and remaining <= 0:
|
||||
stopped = "timeout_s"
|
||||
messages.append(ToolMessage(content="未执行：已达核查的时长上限，不能再查询。", tool_call_id=call_id))
|
||||
continue
|
||||
content = await _run_tool(tool, name, args, call_id, remaining, emit, queries, ledger, owner)
|
||||
if content is None:
|
||||
stopped = "timeout_s"
|
||||
content = "查询因核查的时长上限被中断，没有结果。"
|
||||
if is_query:
|
||||
attempts += 1
|
||||
if queries and queries[-1]["call_id"] == call_id:
|
||||
done[fingerprint] = queries[-1]["alias"]
|
||||
messages.append(ToolMessage(content=content, tool_call_id=call_id))
|
||||
if not stopped and max_queries is not None and attempts >= max_queries:
|
||||
stopped = "max_queries"
|
||||
if stopped:
|
||||
break
|
||||
else:
|
||||
stopped = "max_steps"
|
||||
|
||||
def words_of(reason: str) -> str:
|
||||
return _limit_words(reason, max_queries=max_queries, timeout_s=timeout_s, max_cost_usd=max_cost_usd,
|
||||
daily_max_usd=daily_max_usd, steps=steps)
|
||||
|
||||
words = words_of(stopped) if stopped else None
|
||||
if words:
|
||||
await _announce(emit, EventType.LOG, level="warn", code="verify_limit", limit=stopped,
|
||||
message=f"核查已达上限（{words}），已停止查询，按已有结果收尾")
|
||||
|
||||
# ---- 收尾：按结构交结果。关掉 thinking：结构化输出靠强制工具调用，和思考模式不兼容
|
||||
try:
|
||||
async with SessionLocal() as session:
|
||||
settler, _ = await get_chat_model(session, spec.model_copy(update={"thinking": "off"}))
|
||||
except Exception as e: # noqa: BLE001
|
||||
return outcome(limit=stopped, answered=answered, error=f"核查模型不可用（{describe_exception(e)}）")
|
||||
final = _final_messages(head, task, queries, said, words)
|
||||
try:
|
||||
structured = settler.with_structured_output(schema, include_raw=True)
|
||||
except Exception: # noqa: BLE001 - 模型不支持结构化输出：直接走普通调用，不白算一次调用
|
||||
structured = None
|
||||
remaining = left()
|
||||
grace = None if remaining is None else max(remaining, SETTLE_GRACE_S)
|
||||
finish = None if grace is None else time.monotonic() + grace
|
||||
|
||||
def until_finish() -> float | None:
|
||||
return None if finish is None else finish - time.monotonic()
|
||||
|
||||
parsed: dict[str, Any] | None = None
|
||||
failure: BaseException | None = None
|
||||
if structured is not None:
|
||||
estimate = budget.estimate(_text_of(final) + schema_text, _OUT_FINAL)
|
||||
day, refused = await budget.reserve(estimate, kept=kept)
|
||||
if refused:
|
||||
return outcome(limit=stopped or refused, answered=answered, error=f"已达{words_of(refused)}，没能交结果")
|
||||
got, raw, failure = await ask(final, structured.ainvoke, estimate, day or "", "verify_result", until_finish())
|
||||
parsed = None if failure is not None else _parsed(_split_structured(got)[1])
|
||||
if parsed is None and raw is not None:
|
||||
# 有的网关把 JSON 写在正文里、或者塞进了一个没按结构解析出来的工具调用
|
||||
for candidate in [message_text(raw), *(c.get("args") for c in getattr(raw, "tool_calls", None) or [])]:
|
||||
if (parsed := _parsed(candidate)) is not None:
|
||||
break
|
||||
if parsed is None and not isinstance(failure, asyncio.TimeoutError):
|
||||
# 不支持结构化输出的模型：普通调用一次，结构写进消息，从回复里把 JSON 抠出来。结构化那次已经做过的，
|
||||
# 这一次是补救，照常按上限拦；没做过的，这一次就是收尾，循环里留过钱
|
||||
plain = _plain_messages(final, schema)
|
||||
remaining = until_finish()
|
||||
if remaining is None or remaining > 0:
|
||||
estimate = budget.estimate(_text_of(plain), _OUT_FINAL)
|
||||
day, refused = await budget.reserve(estimate, kept=kept and structured is None)
|
||||
if refused and structured is None:
|
||||
return outcome(limit=stopped or refused, answered=answered,
|
||||
error=f"已达{words_of(refused)}，没能交结果")
|
||||
if not refused:
|
||||
got, _, failure = await ask(plain, settler.ainvoke, estimate, day or "", "verify_result", remaining)
|
||||
parsed = None if failure is not None else _parsed(message_text(got) if got is not None else None)
|
||||
if parsed is None:
|
||||
why = "收尾时超过时长上限，没能交结果" if isinstance(failure, asyncio.TimeoutError) else \
|
||||
f"核查模型没有按结构交结果（{explain_model_error(failure, model_id)}）" if failure is not None else \
|
||||
"核查模型没有按结构交结果"
|
||||
return outcome(limit=stopped, answered=answered, error=why)
|
||||
declared = list((resolved.get("properties") or {}) if _is_object(resolved) else {})
|
||||
if declared and not any(k in parsed for k in declared):
|
||||
# 一个要求的字段都没交：不是「查不到」，是没按结构交。不这么分，字段名对不上的回复会被当成什么都没查到
|
||||
return outcome(limit=stopped, answered=answered,
|
||||
error="核查模型没有按结构交结果（交回的内容里没有要求的字段："
|
||||
+ "、".join(declared[:8]) + ("等" if len(declared) > 8 else "") + "）")
|
||||
|
||||
# ---- 取值、组装理由
|
||||
cells = _Cells(queries)
|
||||
result = cells.walk(resolved, parsed) if _is_object(resolved) else {}
|
||||
rationale = parsed.get(RATIONALE)
|
||||
doc = None
|
||||
if isinstance(rationale, str) and rationale.strip():
|
||||
doc = compose_doc(rationale, build_catalog(nodes={}, ledger=ledger), node_id=owner, run_id=run_id or None)
|
||||
result[RATIONALE] = doc["markdown"] if doc else ""
|
||||
return outcome(result=result, cells=cells.cells, rationale=doc, limit=stopped, answered=answered)
|
||||
|
||||
|
||||
async def _run_tool(tool: Any, name: str, args: dict[str, Any], call_id: str, remaining: float | None,
|
||||
emit: Emit | None, queries: list[dict[str, Any]], ledger: list[dict[str, Any]],
|
||||
owner: str) -> str | None:
|
||||
"""执行一次工具调用，返回交给模型的内容；整个核查的时长到点、这次被中断时返回 None。
|
||||
|
||||
成功的查询编号（Q1、Q2…）、记进本 agent 自己的台账和查询清单，内容前面加「【证据 Qn】」。
|
||||
"""
|
||||
limit = limit_of(tool, name, args)
|
||||
await _announce(emit, EventType.TOOL_START, tool=name, args=args, call_id=call_id, role="verify",
|
||||
**limit_fields(limit))
|
||||
began = time.perf_counter()
|
||||
try:
|
||||
result = await _bounded(run_bounded(tool.ainvoke(args), limit, name), remaining)
|
||||
content = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False, default=str)
|
||||
except ToolTimeout as e:
|
||||
# 这条查询自己的时限到了：交回原因让模型改写，核查照常往下走。要排在下一条前面：ToolTimeout 也是
|
||||
# TimeoutError（asyncio.TimeoutError 就是它），排在后面就被当成整个核查的时长到了
|
||||
await _announce(emit, EventType.TOOL_ERROR, tool=name, call_id=call_id, role="verify", timed_out=True,
|
||||
duration_ms=int((time.perf_counter() - began) * 1000), error=str(e))
|
||||
return str(e)
|
||||
except asyncio.TimeoutError:
|
||||
await _announce(emit, EventType.TOOL_ERROR, tool=name, call_id=call_id, role="verify", timed_out=True,
|
||||
duration_ms=int((time.perf_counter() - began) * 1000),
|
||||
error="查询因核查的时长上限被中断")
|
||||
return None
|
||||
except Exception as e: # noqa: BLE001 - 工具失败交回模型，让它自己改
|
||||
message = f"工具执行失败：{describe_exception(e)}"
|
||||
await _announce(emit, EventType.TOOL_ERROR, tool=name, call_id=call_id, role="verify",
|
||||
duration_ms=int((time.perf_counter() - began) * 1000), error=message[:300])
|
||||
return message
|
||||
elapsed = int((time.perf_counter() - began) * 1000)
|
||||
fields = query_entry_fields(content) if name.startswith(QUERY_PREFIX) else None
|
||||
extra: dict[str, Any] = {}
|
||||
if fields is not None:
|
||||
alias = f"Q{len(queries) + 1}"
|
||||
data = json.loads(content)
|
||||
schema = fields.get("schema_artifact")
|
||||
seen = schema and any(e.get("kind") == "schema" and e.get("artifact") == schema for e in ledger)
|
||||
frozen = schema_ledger_entry(schema, node_id=owner, exec_no=1) if schema and not seen else None
|
||||
shaped = {"schema_artifact": schema, "tables": fields["tables"]} if seen or frozen else {}
|
||||
ledger.append({"kind": "query", "node_id": owner, "exec": 1, "artifact": fields["artifact"],
|
||||
"call_id": call_id, "tool": name, "source": fields["source"], "columns": fields["columns"],
|
||||
"rows": fields["rows"], "truncated": fields["truncated"], **shaped})
|
||||
if frozen:
|
||||
ledger.append(frozen)
|
||||
queries.append({"alias": alias, "call_id": call_id, "tool": name, "source": fields["source"],
|
||||
"sql": str(data.get("sql") or args.get("sql") or ""), "artifact": fields["artifact"],
|
||||
"columns": fields["columns"], "rows": data.get("rows") or [], "row_count": fields["rows"],
|
||||
"truncated": fields["truncated"]})
|
||||
extra = {"query_artifact": fields["artifact"], "alias": alias,
|
||||
**({"schema_artifact": schema} if schema else {}),
|
||||
**({"checks": data["checks"]} if isinstance(data.get("checks"), list) and data["checks"] else {})}
|
||||
content = f"【证据 {alias}】\n{content}"
|
||||
await _announce(emit, EventType.TOOL_END, tool=name, call_id=call_id, duration_ms=elapsed, preview=content[:2000],
|
||||
role="verify", **extra)
|
||||
if len(content) > _TOOL_CHARS:
|
||||
content = content[:_TOOL_CHARS] + "\n…（结果太长，后面的省略了；完整结果在查询快照里，交结果时会按编号列出）"
|
||||
return content
|
||||
|
||||
|
||||
__all__ = ["LIMITS", "RATIONALE", "SETTLE_GRACE_S", "Default", "VerifyOutcome", "result_schema", "run_verify_agent"]
|
||||
@@ -35,7 +35,7 @@ from app.data import catalog as data_catalog
|
||||
from app.data import introspect, table_versions
|
||||
from app.data.sqlcheck import SqlChecker
|
||||
from app.data.engine import SnapshotTampered, masked_columns, query_timeout, run_query
|
||||
from app.data.guard import QueryLimits, SqlRejected, is_write
|
||||
from app.data.guard import QueryLimits, SqlRejected, check, is_write
|
||||
from app.data.tabular import UNSHAPED_NOTE
|
||||
from app.db.models import DataSource
|
||||
from app.tools.registry import ToolContext, Versions
|
||||
@@ -57,6 +57,10 @@ LATE_PIN_CODE = "data_version_pinned"
|
||||
#: 每次查询之后基于数据目录的 SQL 检查最多等这么久。正常一条 SQL 几毫秒；等不到就放弃这一次检查，
|
||||
#: 查询结果照常交回——检查是附加的，不能拖慢查询
|
||||
SQL_CHECK_TIMEOUT_S = 2.0
|
||||
#: force_readonly 的工具拒掉写语句时交回的原因（接在「SQL 被拒绝：」后面）。不说「数据源是只读的」：源本身
|
||||
#: 可能是可写的，是这次调用只许读；也不提「另建一个可写的数据源」——核查用不着写
|
||||
FORCED_READONLY = ("这次查询只允许读取数据，这条语句会修改数据，不被允许。"
|
||||
"只能执行不改数据的 SELECT / WITH / SHOW / DESCRIBE / EXPLAIN")
|
||||
|
||||
|
||||
class _QueryArgs(BaseModel):
|
||||
@@ -165,12 +169,13 @@ _FROZEN_SCHEMA_KEYS = ("schema", "synced_at", "truncated", "total", "import_mode
|
||||
CATALOG_SNAPSHOT_KEY = "catalog"
|
||||
|
||||
|
||||
def _query_description(source: Any) -> str:
|
||||
def _query_description(source: Any, *, readonly: bool | None = None) -> str:
|
||||
"""readonly 不给就按源本身；force_readonly 的工具传 True（源可写，这次调用也只许读）。"""
|
||||
tables = introspect.table_names(source)
|
||||
head = f"在数据源「{source.name}」上执行 SQL 查询。"
|
||||
if source.description:
|
||||
head += f"{source.description}。"
|
||||
head += "只读。" if source.readonly else "可写（写操作需要人工审批）。"
|
||||
head += "只读。" if (source.readonly if readonly is None else readonly) else "可写（写操作需要人工审批）。"
|
||||
if tables:
|
||||
listed = "、".join(tables[:25])
|
||||
more = f" 等 {len(tables)} 张表" if len(tables) > 25 else ""
|
||||
@@ -194,13 +199,16 @@ def _query_description(source: Any) -> str:
|
||||
|
||||
|
||||
async def build_datasource_tools(
|
||||
names: list[str], ctx: ToolContext, session: AsyncSession
|
||||
names: list[str], ctx: ToolContext, session: AsyncSession, *, force_readonly: bool = False
|
||||
) -> list[StructuredTool]:
|
||||
"""按工具名反查数据源并构造工具。build_tools 的第四个来源。
|
||||
|
||||
上传源先按 ctx.data_versions 解析成绑定快照的视图（见 _resolve），工具描述、表结构、
|
||||
查询都按那个版本来。固定的版本找不到时建一个同名的替身（_unavailable_tool），
|
||||
每次调用都说明原因，不拿当前版本顶上。
|
||||
|
||||
force_readonly：查询工具一律按只读守卫检查，可写源上的写语句也拒绝（见 _make_query_tool）。
|
||||
核查（engine/verify_agent.py）建工具时传它：那一路没有审批关卡，绝不能写。
|
||||
"""
|
||||
wanted: dict[str, list[str]] = {}
|
||||
for name in names:
|
||||
@@ -230,7 +238,7 @@ async def build_datasource_tools(
|
||||
entries = await _catalog_of(row)
|
||||
for prefix in prefixes:
|
||||
if prefix == QUERY_PREFIX:
|
||||
tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries))
|
||||
tools.append(_make_query_tool(source, ctx, fixed=fixed, catalog=entries, force_readonly=force_readonly))
|
||||
else:
|
||||
tools.append(_make_schema_tool(source, fixed=fixed, catalog=entries))
|
||||
return tools
|
||||
@@ -441,11 +449,15 @@ def _unavailable_tool(row: DataSource, prefix: str, reason: str) -> StructuredTo
|
||||
|
||||
|
||||
def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
|
||||
catalog: Mapping[str, data_catalog.CatalogEntry] | None = None) -> StructuredTool:
|
||||
catalog: Mapping[str, data_catalog.CatalogEntry] | None = None,
|
||||
force_readonly: bool = False) -> StructuredTool:
|
||||
"""source 是手工源的 DataSource，或上传源绑定快照的 SourceView（字段同名）。
|
||||
|
||||
fixed：上传源的版本是这次运行固定的（找不到时的说法不同）。
|
||||
catalog：建工具时读到的数据目录，随表结构快照冻结（_store_schema）。
|
||||
force_readonly：不管源可不可写，一律按只读守卫检查，写语句交回 FORCED_READONLY、不执行。只靠守卫这一道：
|
||||
连接池按源缓存、不分只读与否（engine.cache_key），可写源借不到连接层只读的连接——别为此另起一个只读池，
|
||||
缓存键相同，会把这个源平时用的可写连接也换成只读的。
|
||||
"""
|
||||
# 数据源在 options.query_timeout_s 里配了就用它，否则取缺省。数据库按它停下语句，
|
||||
# 引擎按它告诉界面上限是多少（metadata["timeout_s"]）
|
||||
@@ -463,6 +475,15 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
|
||||
async def _run(sql: str, limit: int | None = None) -> str:
|
||||
if gone := _vanished(source, fixed):
|
||||
return f"{DATA_UNAVAILABLE}{gone}"
|
||||
if force_readonly and not source.readonly:
|
||||
# 只读源由 run_query 自己按只读检查（原因照旧说「数据源是只读的」）；可写源在这里先按只读判一遍。
|
||||
# 先过一遍和 run_query 相同的检查：多条语句、结构变更这类问题照原来的说法交回
|
||||
try:
|
||||
statement = check(sql, readonly=False, source_name=source.name, dialect=(source.kind or "").lower())
|
||||
except SqlRejected as e:
|
||||
return f"SQL 被拒绝：{e}"
|
||||
if is_write(statement, dialect=(source.kind or "").lower()):
|
||||
return f"SQL 被拒绝：{FORCED_READONLY}"
|
||||
limits = QueryLimits(timeout_seconds=seconds, **({"max_rows": int(limit)} if limit else {}))
|
||||
try:
|
||||
result = await run_query(source, sql, limits=limits)
|
||||
@@ -520,12 +541,14 @@ def _make_query_tool(source: Any, ctx: ToolContext, *, fixed: bool = False,
|
||||
|
||||
return StructuredTool(
|
||||
name=f"{QUERY_PREFIX}{source.name}",
|
||||
description=_query_description(source),
|
||||
description=_query_description(source, readonly=True if force_readonly else None),
|
||||
args_schema=_QueryArgs,
|
||||
coroutine=_run,
|
||||
func=None,
|
||||
# 同一个工具，危不危险看这一次的 SQL。审批关卡经 registry.call_is_dangerous 问它
|
||||
metadata={"dangerous_if": lambda args: is_dangerous_call(source, str(args.get("sql") or "")),
|
||||
# 同一个工具，危不危险看这一次的 SQL。审批关卡经 registry.call_is_dangerous 问它。
|
||||
# force_readonly 的工具写不了，没有要人批的调用
|
||||
metadata={"dangerous_if": (lambda args: False) if force_readonly
|
||||
else (lambda args: is_dangerous_call(source, str(args.get("sql") or ""))),
|
||||
"timeout_s": seconds},
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
"""数据源查询工具的 force_readonly：核查（双盲复算、反方审稿、审批质询）只许读数据。
|
||||
|
||||
核查的 agent 是另起的一路，没有审批关卡兜底；它要查的数据源却可能是可写的。所以建工具时传 force_readonly=True：
|
||||
不管源本身可不可写，一律按只读守卫检查——可写源上的写语句也拒绝，原因交回模型（它据此改写），数据原样不动；
|
||||
查询照常。工具描述跟着写「只读」，审批判定也不再把写语句当成要人批的危险调用（反正执行不了）。
|
||||
不传的时候一切照旧：可写源上的写语句仍是要审批的危险调用。
|
||||
|
||||
夹具是一个现造的可写 SQLite 文件，表和数字都是虚构的。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from app.data import introspect
|
||||
from app.data.engine import engines
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import DataSource
|
||||
from app.tools.datasource import build_datasource_tools
|
||||
from app.tools.registry import ToolContext, Versions
|
||||
|
||||
|
||||
def rows_in(path: str) -> list[tuple]:
|
||||
db = sqlite3.connect(path)
|
||||
try:
|
||||
return db.execute("SELECT id, amount FROM orders ORDER BY id").fetchall()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def writable(tmp_path):
|
||||
"""一个可写的手工源：orders 两行，金额 120 和 80。"""
|
||||
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
||||
db = sqlite3.connect(path)
|
||||
db.executescript("CREATE TABLE orders (id INTEGER PRIMARY KEY, amount INTEGER);"
|
||||
"INSERT INTO orders VALUES (1, 120), (2, 80);")
|
||||
db.commit()
|
||||
db.close()
|
||||
name = f"shop_{uuid.uuid4().hex[:8]}"
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True)
|
||||
session.add(row)
|
||||
await session.flush()
|
||||
row.schema_cache = await introspect.introspect(row)
|
||||
await session.commit()
|
||||
source_id = row.id
|
||||
yield name, str(path)
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
async def tools(name: str, **kwargs) -> dict:
|
||||
async with SessionLocal() as session:
|
||||
built = await build_datasource_tools([f"db_query__{name}", f"db_schema__{name}"],
|
||||
ToolContext(data_versions=Versions.CURRENT), session, **kwargs)
|
||||
return {t.name: t for t in built}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", [
|
||||
"DELETE FROM orders",
|
||||
"UPDATE orders SET amount = 0",
|
||||
"INSERT INTO orders VALUES (3, 1)",
|
||||
"WITH k AS (SELECT 1) DELETE FROM orders WHERE id = 1",
|
||||
])
|
||||
async def test_writes_are_refused_on_a_writable_source(writable, sql):
|
||||
name, path = writable
|
||||
query = (await tools(name, force_readonly=True))[f"db_query__{name}"]
|
||||
got = await query.ainvoke({"sql": sql})
|
||||
assert got.startswith("SQL 被拒绝：") and "只允许读取数据" in got, got
|
||||
assert rows_in(path) == [(1, 120), (2, 80)], "数据原样不动"
|
||||
|
||||
|
||||
async def test_structural_changes_are_refused_with_the_usual_reason(writable):
|
||||
name, path = writable
|
||||
query = (await tools(name, force_readonly=True))[f"db_query__{name}"]
|
||||
got = await query.ainvoke({"sql": "DROP TABLE orders"})
|
||||
assert got.startswith("SQL 被拒绝：") and "结构变更" in got
|
||||
assert rows_in(path) == [(1, 120), (2, 80)]
|
||||
|
||||
|
||||
async def test_queries_still_work_and_land_as_snapshots(writable):
|
||||
name, _ = writable
|
||||
built = await tools(name, force_readonly=True)
|
||||
got = json.loads(await built[f"db_query__{name}"].ainvoke({"sql": "SELECT SUM(amount) AS total FROM orders"}))
|
||||
assert got["rows"] == [[200]] and got["artifact"], "查询照常，结果照常落成查询快照"
|
||||
assert "orders" in await built[f"db_schema__{name}"].ainvoke({})
|
||||
|
||||
|
||||
async def test_the_tool_says_read_only_and_needs_no_approval(writable):
|
||||
name, _ = writable
|
||||
forced = (await tools(name, force_readonly=True))[f"db_query__{name}"]
|
||||
assert "只读。" in forced.description and "可写" not in forced.description
|
||||
assert forced.metadata["dangerous_if"]({"sql": "DELETE FROM orders"}) is False, "写语句反正执行不了，不用找人批"
|
||||
|
||||
|
||||
async def test_without_the_flag_nothing_changes(writable):
|
||||
name, path = writable
|
||||
normal = (await tools(name))[f"db_query__{name}"]
|
||||
assert "可写（写操作需要人工审批）" in normal.description
|
||||
assert normal.metadata["dangerous_if"]({"sql": "DELETE FROM orders"}) is True
|
||||
got = json.loads(await normal.ainvoke({"sql": "UPDATE orders SET amount = 90 WHERE id = 2"}))
|
||||
assert got["rows"] == [[1]], "不开 force_readonly 时可写源照常能写（审批由节点把关）"
|
||||
assert rows_in(path) == [(1, 120), (2, 90)]
|
||||
@@ -0,0 +1,714 @@
|
||||
"""SQL 语义（data/sql_semantics.py）：两条 SQL 算的是不是一回事，不是的话差在哪。
|
||||
|
||||
- normalize_sql：字面量换成占位符后的规范文本 + 字面量。只有字面量不同的两条 SQL 规范文本相同（定时运行每期的日期
|
||||
不同，不能算模板变了）；关键字大小写、空白、注释也不算不同；解析不了的退回逐词规范化，照样认得出只改了字面量
|
||||
- facets：表、连接（类型 / 关联进来的表 / ON 条件）、WHERE 合取式、时间条件（列 / 运算符 / 取值）、分组、聚合
|
||||
（函数 / 参数 / 去重）、HAVING、取哪几行（各段的 LIMIT / OFFSET / 排序）。别名换成真实表名（条件里套着的子查询、
|
||||
HAVING 里的别名也是），CTE、子查询按位置编号、不按起的名字比，条件里的子查询选出来的列名不算；比较的两边排好
|
||||
顺序；解析不了带 parse_error，不抛
|
||||
- diff_facets：每处不同一句中文：「A 按支付时间 paid_at 筛时间，B 按下单时间 created_at」「A 排除了
|
||||
is_deleted = 1 的行，B 没有」「A 对 order_id 去重计数，B 没去重」……写法不同、口径相同的不报，一边多出一个排在
|
||||
前面的 CTE、编号错开了也不报
|
||||
|
||||
夹具全是虚构的表：订单 orders、退款 refunds、门店 stores、用户 users、付款 payments，以及一张中文表「订单」。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from app.data.sql_semantics import (
|
||||
Facets,
|
||||
catalog_labels,
|
||||
diff_facets,
|
||||
facets,
|
||||
normalize_sql,
|
||||
)
|
||||
|
||||
DIALECTS = ("sqlite", "mysql", "postgres")
|
||||
|
||||
|
||||
def texts(a: str, b: str, dialect: str = "sqlite", **kwargs) -> list[str]:
|
||||
return [d.text for d in diff_facets(facets(a, dialect), facets(b, dialect), **kwargs)]
|
||||
|
||||
|
||||
def aspects(a: str, b: str, dialect: str = "sqlite") -> list[str]:
|
||||
return [d.aspect for d in diff_facets(facets(a, dialect), facets(b, dialect))]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# normalize_sql
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_only_literals_differ_means_the_same_template():
|
||||
week1 = normalize_sql("SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-01' AND store_id = 3 LIMIT 10")
|
||||
week2 = normalize_sql("SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-08' AND store_id = 7 LIMIT 10")
|
||||
assert week1.same_shape(week2) and week1.text == week2.text
|
||||
assert week1.text == "SELECT SUM(amount) FROM orders WHERE paid_at >= ? AND store_id = ? LIMIT ?"
|
||||
# 字面量按规范文本里的先后排：语法树里 LIMIT 存在 WHERE 前面，不能因此把 10 排到第一个
|
||||
assert week1.literals == ("'2026-09-01'", "3", "10")
|
||||
assert week2.literals == ("'2026-09-08'", "7", "10")
|
||||
assert week1.parse_error is None
|
||||
|
||||
|
||||
def test_case_whitespace_and_comments_do_not_count():
|
||||
a = normalize_sql("select Amount\nfrom ORDERS -- 本周\nwhere paid_at >= '2026-09-01'", "sqlite")
|
||||
b = normalize_sql("SELECT amount FROM orders /* 上周 */ WHERE paid_at >= '2026-08-25'", "sqlite")
|
||||
assert a.text == b.text and a.literals != b.literals
|
||||
|
||||
|
||||
def test_a_structural_change_is_a_different_template():
|
||||
base = normalize_sql("SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-01'")
|
||||
for changed in ("SELECT SUM(amount) FROM orders WHERE created_at >= '2026-09-01'",
|
||||
"SELECT SUM(amount) FROM orders WHERE paid_at > '2026-09-01'",
|
||||
"SELECT AVG(amount) FROM orders WHERE paid_at >= '2026-09-01'",
|
||||
"SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-01' AND is_deleted = 0"):
|
||||
assert not normalize_sql(changed).same_shape(base), changed
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dialect,a,b", [
|
||||
("sqlite", "SELECT strftime('%Y-%m', paid_at) AS m, SUM(amount) AS \"销售额\" FROM \"订单\" "
|
||||
"WHERE \"支付日期\" BETWEEN '2026-09-01' AND '2026-09-30' GROUP BY 1",
|
||||
"SELECT strftime('%Y-%m', paid_at) AS m, SUM(amount) AS \"销售额\" FROM \"订单\" "
|
||||
"WHERE \"支付日期\" BETWEEN '2026-10-01' AND '2026-10-31' GROUP BY 1"),
|
||||
("mysql", "SELECT SUM(o.`金额`) AS `销售额` FROM `订单` o WHERE o.paid_at >= \"2026-09-01\" LIMIT 5, 10",
|
||||
"SELECT SUM(o.`金额`) AS `销售额` FROM `订单` o WHERE o.paid_at >= \"2026-09-08\" LIMIT 5, 20"),
|
||||
("postgres", "SELECT SUM(amount) AS \"销售额\" FROM orders WHERE paid_at::date >= CURRENT_DATE - INTERVAL '7 days'",
|
||||
"SELECT SUM(amount) AS \"销售额\" FROM orders WHERE paid_at::date >= CURRENT_DATE - INTERVAL '14 days'"),
|
||||
])
|
||||
def test_each_dialect_reads_its_own_syntax(dialect, a, b):
|
||||
na, nb = normalize_sql(a, dialect), normalize_sql(b, dialect)
|
||||
assert na.parse_error is None and nb.parse_error is None
|
||||
assert na.text == nb.text and na.literals != nb.literals
|
||||
# 带引号的中文别名原样留在规范文本里：别名不是字面量
|
||||
assert "销售额" in na.text
|
||||
|
||||
|
||||
def test_the_data_source_kind_is_accepted_as_the_dialect():
|
||||
assert normalize_sql("SELECT 1 FROM t WHERE a = 'x'", "postgresql").text == \
|
||||
normalize_sql("SELECT 1 FROM t WHERE a = 'y'", "postgres").text
|
||||
|
||||
|
||||
def test_unparsable_sql_still_tells_a_literal_only_change():
|
||||
a = normalize_sql("SELECT amount FROM orders WHERE paid_at >= '2026-09-01' AND AND", "sqlite")
|
||||
b = normalize_sql("select amount from ORDERS where paid_at >= '2026-09-08' and and", "sqlite")
|
||||
assert a.parse_error and "无法解析" in a.parse_error
|
||||
assert a.text == b.text and a.literals == ("'2026-09-01'",) and b.literals == ("'2026-09-08'",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql", ["", " ", None, "SELECT 1; SELECT 2"])
|
||||
def test_nothing_to_normalize_never_raises(sql):
|
||||
got = normalize_sql(sql, "sqlite")
|
||||
assert got.parse_error
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# facets：各个方面
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tables_are_the_real_tables_without_cte_names():
|
||||
f = facets("WITH paid AS (SELECT * FROM orders WHERE status = 'paid') "
|
||||
"SELECT s.name, SUM(p.amount) FROM paid p JOIN stores s ON s.id = p.store_id GROUP BY s.name")
|
||||
assert [t.text for t in f.tables] == ["orders", "stores"] or [t.text for t in f.tables] == ["stores", "orders"]
|
||||
assert "paid" not in {t.key for t in f.tables}
|
||||
|
||||
|
||||
def test_joins_carry_kind_both_sides_and_a_normalized_condition():
|
||||
f = facets("SELECT o.id FROM orders o LEFT JOIN refunds r ON r.order_id = o.id AND r.state = 'done'")
|
||||
[join] = f.joins
|
||||
assert join.kind == "left" and join.table.text == "refunds" and join.other == ("orders",)
|
||||
# 别名换成真实表名，两边排好顺序，几个条件按 key 排序后用 AND 连起来
|
||||
assert join.on.text == "orders.id = refunds.order_id AND refunds.state = 'done'"
|
||||
same = facets("SELECT orders.id FROM orders LEFT OUTER JOIN refunds ON refunds.state = 'done' "
|
||||
"AND orders.id = refunds.order_id")
|
||||
assert same.joins[0].on.key == join.on.key and same.joins[0].kind == "left"
|
||||
|
||||
|
||||
def test_a_comma_join_is_the_same_as_an_inner_join():
|
||||
comma = facets("SELECT a.x FROM orders a, stores b WHERE a.store_id = b.id AND b.city = '虚构市'")
|
||||
inner = facets("SELECT orders.x FROM orders JOIN stores ON stores.id = orders.store_id WHERE stores.city = '虚构市'")
|
||||
assert comma.joins == inner.joins and comma.joins[0].kind == "inner"
|
||||
assert comma.joins[0].table.text == "stores" and comma.joins[0].other == ("orders",)
|
||||
cross = facets("SELECT 1 FROM orders o CROSS JOIN stores s")
|
||||
assert (cross.joins[0].kind, cross.joins[0].on, cross.joins[0].other) == ("cross", None, ("orders",))
|
||||
assert [p.key for p in comma.where] == [p.key for p in inner.where] == ["stores.city = '虚构市'"]
|
||||
|
||||
|
||||
def test_where_conjuncts_are_normalized():
|
||||
f = facets("SELECT 1 FROM orders o WHERE 0 = o.is_deleted AND o.status IN ('shipped', 'paid') "
|
||||
"AND (o.channel = 'app' OR o.channel = 'web')")
|
||||
assert [p.text for p in f.where] == ["is_deleted = 0", "status IN ('paid', 'shipped')",
|
||||
"channel = 'app' OR channel = 'web'"]
|
||||
assert f.where[0].columns == (("orders", "is_deleted"),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dialect", DIALECTS)
|
||||
def test_exclusions_are_read_in_every_dialect(dialect):
|
||||
f = facets("SELECT 1 FROM orders WHERE is_deleted <> 1 AND store_id NOT IN (8, 9) AND refund_id IS NOT NULL "
|
||||
"AND note NOT LIKE '%测试%'", dialect)
|
||||
assert [p.excludes for p in f.where] == ["is_deleted = 1", "store_id IN (8, 9)", "refund_id 为空",
|
||||
"note LIKE '%测试%'"]
|
||||
|
||||
|
||||
def test_time_predicates_have_column_operator_and_value():
|
||||
f = facets("SELECT SUM(amount) FROM orders o WHERE o.paid_at >= '2026-09-01' AND '2026-10-01' > o.paid_at "
|
||||
"AND o.store_id = 3")
|
||||
assert [(t.column.text, t.op, t.value) for t in f.time] == [("paid_at", ">=", "'2026-09-01'"),
|
||||
("paid_at", "<", "'2026-10-01'")]
|
||||
between = facets("SELECT 1 FROM \"订单\" WHERE \"支付日期\" BETWEEN '2026-09-01' AND '2026-09-30'")
|
||||
[t] = between.time
|
||||
assert (t.column.text, t.op, t.value) == ('"支付日期"', "between", "'2026-09-01' AND '2026-09-30'")
|
||||
# 列套着日期函数也认；按月份字符串筛也认
|
||||
month = facets("SELECT 1 FROM orders WHERE strftime('%Y-%m', created_at) = '2026-09'")
|
||||
assert [(t.column.text, t.op) for t in month.time] == [("created_at", "=")]
|
||||
assert facets("SELECT 1 FROM orders WHERE store_id = 3 AND amount > 100").time == ()
|
||||
|
||||
|
||||
def test_group_by_position_and_alias_mean_the_expression():
|
||||
explicit = facets("SELECT o.region, SUM(o.amount) FROM orders o GROUP BY o.region")
|
||||
by_position = facets("SELECT region AS r, SUM(amount) FROM orders GROUP BY 1")
|
||||
by_alias = facets("SELECT region AS r, SUM(amount) FROM orders GROUP BY r")
|
||||
assert [p.key for p in explicit.group_by] == [p.key for p in by_position.group_by] == \
|
||||
[p.key for p in by_alias.group_by] == ["orders.region"]
|
||||
|
||||
|
||||
def test_aggregates_have_function_argument_and_distinct():
|
||||
f = facets("SELECT COUNT(*), COUNT(1), COUNT(DISTINCT o.order_id), SUM(o.amount), "
|
||||
"SUM(o.amount) OVER (PARTITION BY o.region) FROM orders o")
|
||||
assert [(g.func, g.arg.text, g.distinct) for g in f.aggregates] == [
|
||||
("count", "*", False), ("count", "order_id", True), ("sum", "amount", False)], \
|
||||
"COUNT(1) 和 COUNT(*) 是同一个；窗口函数不算聚合"
|
||||
|
||||
|
||||
def test_having_and_limit():
|
||||
f = facets("SELECT region, SUM(amount) FROM orders GROUP BY region HAVING 100 < SUM(amount) LIMIT 10")
|
||||
assert [p.text for p in f.having] == ["SUM(amount) > 100"] and f.limit == "10"
|
||||
assert facets("SELECT a FROM orders LIMIT 5, 20", "mysql").limit == "20"
|
||||
|
||||
|
||||
def test_having_reads_select_aliases():
|
||||
"""HAVING 里引用 SELECT 的别名（MySQL、SQLite 的写法）换成别名指的表达式，和 GROUP BY 别名一样。"""
|
||||
by_alias = facets("SELECT region AS r, SUM(amount) AS s FROM orders GROUP BY r HAVING s > 100")
|
||||
explicit = facets("SELECT o.region, SUM(o.amount) FROM orders o GROUP BY o.region HAVING SUM(o.amount) > 100")
|
||||
assert [p.key for p in by_alias.having] == [p.key for p in explicit.having] == ["SUM(orders.amount) > 100"]
|
||||
|
||||
|
||||
def test_limits_are_collected_per_segment_with_offset_and_order():
|
||||
"""LIMIT 逐段收集（CTE、FROM 里的子查询各算一段），连同 OFFSET 和这一段的排序：只取前几行时，按什么排、
|
||||
跳过几行决定了取到的是哪几行。OFFSET 0 等于没写。"""
|
||||
f = facets("WITH top AS (SELECT amount FROM orders ORDER BY amount DESC LIMIT 10) "
|
||||
"SELECT SUM(amount) FROM top LIMIT 1")
|
||||
assert [(x.scope, x.count, x.offset) for x in f.limits] == [(None, "1", None), ("top", "10", None)]
|
||||
[(key, desc)] = f.limits[1].order
|
||||
assert key.text == "amount" and desc is True
|
||||
mysql = facets("SELECT a FROM orders ORDER BY 1 LIMIT 5, 20", "mysql")
|
||||
assert (mysql.limits[0].count, mysql.limits[0].offset) == ("20", "5")
|
||||
assert [(k.text, d) for k, d in mysql.limits[0].order] == [("a", False)], "ORDER BY 1 换成 SELECT 里的第一项"
|
||||
assert facets("SELECT a FROM orders LIMIT 10 OFFSET 0").limits[0].offset is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dialect,sql", [
|
||||
("sqlite", 'SELECT SUM("金额") AS "销售额" FROM "订单" WHERE "是否作废" = 0'),
|
||||
("mysql", "SELECT SUM(`金额`) AS `销售额` FROM `订单` WHERE `是否作废` = 0"),
|
||||
("mysql", 'SELECT SUM(`金额`) AS "销售额" FROM `订单` WHERE `是否作废` = 0'),
|
||||
("postgres", 'SELECT SUM("金额") AS "销售额" FROM "订单" WHERE "是否作废" = 0'),
|
||||
])
|
||||
def test_quoted_chinese_names_and_aliases(dialect, sql):
|
||||
f = facets(sql, dialect)
|
||||
assert f.parse_error is None
|
||||
assert [t.key for t in f.tables] == ["订单"]
|
||||
assert [(g.func, g.arg.key) for g in f.aggregates] == [("sum", '"订单"."金额"')] or \
|
||||
[(g.func, g.arg.columns) for g in f.aggregates] == [("sum", (("订单", "金额"),))]
|
||||
assert f.where[0].columns == (("订单", "是否作废"),)
|
||||
|
||||
|
||||
def test_a_cte_that_passes_columns_through_maps_them_to_the_table():
|
||||
wrapped = facets("WITH v AS (SELECT * FROM orders WHERE is_deleted = 0) SELECT SUM(v.amount) FROM v")
|
||||
direct = facets("SELECT SUM(amount) FROM orders WHERE is_deleted = 0")
|
||||
assert wrapped.aggregates == direct.aggregates and wrapped.where == direct.where
|
||||
|
||||
|
||||
def test_a_correlated_subquery_reads_the_outer_alias():
|
||||
"""子查询里引用外层的别名（o.id）到外层去找：写成真实表名，两张表的列互相比的条件也写出表名。"""
|
||||
f = facets("SELECT o.id FROM orders o WHERE NOT EXISTS (SELECT 1 FROM refunds r WHERE r.order_id = o.id)")
|
||||
assert "orders.id = refunds.order_id" in [p.text for p in f.where]
|
||||
plain = facets("SELECT orders.id FROM orders WHERE NOT EXISTS "
|
||||
"(SELECT 1 FROM refunds WHERE orders.id = refunds.order_id)")
|
||||
assert plain.where[1].key == f.where[1].key
|
||||
# 外层那个条件（NOT EXISTS 整句）也一样：子查询按它自己那一段的来源写，别名换成真实表名
|
||||
assert plain.where[0].key == f.where[0].key
|
||||
assert f.where[0].excludes == "EXISTS(SELECT 1 FROM refunds WHERE orders.id = refunds.order_id)"
|
||||
|
||||
|
||||
def test_a_subquery_in_a_condition_is_written_with_its_own_tables():
|
||||
f = facets("SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT u.id FROM users u WHERE u.vip = 1)")
|
||||
assert f.where[0].text == "user_id IN (SELECT id FROM users WHERE vip = 1)"
|
||||
assert f.where[0].key == "orders.user_id IN (SELECT users.id FROM users WHERE users.vip = 1)"
|
||||
|
||||
|
||||
def test_the_schema_tells_which_table_an_unqualified_column_belongs_to():
|
||||
sql = "SELECT SUM(amount) FROM orders o JOIN stores s ON s.id = o.store_id WHERE city = '虚构市'"
|
||||
assert facets(sql).where[0].columns == ((None, "city"),), "不给表结构：几张表都可能有这一列，照写列名"
|
||||
schema = {"tables": {"orders": {"columns": [{"name": "id"}, {"name": "amount"}, {"name": "store_id"}]},
|
||||
"stores": {"columns": [{"name": "id"}, {"name": "city"}]}}}
|
||||
assert facets(sql, schema=schema).where[0].columns == (("stores", "city"),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sql,words", [
|
||||
("SELECT FROM WHERE", "无法解析"),
|
||||
("SELECT 1; SELECT 2", "2 条"),
|
||||
("UPDATE orders SET amount = 0", "不是查询"),
|
||||
("", "为空"),
|
||||
])
|
||||
def test_unparsable_sql_carries_parse_error_and_never_raises(sql, words):
|
||||
f = facets(sql, "sqlite")
|
||||
assert f.parse_error and words in f.parse_error
|
||||
assert (f.tables, f.joins, f.where, f.aggregates) == ((), (), (), ())
|
||||
|
||||
|
||||
def test_facets_are_json_ready():
|
||||
f = facets("SELECT COUNT(DISTINCT order_id) FROM orders WHERE paid_at >= '2026-09-01' GROUP BY region")
|
||||
data = json.loads(json.dumps(f.as_dict(), ensure_ascii=False))
|
||||
assert data["aggregates"][0]["distinct"] is True and data["time"][0]["op"] == ">="
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# diff_facets：每个方面一对
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
BASE = "SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-01' AND paid_at < '2026-10-01'"
|
||||
|
||||
|
||||
def test_time_column_differs_with_catalog_names():
|
||||
a = "SELECT SUM(amount) FROM orders WHERE paid_at >= '2026-09-01' AND paid_at < '2026-10-01'"
|
||||
b = "SELECT SUM(amount) FROM orders WHERE created_at >= '2026-09-01' AND created_at < '2026-10-01'"
|
||||
[diff] = diff_facets(facets(a), facets(b), names={"orders.paid_at": "支付时间", "created_at": "下单时间"})
|
||||
assert diff.aspect == "time"
|
||||
assert diff.text == "A 按支付时间 paid_at 筛时间，B 按下单时间 created_at"
|
||||
assert (diff.a, diff.b) == ("paid_at", "created_at")
|
||||
assert diff.detail == {"a": [{"table": "orders", "column": "paid_at"}],
|
||||
"b": [{"table": "orders", "column": "created_at"}]}
|
||||
assert texts(a, b) == ["A 按 paid_at 筛时间，B 按 created_at"], "没有中文名就只写列名"
|
||||
|
||||
|
||||
def test_time_filter_on_one_side_only():
|
||||
assert texts(BASE, "SELECT SUM(amount) FROM orders") == ["A 按 paid_at 筛时间，B 没有按时间筛选"]
|
||||
assert texts("SELECT SUM(amount) FROM orders", BASE) == ["A 没有按时间筛选，B 按 paid_at 筛时间"]
|
||||
|
||||
|
||||
def test_same_time_column_different_range():
|
||||
other = "SELECT SUM(amount) FROM orders WHERE paid_at BETWEEN '2026-09-01' AND '2026-09-30'"
|
||||
[diff] = diff_facets(facets(BASE), facets(other))
|
||||
assert diff.aspect == "time_range"
|
||||
assert diff.text == ("A 的时间范围是 paid_at >= '2026-09-01' AND paid_at < '2026-10-01'，"
|
||||
"B 是 paid_at BETWEEN '2026-09-01' AND '2026-09-30'")
|
||||
|
||||
|
||||
def test_only_literals_differ_is_still_a_different_range():
|
||||
"""规范文本相同（同一个模板），但比口径时日期不同就是不同的范围：两个函数回答的是两个问题。"""
|
||||
later = BASE.replace("2026-09-01", "2026-09-08").replace("2026-10-01", "2026-10-08")
|
||||
assert normalize_sql(BASE).same_shape(normalize_sql(later))
|
||||
assert aspects(BASE, later) == ["time_range"]
|
||||
|
||||
|
||||
def test_an_exclusion_on_one_side():
|
||||
assert texts("SELECT SUM(amount) FROM orders WHERE is_deleted <> 1", "SELECT SUM(amount) FROM orders") == \
|
||||
["A 排除了 is_deleted = 1 的行，B 没有"]
|
||||
assert texts("SELECT SUM(amount) FROM orders", "SELECT SUM(amount) FROM orders WHERE refund_id IS NOT NULL") == \
|
||||
["B 排除了 refund_id 为空的行，A 没有"]
|
||||
|
||||
|
||||
def test_a_kept_condition_on_one_side():
|
||||
assert texts("SELECT SUM(amount) FROM orders", "SELECT SUM(amount) FROM orders WHERE status = 'paid'") == \
|
||||
["B 只保留 status = 'paid' 的行，A 没有这个条件"]
|
||||
|
||||
|
||||
def test_two_conditions_on_the_same_column_are_paired():
|
||||
[diff] = diff_facets(facets("SELECT SUM(amount) FROM orders WHERE status = 'paid'"),
|
||||
facets("SELECT SUM(amount) FROM orders WHERE status IN ('paid', 'shipped')"))
|
||||
assert (diff.aspect, diff.text) == ("filter", "A 按 status = 'paid' 筛选，B 按 status IN ('paid', 'shipped') 筛选")
|
||||
assert (diff.a, diff.b) == ("status = 'paid'", "status IN ('paid', 'shipped')")
|
||||
|
||||
|
||||
def test_distinct_count_versus_plain_count():
|
||||
assert texts("SELECT COUNT(DISTINCT order_id) FROM orders", "SELECT COUNT(order_id) FROM orders") == \
|
||||
["A 对 order_id 去重计数，B 没去重"]
|
||||
assert texts("SELECT COUNT(order_id) FROM orders", "SELECT COUNT(DISTINCT order_id) FROM orders") == \
|
||||
["A 对 order_id 计数时没去重，B 去重计数"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("a,b,text", [
|
||||
("SELECT SUM(amount) FROM orders", "SELECT AVG(amount) FROM orders", "A 对 amount 求和，B 求平均"),
|
||||
("SELECT SUM(paid_amount) FROM orders", "SELECT SUM(amount) FROM orders", "A 对 paid_amount 求和，B 对 amount 求和"),
|
||||
("SELECT COUNT(*) FROM orders", "SELECT COUNT(DISTINCT order_id) FROM orders",
|
||||
"A 数的是行数，B 对 order_id 去重计数"),
|
||||
])
|
||||
def test_other_aggregate_differences(a, b, text):
|
||||
assert texts(a, b) == [text]
|
||||
|
||||
|
||||
def test_join_differences():
|
||||
left = "SELECT SUM(o.amount) FROM orders o LEFT JOIN refunds r ON r.order_id = o.id"
|
||||
inner = "SELECT SUM(o.amount) FROM orders o JOIN refunds r ON r.order_id = o.id"
|
||||
other_key = "SELECT SUM(o.amount) FROM orders o LEFT JOIN refunds r ON r.user_id = o.user_id"
|
||||
plain = "SELECT SUM(amount) FROM orders"
|
||||
assert texts(left, inner) == ["A 用 LEFT JOIN 关联 refunds 表，B 用 INNER JOIN"]
|
||||
assert texts(left, other_key) == ["A 关联 refunds 表的条件是 orders.id = refunds.order_id，"
|
||||
"B 是 orders.user_id = refunds.user_id"]
|
||||
assert texts(left, plain) == ["A 关联了 refunds 表，B 没有"]
|
||||
|
||||
|
||||
def test_a_different_source_table_is_said_once():
|
||||
"""表换了一张：说一次「查的是哪张表」就够了，同一个聚合、同一个条件不再因为表名不同各报一遍。"""
|
||||
a = "SELECT SUM(amount) FROM orders WHERE is_deleted = 0 AND paid_at >= '2026-09-01'"
|
||||
b = "SELECT SUM(amount) FROM orders_archive WHERE is_deleted = 0 AND paid_at >= '2026-09-01'"
|
||||
[diff] = diff_facets(facets(a), facets(b))
|
||||
assert (diff.aspect, diff.text) == ("tables", "A 查的是 orders 表，B 查的是 orders_archive 表")
|
||||
assert aspects(a, b.replace("is_deleted = 0", "is_deleted = 1")) == ["filter", "tables"], "换表之外的不同照报"
|
||||
|
||||
|
||||
def test_two_columns_that_read_the_same_are_told_apart():
|
||||
"""各自只写列名时两边读起来一样，其实是两张表的列：改写成带表名的写法，不说「A 对 amount 求和，B 对 amount 求和」。"""
|
||||
a = "SELECT SUM(o.amount) FROM orders o JOIN refunds r ON r.order_id = o.id"
|
||||
b = "SELECT SUM(r.amount) FROM orders o JOIN refunds r ON r.order_id = o.id"
|
||||
assert texts(a, b) == ["A 对 orders.amount 求和，B 对 refunds.amount 求和"]
|
||||
|
||||
|
||||
def test_group_by_difference():
|
||||
assert texts("SELECT region, SUM(amount) FROM orders GROUP BY region",
|
||||
"SELECT region, city, SUM(amount) FROM orders GROUP BY region, city") == \
|
||||
["A 按 region 分组，B 按 region、city 分组"]
|
||||
assert texts("SELECT SUM(amount) FROM orders", "SELECT region, SUM(amount) FROM orders GROUP BY region") == \
|
||||
["A 没有分组，B 按 region 分组"]
|
||||
|
||||
|
||||
def test_having_difference():
|
||||
assert texts("SELECT region, SUM(amount) FROM orders GROUP BY region HAVING SUM(amount) > 100",
|
||||
"SELECT region, SUM(amount) FROM orders GROUP BY region") == \
|
||||
["A 分组后只保留 SUM(amount) > 100 的组，B 没有这个条件"]
|
||||
|
||||
|
||||
def test_limit_difference():
|
||||
assert texts("SELECT amount FROM orders LIMIT 10", "SELECT amount FROM orders") == ["A 只取前 10 行，B 不限行数"]
|
||||
assert texts("SELECT amount FROM orders LIMIT 10", "SELECT amount FROM orders LIMIT 100") == \
|
||||
["A 只取前 10 行，B 只取前 100 行"]
|
||||
|
||||
|
||||
def test_an_unparsable_side_is_reported_once_without_item_by_item_comparison():
|
||||
[diff] = diff_facets(facets("SELECT FROM WHERE"), facets(BASE))
|
||||
assert diff.aspect == "parse" and diff.text.startswith("A：SQL 无法解析") and "没有逐项比较" in diff.text
|
||||
assert diff.b is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("a,b", [
|
||||
# 别名、比较两边的顺序、大小写
|
||||
("SELECT SUM(o.amount) FROM orders o WHERE o.paid_at >= '2026-09-01' AND o.is_deleted = 0",
|
||||
"select sum(amount) from orders where 0 = is_deleted and '2026-09-01' <= paid_at"),
|
||||
# 有效记录条件写在 CTE 里、还是写在外层
|
||||
("WITH v AS (SELECT * FROM orders WHERE is_deleted = 0) SELECT COUNT(DISTINCT v.order_id) FROM v",
|
||||
"SELECT COUNT(DISTINCT order_id) FROM orders WHERE is_deleted = 0"),
|
||||
# 逗号连接和 INNER JOIN
|
||||
("SELECT SUM(o.amount) FROM orders o, stores s WHERE o.store_id = s.id AND s.city = '虚构市'",
|
||||
"SELECT SUM(orders.amount) FROM orders INNER JOIN stores ON stores.id = orders.store_id "
|
||||
"WHERE stores.city = '虚构市'"),
|
||||
# 条件里套着的子查询用了别名：NOT EXISTS、IN (SELECT …)
|
||||
("SELECT o.id FROM orders o WHERE NOT EXISTS (SELECT 1 FROM refunds r WHERE r.order_id = o.id)",
|
||||
"SELECT orders.id FROM orders WHERE NOT EXISTS (SELECT 1 FROM refunds WHERE refunds.order_id = orders.id)"),
|
||||
("SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT u.id FROM users u WHERE u.vip = 1)",
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN (SELECT id FROM users WHERE vip = 1)"),
|
||||
# HAVING 里用 SELECT 的别名
|
||||
("SELECT o.region, SUM(o.amount) FROM orders o GROUP BY o.region HAVING SUM(o.amount) > 100",
|
||||
"SELECT region AS r, SUM(amount) AS s FROM orders GROUP BY r HAVING s > 100"),
|
||||
# OFFSET 0 等于没写；取前几行时 ORDER BY 1 和写出列名是同一个排序
|
||||
("SELECT amount FROM orders ORDER BY 1 DESC LIMIT 10 OFFSET 0",
|
||||
"SELECT o.amount FROM orders o ORDER BY o.amount DESC LIMIT 10"),
|
||||
])
|
||||
def test_the_same_meaning_written_differently_is_no_difference(a, b):
|
||||
assert texts(a, b) == []
|
||||
|
||||
|
||||
def test_a_real_difference_inside_a_subquery_is_still_reported():
|
||||
a = "SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT u.id FROM users u WHERE u.vip = 1)"
|
||||
b = "SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT u.referrer_id FROM users u WHERE u.vip = 1)"
|
||||
assert texts(a, b) == ["A 按 user_id IN (SELECT id FROM users WHERE vip = 1) 筛选，"
|
||||
"B 按 user_id IN (SELECT referrer_id FROM users WHERE vip = 1) 筛选"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dialect,a,b", [
|
||||
# 条件里的子查询给选出来的列起了名字：IN、标量子查询只看取值，列名不影响口径
|
||||
("sqlite",
|
||||
"SELECT SUM(o.amount) FROM orders o WHERE o.amount > "
|
||||
"(SELECT AVG(o2.amount) AS avg_amount FROM orders o2 WHERE o2.is_deleted = 0)",
|
||||
"SELECT SUM(amount) FROM orders WHERE amount > (SELECT AVG(amount) FROM orders WHERE is_deleted = 0)"),
|
||||
("sqlite",
|
||||
"SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT u.id AS uid FROM users u WHERE u.vip = 1)",
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN (SELECT id FROM users WHERE vip = 1)"),
|
||||
# 条件里的子查询从 FROM 里的子查询取数，两边给它起的别名不同
|
||||
("sqlite",
|
||||
"SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT s.user_id FROM "
|
||||
"(SELECT user_id, SUM(amount) AS total FROM payments GROUP BY user_id) s WHERE s.total > 100)",
|
||||
"SELECT SUM(o.amount) FROM orders o WHERE o.user_id IN (SELECT agg.user_id FROM "
|
||||
"(SELECT user_id, SUM(amount) AS total FROM payments GROUP BY user_id) agg WHERE agg.total > 100)"),
|
||||
# CTE 起的名字不同，外层按 CTE 算出来的列筛选
|
||||
("sqlite",
|
||||
"WITH user_totals AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE is_deleted = 0 GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM user_totals WHERE total > 100",
|
||||
"WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE is_deleted = 0 GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM t WHERE t.total > 100"),
|
||||
# 写成 CTE 还是写成 FROM 里的子查询
|
||||
("postgres",
|
||||
"WITH user_totals AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM user_totals WHERE total > 100",
|
||||
"SELECT COUNT(*) FROM (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) AS agg "
|
||||
"WHERE agg.total > 100"),
|
||||
# LATERAL 子查询的别名不同：别名挂在 LATERAL 上，不在子查询上
|
||||
("postgres",
|
||||
"SELECT SUM(o.amount) FROM orders o, LATERAL (SELECT COUNT(*) AS n FROM refunds r WHERE r.order_id = o.id) l "
|
||||
"WHERE l.n > 0",
|
||||
"SELECT SUM(o.amount) FROM orders o CROSS JOIN LATERAL (SELECT COUNT(*) AS n FROM refunds r "
|
||||
"WHERE r.order_id = o.id) AS x WHERE x.n > 0"),
|
||||
# 关联进来的是子查询，两边的别名不同（一边是带反引号的中文别名）
|
||||
("mysql",
|
||||
"SELECT SUM(o.amount) FROM orders o JOIN (SELECT order_id, COUNT(*) AS n FROM refunds GROUP BY order_id) r "
|
||||
"ON r.order_id = o.id WHERE r.n > 1",
|
||||
"SELECT SUM(o.amount) FROM orders o JOIN (SELECT order_id, COUNT(*) AS n FROM refunds GROUP BY order_id) "
|
||||
"`退款` ON o.id = `退款`.order_id WHERE `退款`.n > 1"),
|
||||
# 条件里的子查询用到 CTE，两边给 CTE 起的名字不同
|
||||
("sqlite",
|
||||
"WITH vip AS (SELECT id FROM users WHERE grade >= 3) "
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN (SELECT id FROM vip)",
|
||||
"WITH good_users AS (SELECT id FROM users WHERE grade >= 3) "
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN (SELECT g.id FROM good_users g)"),
|
||||
# CTE 定义在条件里的子查询里：定义连同子查询整个写进外层的条件，定义上的名字也不算
|
||||
("postgres",
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN "
|
||||
"(WITH t AS (SELECT id FROM users WHERE grade >= 3) SELECT id FROM t)",
|
||||
"SELECT SUM(amount) FROM orders WHERE user_id IN "
|
||||
"(WITH g AS (SELECT id FROM users WHERE grade >= 3) SELECT g.id FROM g)"),
|
||||
# 两个 CTE 定义的先后、起的名字都不同，按用到的先后它们是同一对
|
||||
("sqlite",
|
||||
"WITH paid AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE status = 'paid' GROUP BY user_id), "
|
||||
"refunded AS (SELECT user_id, SUM(amount) AS total FROM refunds GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM paid p JOIN refunded r ON r.user_id = p.user_id WHERE p.total > r.total",
|
||||
"WITH r AS (SELECT user_id, SUM(amount) AS total FROM refunds GROUP BY user_id), "
|
||||
"p AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE status = 'paid' GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM p JOIN r ON p.user_id = r.user_id WHERE p.total > r.total"),
|
||||
# CTE 定义的先后和用到的先后不一致，另一边又把它们写成了 FROM 里的子查询：按用到的先后编号才对得上
|
||||
("sqlite",
|
||||
"WITH c AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id), "
|
||||
"d AS (SELECT user_id, COUNT(*) AS n FROM refunds GROUP BY user_id), "
|
||||
"a AS (SELECT * FROM c WHERE total > 100), b AS (SELECT * FROM d WHERE n > 1) "
|
||||
"SELECT COUNT(*) FROM b JOIN a ON a.user_id = b.user_id",
|
||||
"SELECT COUNT(*) FROM (SELECT * FROM (SELECT user_id, COUNT(*) AS n FROM refunds GROUP BY user_id) x "
|
||||
"WHERE n > 1) b JOIN (SELECT * FROM (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) y "
|
||||
"WHERE total > 100) a ON a.user_id = b.user_id"),
|
||||
# 同上，两个 CTE 还算出了同名的列、外层拿它们互相比：不写编号的写法分不出是哪个的 total，编号得对得上
|
||||
("sqlite",
|
||||
"WITH c AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id), "
|
||||
"d AS (SELECT user_id, SUM(amount) AS total FROM refunds GROUP BY user_id), "
|
||||
"a AS (SELECT * FROM c), b AS (SELECT * FROM d) "
|
||||
"SELECT COUNT(*) FROM b JOIN a ON a.user_id = b.user_id WHERE a.total > b.total",
|
||||
"SELECT COUNT(*) FROM (SELECT * FROM (SELECT user_id, SUM(amount) AS total FROM refunds GROUP BY user_id) x) b "
|
||||
"JOIN (SELECT * FROM (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) y) a "
|
||||
"ON a.user_id = b.user_id WHERE a.total > b.total"),
|
||||
# 先用一个 CTE 筛一遍再汇总，还是直接在汇总的 CTE 里筛：里层多一段，不挪动外层那一段的编号
|
||||
("sqlite",
|
||||
"WITH base AS (SELECT * FROM orders WHERE is_deleted = 0), "
|
||||
"per_user AS (SELECT user_id, SUM(amount) AS total FROM base GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM per_user WHERE total > 100",
|
||||
"WITH totals AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE is_deleted = 0 GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM totals WHERE total > 100"),
|
||||
])
|
||||
def test_names_given_to_derived_tables_are_no_difference(dialect, a, b):
|
||||
"""CTE 名、子查询的别名、条件里的子查询选出来的列名都是各自起的（双盲复算的盲算更是）：只因起的名字不同，
|
||||
不算换了口径。"""
|
||||
assert texts(a, b, dialect) == []
|
||||
|
||||
|
||||
def test_derived_tables_are_compared_by_position_and_shown_by_name():
|
||||
"""派生来源（CTE、FROM 里的子查询）比较用的写法按它在语句里的位置编号，不写名字；给人看的写法照原样写名字。"""
|
||||
a = facets("WITH user_totals AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM user_totals WHERE total > 100")
|
||||
b = facets("SELECT COUNT(*) FROM (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) agg "
|
||||
"WHERE agg.total > 100")
|
||||
assert a.where[0].key == b.where[0].key
|
||||
assert "user_totals" not in a.where[0].key and "agg" not in b.where[0].key
|
||||
assert (a.where[0].text, a.where[0].full) == ("total > 100", "user_totals.total > 100")
|
||||
assert (b.where[0].text, b.where[0].full) == ("total > 100", "agg.total > 100")
|
||||
joined = facets("SELECT SUM(o.amount) FROM orders o JOIN (SELECT order_id, COUNT(*) AS n FROM refunds "
|
||||
"GROUP BY order_id) r ON r.order_id = o.id")
|
||||
[join] = joined.joins
|
||||
assert join.table.text == "r" and join.on.text == "orders.id = refunds.order_id"
|
||||
|
||||
|
||||
SHIFTED_A = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM t WHERE total > 100 AND user_id IN (SELECT id FROM users WHERE vip = 1)")
|
||||
SHIFTED_B = ("WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM vip JOIN t ON t.user_id = vip.id WHERE t.total > 100")
|
||||
|
||||
|
||||
def test_a_shifted_derived_number_does_not_split_one_condition_in_two():
|
||||
"""编号只看位置：一边多出一个排在前面的 CTE（筛 VIP 用户），后面那个 CTE 的编号就错开了。按编号比不上的再按
|
||||
不写编号的写法一对一地配，同一个条件不会被说成两边各有一个；真正的不同（VIP 用 IN 筛还是关联进来）照报。"""
|
||||
found = diff_facets(facets(SHIFTED_A), facets(SHIFTED_B), labels=("原流程", "盲算"))
|
||||
filters = [d.text for d in found if d.aspect == "filter"]
|
||||
assert filters == ["原流程只保留 user_id IN (SELECT id FROM users WHERE vip = 1) 的行，盲算没有这个条件"]
|
||||
# 错开之后的条件真不同，也照样配成一对说
|
||||
other = SHIFTED_B.replace("t.total > 100", "t.total > 200")
|
||||
filters = [d.text for d in diff_facets(facets(SHIFTED_A), facets(other), labels=("原流程", "盲算"))
|
||||
if d.aspect == "filter"]
|
||||
assert filters[0] == "原流程按 total > 100 筛选，盲算按 total > 200 筛选" and len(filters) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("aspect,a,b", [
|
||||
# 时间条件、分组键、聚合、HAVING、取前几行时的排序都是按 CTE 算出来的列
|
||||
("time",
|
||||
"WITH t AS (SELECT user_id, DATE(paid_at) AS day, amount FROM orders) "
|
||||
"SELECT SUM(amount) FROM t WHERE day >= '2026-09-01' AND user_id IN (SELECT id FROM users WHERE vip = 1)",
|
||||
"WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, DATE(paid_at) AS day, amount FROM orders) "
|
||||
"SELECT SUM(t.amount) FROM vip JOIN t ON t.user_id = vip.id WHERE t.day >= '2026-09-01'"),
|
||||
("group_by",
|
||||
"WITH t AS (SELECT user_id, SUM(amount) AS total, COUNT(*) AS n FROM orders GROUP BY user_id) "
|
||||
"SELECT n, COUNT(*) FROM t WHERE user_id IN (SELECT id FROM users WHERE vip = 1) GROUP BY n",
|
||||
"WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total, COUNT(*) AS n FROM orders GROUP BY user_id) "
|
||||
"SELECT t.n, COUNT(*) FROM vip JOIN t ON t.user_id = vip.id GROUP BY t.n"),
|
||||
("aggregate",
|
||||
"WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT AVG(total) FROM t WHERE user_id IN (SELECT id FROM users WHERE vip = 1)",
|
||||
"WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT AVG(t.total) FROM vip JOIN t ON t.user_id = vip.id"),
|
||||
("having",
|
||||
"WITH t AS (SELECT store_id, user_id, SUM(amount) AS total FROM orders GROUP BY store_id, user_id) "
|
||||
"SELECT store_id FROM t WHERE user_id IN (SELECT id FROM users WHERE vip = 1) GROUP BY store_id "
|
||||
"HAVING MAX(total) > 100",
|
||||
"WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT store_id, user_id, SUM(amount) AS total FROM orders GROUP BY store_id, user_id) "
|
||||
"SELECT t.store_id FROM vip JOIN t ON t.user_id = vip.id GROUP BY t.store_id HAVING MAX(t.total) > 100"),
|
||||
("limit",
|
||||
"WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT user_id FROM t WHERE user_id IN (SELECT id FROM users WHERE vip = 1) ORDER BY total DESC LIMIT 10",
|
||||
"WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT t.user_id FROM vip JOIN t ON t.user_id = vip.id ORDER BY t.total DESC LIMIT 10"),
|
||||
])
|
||||
def test_a_shifted_derived_number_is_no_difference_in_any_aspect(aspect, a, b):
|
||||
"""编号错开之后，按 CTE 算出来的列写的时间条件、分组键、聚合、HAVING、排序也不算不同：只剩 VIP 那处真正的
|
||||
不同。"""
|
||||
assert not [x for x in aspects(a, b) if x.startswith(aspect)]
|
||||
|
||||
|
||||
def test_a_shifted_number_inside_a_condition_subquery_is_no_difference():
|
||||
"""条件里的子查询用到的 CTE 编号错开了：子查询整个写进外层条件的那一份，按不写编号的写法也写「@」，照样配得上。"""
|
||||
a = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM users u WHERE u.vip = 1 AND u.id IN (SELECT user_id FROM t WHERE total > 100)")
|
||||
b = ("WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM vip WHERE vip.id IN (SELECT user_id FROM t WHERE total > 100)")
|
||||
assert texts(a, b) == []
|
||||
|
||||
|
||||
def test_a_shifted_derived_join_is_paired_with_its_counterpart():
|
||||
"""关联进来的 CTE 编号错开了（盲算多关联了一个排在前面的 CTE）：先把完全一样的关联配掉，再按关联的是哪张表配，
|
||||
不会把原流程的 t 和盲算的 vip 当成同一处关联、说成「关联条件不同」。"""
|
||||
a = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM users u JOIN t ON t.user_id = u.id WHERE t.total > 100")
|
||||
b = ("WITH vip AS (SELECT id FROM users WHERE vip = 1), "
|
||||
"t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM users u JOIN vip ON vip.id = u.id JOIN t ON t.user_id = u.id WHERE t.total > 100")
|
||||
assert texts(a, b, labels=("原流程", "盲算")) == ["盲算只保留 vip = 1 的行，原流程没有这个条件",
|
||||
"盲算关联了 vip 表，原流程没有"]
|
||||
|
||||
|
||||
def test_swapped_join_order_of_two_derived_tables_is_no_difference():
|
||||
"""两个 CTE 关联的先后对调，编号也跟着对调：拿它们的列互相比的条件，两边先按不写编号的写法排好两边的先后，
|
||||
才配得上（只按编号排，同一个比较会排成两个方向）。"""
|
||||
ctes = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id), "
|
||||
"u AS (SELECT user_id, COUNT(*) AS cnt FROM refunds GROUP BY user_id) ")
|
||||
assert texts(ctes + "SELECT COUNT(*) FROM t JOIN u ON u.user_id = t.user_id WHERE t.total > u.cnt",
|
||||
ctes + "SELECT COUNT(*) FROM u JOIN t ON t.user_id = u.user_id WHERE t.total > u.cnt") == []
|
||||
|
||||
|
||||
def test_a_reversed_comparison_between_two_derived_tables_is_still_reported():
|
||||
"""两个 CTE 算出了同名的列：编号对得上时照样分得清是哪个 CTE 的，比较的方向反了要报出来（不写编号的写法只在
|
||||
按编号比不上时才拿来配，不会把它吞掉）。"""
|
||||
a = ("WITH paid AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id), "
|
||||
"refunded AS (SELECT user_id, SUM(amount) AS total FROM refunds GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM paid p JOIN refunded r ON r.user_id = p.user_id WHERE p.total > r.total")
|
||||
b = a.replace("WHERE p.total > r.total", "WHERE r.total > p.total")
|
||||
assert texts(a, b) == ["A 只保留 paid.total > refunded.total 的行，B 没有这个条件",
|
||||
"B 只保留 paid.total < refunded.total 的行，A 没有这个条件"]
|
||||
|
||||
|
||||
def test_a_real_difference_next_to_renamed_derived_tables_is_reported_once():
|
||||
"""派生来源起的名字不同、里面又真有不同时，只报真正不同的那一处：CTE 里少一个条件，外层按 CTE 算出来的列
|
||||
筛选的条件不跟着报（不然两边各报一句「只保留 total > 100 的行」，自相矛盾）；外层的条件本身不同，两边配成
|
||||
一对说。"""
|
||||
a = ("WITH user_totals AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE is_deleted = 0 "
|
||||
"GROUP BY user_id) SELECT COUNT(*) FROM user_totals WHERE total > 100")
|
||||
inner = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM t WHERE total > 100")
|
||||
assert texts(a, inner, labels=("原流程", "盲算")) == ["原流程只保留 is_deleted = 0 的行，盲算没有这个条件"]
|
||||
outer = ("WITH t AS (SELECT user_id, SUM(amount) AS total FROM orders WHERE is_deleted = 0 GROUP BY user_id) "
|
||||
"SELECT COUNT(*) FROM t WHERE total > 200")
|
||||
assert texts(a, outer, labels=("原流程", "盲算")) == ["原流程按 total > 100 筛选，盲算按 total > 200 筛选"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("a,b,text", [
|
||||
("WITH top AS (SELECT amount FROM orders ORDER BY amount DESC LIMIT 10) SELECT SUM(amount) FROM top",
|
||||
"WITH top AS (SELECT amount FROM orders ORDER BY amount DESC LIMIT 100) SELECT SUM(amount) FROM top",
|
||||
"A 在 top 里只取前 10 行，B 只取前 100 行"),
|
||||
("SELECT amount FROM orders LIMIT 10 OFFSET 0", "SELECT amount FROM orders LIMIT 10 OFFSET 10",
|
||||
"A 不跳过行，B 跳过前 10 行"),
|
||||
("SELECT amount FROM orders ORDER BY 1 DESC LIMIT 10", "SELECT amount FROM orders ORDER BY 1 ASC LIMIT 10",
|
||||
"A 取行前按 amount 降序排，B 按 amount 升序排"),
|
||||
("SELECT amount FROM orders ORDER BY amount DESC LIMIT 10", "SELECT amount FROM orders LIMIT 10",
|
||||
"A 取行前按 amount 降序排，B 没有排序"),
|
||||
])
|
||||
def test_which_rows_are_taken_is_compared(a, b, text):
|
||||
"""取哪几行：每一段的行数、跳过的行数，以及取前几行时的排序。排序只在有 LIMIT / OFFSET 的段里比：
|
||||
没有取前几行，排序不影响算出来的数。"""
|
||||
assert texts(a, b) == [text]
|
||||
assert texts("SELECT SUM(amount) FROM orders ORDER BY 1 DESC", "SELECT SUM(amount) FROM orders") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dialect", DIALECTS)
|
||||
def test_each_dialect_gives_the_same_three_findings(dialect):
|
||||
q = {"sqlite": '"', "postgres": '"', "mysql": "`"}[dialect]
|
||||
a = (f"SELECT COUNT(DISTINCT o.order_id) AS {q}订单数{q} FROM orders o "
|
||||
"WHERE o.paid_at >= '2026-09-01' AND o.is_deleted <> 1")
|
||||
b = f"SELECT COUNT(o.order_id) AS {q}订单数{q} FROM orders o WHERE o.created_at >= '2026-09-01'"
|
||||
assert texts(a, b, dialect) == ["A 按 paid_at 筛时间，B 按 created_at", "A 排除了 is_deleted = 1 的行，B 没有",
|
||||
"A 对 order_id 去重计数，B 没去重"]
|
||||
|
||||
|
||||
def test_side_labels_in_chinese_have_no_stray_spaces():
|
||||
a = "SELECT COUNT(DISTINCT order_id) FROM orders WHERE paid_at >= '2026-09-01'"
|
||||
b = "SELECT COUNT(order_id) FROM orders WHERE created_at >= '2026-09-01'"
|
||||
assert texts(a, b, labels=("原流程", "盲算"), names={"paid_at": "支付时间"}) == [
|
||||
"原流程按支付时间 paid_at 筛时间，盲算按 created_at", "原流程对 order_id 去重计数，盲算没去重"]
|
||||
|
||||
|
||||
def test_catalog_labels_come_from_visible_catalog_items():
|
||||
notes = {
|
||||
"orders": {"columns": {"paid_at": {"label": {"value": "支付时间", "status": "confirmed"}},
|
||||
"status": {"label": {"value": "状态", "status": "proposed"}},
|
||||
"amount": {"label": {"value": "错的名字", "status": "rejected"}}}},
|
||||
"refunds": {"columns": {"status": {"label": {"value": "退款状态", "status": "confirmed"}}}},
|
||||
}
|
||||
labels = catalog_labels(notes)
|
||||
assert labels["orders.paid_at"] == labels["paid_at"] == "支付时间"
|
||||
assert labels["orders.status"] == "状态" and labels["refunds.status"] == "退款状态"
|
||||
assert "status" not in labels, "两张表叫法不同的列不按光列名记"
|
||||
assert "orders.amount" not in labels and "amount" not in labels, "被驳回的不算"
|
||||
|
||||
|
||||
def test_a_facets_object_compares_with_itself_cleanly():
|
||||
f = facets(BASE)
|
||||
assert isinstance(f, Facets) and diff_facets(f, f) == []
|
||||
@@ -0,0 +1,794 @@
|
||||
"""核查用的只读 agent（engine/verify_agent.py）：查库、交回有出处的结果。双盲复算、反方审稿、审批质询共用。
|
||||
|
||||
- 工具只有给定数据源的 db_schema / db_query，建的时候 force_readonly：可写源上的写语句也被拒绝，数据不动
|
||||
- 每次成功的查询是一份查询快照，按这个 agent 自己的顺序编 Q1、Q2…；节点里调用时发 tool.end（带 query_artifact、
|
||||
role: verify），但不往证据台账里写 query：核查的查询进了报告的写作目录，写作者就会把它们当成可以引用的证据
|
||||
- 最后按调用方的结构交结果：数值字段一律是格子引用 {alias, row, column}，由代码从快照取值，模型自己写的数字
|
||||
不采信（type 列表、anyOf、$ref、allOf 里的数也是；认不出的结构下写的数一律丢掉）；没交的字段和交了 null 分开记；
|
||||
理由里的数字要写成 [[v:Qn.rX.列]]，用本 agent 自己的台账组装成文档，裸数字会被标出
|
||||
- 查询次数、时长、金额（单次和每日）任何一项到顶就停止查询、按已有结果收尾；收尾的钱在查询之前就按最贵的样子
|
||||
留好；每日上限缺省按设置；模型不在价格目录里时金额上限无效，照裁判的说法记一条说明
|
||||
- 永远不抛：库写不进去、模型不支持结构化输出，都落在返回值里
|
||||
- 上传表格按调用方给的数据版本查（运行固定的那一版），不拿当前版本顶上
|
||||
|
||||
模型一律是 MockChatModel 剧本，数据源是现造的 SQLite 文件，表名、数字都是虚构的。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from sqlalchemy import delete, select
|
||||
|
||||
from app.core import artifact_store
|
||||
from app.core.config import settings
|
||||
from app.core.events import EventType
|
||||
from app.data import introspect, table_versions
|
||||
from app.data.engine import engines
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import DataSource, Setting
|
||||
from app.engine import spend, verify_agent
|
||||
from app.engine.toolcalls import Limit, ToolTimeout
|
||||
from app.engine.verify_agent import result_schema, run_verify_agent
|
||||
from app.providers import catalog as pricing
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured
|
||||
from app.providers.mock_model import MockChatModel
|
||||
from app.tools.registry import Versions
|
||||
|
||||
SPEC = ModelSpec(model="mock-fast")
|
||||
SCHEMA = {"type": "object", "properties": {
|
||||
"value": {"type": "number", "description": "未作废订单的金额合计"},
|
||||
"verdict": {"type": "string", "enum": ["一致", "不一致", "无法判断"]},
|
||||
}, "required": ["value", "verdict"]}
|
||||
VALID_SUM = "SELECT SUM(amount) AS total FROM orders WHERE is_deleted = 0"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def clean_spend():
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key.in_([spend.SPEND_KEY, "verify"])))
|
||||
await session.commit()
|
||||
yield
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key == "verify"))
|
||||
await session.commit()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def shop(tmp_path):
|
||||
"""可写的手工源：四笔订单，金额 120、80、50（已作废）、30；未作废的合计 230。返回 (源名, 文件路径)。"""
|
||||
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
||||
db = sqlite3.connect(path)
|
||||
db.executescript(
|
||||
"CREATE TABLE orders (id INTEGER PRIMARY KEY, amount INTEGER, is_deleted INTEGER, paid_at TEXT);"
|
||||
"INSERT INTO orders VALUES (1, 120, 0, '2026-09-02'), (2, 80, 0, '2026-09-03'),"
|
||||
" (3, 50, 1, '2026-09-03'), (4, 30, 0, '2026-09-05');")
|
||||
db.commit()
|
||||
db.close()
|
||||
name = f"shop_{uuid.uuid4().hex[:8]}"
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True)
|
||||
session.add(row)
|
||||
await session.flush()
|
||||
row.schema_cache = await introspect.introspect(row)
|
||||
await session.commit()
|
||||
source_id = row.id
|
||||
yield name, str(path)
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
class Script:
|
||||
"""核查模型的剧本。steps：循环里每一轮发的工具调用 [(工具名, 参数, 调用 id)]，或者一段说明文字；剧本用完就收手。
|
||||
final：交结果那一轮返回的 JSON（可以是按消息算的函数）。calls 按顺序记下每次调用的种类和收到的消息。"""
|
||||
|
||||
def __init__(self, monkeypatch, steps, final, *, model_id: str = "mock-fast", latency: float = 0.0):
|
||||
self.calls: list[tuple[str, list]] = []
|
||||
self.formats: list[dict] = []
|
||||
self.specs: list[ModelSpec] = []
|
||||
pending = list(steps)
|
||||
|
||||
def decide(model, messages):
|
||||
if model.response_format:
|
||||
self.calls.append(("final", list(messages)))
|
||||
self.formats.append(model.response_format)
|
||||
body = final(messages) if callable(final) else final
|
||||
return AIMessage(content=json.dumps(body, ensure_ascii=False))
|
||||
self.calls.append(("step", list(messages)))
|
||||
if not pending:
|
||||
return AIMessage(content="查完了。")
|
||||
step = pending.pop(0)
|
||||
if isinstance(step, str):
|
||||
return AIMessage(content=step)
|
||||
return AIMessage(content="查一下。", tool_calls=[{"name": n, "args": a, "id": i} for n, a, i in step])
|
||||
|
||||
monkeypatch.setattr(MockChatModel, "_decide", decide)
|
||||
|
||||
async def fake_model(session, spec):
|
||||
self.specs.append(spec)
|
||||
return MockChatModel(model_name=model_id, latency=latency), model_id
|
||||
|
||||
monkeypatch.setattr(verify_agent, "get_chat_model", fake_model)
|
||||
|
||||
def tool_results(self) -> dict[str, str]:
|
||||
"""模型在循环里看到的工具结果：{调用 id: 内容}。"""
|
||||
out: dict[str, str] = {}
|
||||
for _, messages in self.calls:
|
||||
for m in messages:
|
||||
if isinstance(m, ToolMessage):
|
||||
out[m.tool_call_id] = str(m.content)
|
||||
return out
|
||||
|
||||
@property
|
||||
def kinds(self) -> list[str]:
|
||||
return [k for k, _ in self.calls]
|
||||
|
||||
|
||||
def q(name: str, sql: str, call_id: str) -> list[tuple]:
|
||||
return [(f"db_query__{name}", {"sql": sql}, call_id)]
|
||||
|
||||
|
||||
async def run(name: str, **kwargs):
|
||||
base = dict(spec=SPEC, system="核对口径卡里的「销售额」。", task="算一遍未作废订单的金额合计。", sources=[name],
|
||||
output_schema=SCHEMA, data_versions=Versions.CURRENT)
|
||||
return await run_verify_agent(**{**base, **kwargs})
|
||||
|
||||
|
||||
def counts(path: str) -> tuple[int, int]:
|
||||
db = sqlite3.connect(path)
|
||||
try:
|
||||
return db.execute("SELECT COUNT(*), SUM(amount) FROM orders").fetchone()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 正常查询、取格子
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_cells_are_taken_from_the_snapshot(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
script = Script(monkeypatch, [[(f"db_schema__{name}", {"table": "orders"}, "s1")], q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致",
|
||||
"rationale": "未作废订单的金额合计是 [[v:Q1.r0.total]]。[[see:Q1]]"})
|
||||
out = await run(name)
|
||||
assert out.error is None and out.limit is None
|
||||
assert out.result["value"] == 230 and out.result["verdict"] == "一致"
|
||||
[cell] = out.cells
|
||||
assert cell["status"] == "resolved" and cell["ref"] == "Q1.r0.total" and cell["value"] == 230
|
||||
assert cell["path"] == "value" and cell["eid"].startswith("ev:cell:")
|
||||
[query] = out.queries
|
||||
assert query["alias"] == "Q1" and query["row_count"] == 1 and query["truncated"] is False
|
||||
assert "SUM(amount)" in query["sql"]
|
||||
snapshot = artifact_store.load(query["artifact"])
|
||||
assert snapshot["rows"] == [[230]] and cell["artifact"] == query["artifact"]
|
||||
# 理由按本 agent 自己的台账组装：引用渲染成快照里的值，没有违规
|
||||
assert out.rationale["markdown"] == "未作废订单的金额合计是 230。" and out.rationale["violations"] == []
|
||||
assert out.result["rationale"] == "未作废订单的金额合计是 230。"
|
||||
# 结果里的数值要求写格子引用：交给模型的结构里，value 是 {alias, row, column}
|
||||
assert "alias" in json.dumps(script.formats[0]) and "rationale" in script.formats[0]["properties"]
|
||||
assert script.kinds == ["step", "step", "step", "final"]
|
||||
assert out.usage["calls"] == 4 and out.usage["input_tokens"] > 0
|
||||
assert script.tool_results()["c1"].startswith("【证据 Q1】")
|
||||
|
||||
|
||||
async def test_events_mark_the_queries_as_verification_and_nothing_goes_to_the_ledger(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1"), q(name, "SELEC broken", "c2")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "见 [[see:Q1]]。"})
|
||||
events: list[tuple] = []
|
||||
out = await run(name, emit=lambda kind, **data: events.append((kind, data)), run_id="run-x", node_id="check")
|
||||
ends = {d["call_id"]: d for k, d in events if k == EventType.TOOL_END}
|
||||
starts = [d for k, d in events if k == EventType.TOOL_START]
|
||||
assert [d["call_id"] for d in starts] == ["c1", "c2"], "每个工具调用一条 tool.start，和 tool.end 配对"
|
||||
assert all(d["role"] == "verify" for d in starts)
|
||||
end = ends["c1"]
|
||||
assert end["role"] == "verify" and end["query_artifact"] == out.queries[0]["artifact"] and end["alias"] == "Q1"
|
||||
# 被拒的语句照常发 tool.end（和 Agent 节点一样），运行面板按「SQL 被拒绝」开头标成失败；它没有查询快照
|
||||
assert ends["c2"]["role"] == "verify" and ends["c2"]["preview"].startswith("SQL 被拒绝")
|
||||
assert "query_artifact" not in ends["c2"] and "alias" not in ends["c2"]
|
||||
llm_ends = [d for k, d in events if k == EventType.LLM_END]
|
||||
assert len(llm_ends) == out.usage["calls"] and all(d["role"] == "verify" for d in llm_ends)
|
||||
# 工件归属这次运行、这个节点
|
||||
assert artifact_store.load(out.queries[0]["artifact"])["rows"] == [[230]]
|
||||
# 只交回结果，不交台账：没有任何 kind: query 的条目能流进 state["evidence"]
|
||||
assert "evidence" not in out.as_dict() and all("kind" not in qd for qd in out.queries)
|
||||
|
||||
|
||||
async def test_without_emit_nothing_is_announced(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "见 [[see:Q1]]。"})
|
||||
out = await run(name)
|
||||
assert out.result["value"] == 230, "质询接口这类不在节点里的调用不传 emit，照常工作"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 只读
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_writes_are_refused_on_a_writable_source(monkeypatch, shop):
|
||||
name, path = shop
|
||||
script = Script(monkeypatch, [q(name, "DELETE FROM orders WHERE is_deleted = 1", "w1"),
|
||||
q(name, "WITH k AS (SELECT 1) UPDATE orders SET amount = 0", "w2"),
|
||||
q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致",
|
||||
"rationale": "见 [[see:Q1]]。"})
|
||||
out = await run(name)
|
||||
seen = script.tool_results()
|
||||
assert seen["w1"].startswith("SQL 被拒绝：") and seen["w2"].startswith("SQL 被拒绝：")
|
||||
assert counts(path) == (4, 280), "数据原样不动"
|
||||
assert [qd["alias"] for qd in out.queries] == ["Q1"], "被拒的语句不编号"
|
||||
assert out.result["value"] == 230
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 上限
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_the_query_limit_stops_the_loop_and_settles(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
script = Script(monkeypatch, [q(name, "SELECT COUNT(*) AS n FROM orders", "c1"),
|
||||
[(f"db_query__{name}", {"sql": "SELECT MAX(amount) AS m FROM orders"}, "c2"),
|
||||
(f"db_query__{name}", {"sql": VALID_SUM}, "c3")],
|
||||
q(name, "SELECT MIN(amount) AS m FROM orders", "c4")],
|
||||
{"value": {"alias": "Q2", "row": 0, "column": "m"}, "verdict": "无法判断",
|
||||
"rationale": "查询次数用完了，最大一笔是 [[v:Q2.r0.m]]。"})
|
||||
out = await run(name, max_queries=2)
|
||||
assert out.limit == "max_queries"
|
||||
assert [qd["alias"] for qd in out.queries] == ["Q1", "Q2"]
|
||||
assert all("is_deleted" not in qd["sql"] for qd in out.queries), "同一轮里超出上限的那条没有执行"
|
||||
assert script.kinds == ["step", "step", "final"], "到顶就收尾，不再给模型下一轮"
|
||||
assert out.result["value"] == 120
|
||||
final_prompt = str(script.calls[-1][1][-1].content)
|
||||
assert "已达上限（最多查询 2 次）" in final_prompt, "收尾时告诉模型为什么不能再查"
|
||||
|
||||
|
||||
async def test_the_time_limit_stops_the_loop_and_still_settles(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
monkeypatch.setattr(verify_agent, "SETTLE_GRACE_S", 2.0)
|
||||
endless = [q(name, f"SELECT COUNT(*) AS n FROM orders WHERE id > {i}", f"c{i}") for i in range(50)]
|
||||
script = Script(monkeypatch, endless,
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "n"}, "verdict": "无法判断", "rationale": "超时了。"},
|
||||
latency=0.15)
|
||||
started = time.monotonic()
|
||||
out = await run(name, timeout_s=0.5, max_queries=None)
|
||||
assert out.limit == "timeout_s"
|
||||
assert script.kinds[-1] == "final", "到点也要收尾交结果"
|
||||
assert time.monotonic() - started < 0.5 + 2.0 + 1.0
|
||||
assert 1 <= len(out.queries) < 10
|
||||
|
||||
|
||||
async def test_a_model_call_that_runs_past_the_deadline_is_cut(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
monkeypatch.setattr(verify_agent, "SETTLE_GRACE_S", 0.3)
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": None, "verdict": "无法判断", "rationale": "没查到。"}, latency=1.0)
|
||||
started = time.monotonic()
|
||||
out = await run(name, timeout_s=0.2)
|
||||
assert out.limit == "timeout_s" and out.queries == []
|
||||
assert time.monotonic() - started < 1.5, "超时的那次调用当场放弃，收尾那次也只等宽限的时长"
|
||||
assert out.result is None and out.error, "收尾那次也没赶上：照实说没交成结果"
|
||||
|
||||
|
||||
async def test_a_query_that_hits_its_own_time_limit_is_not_the_overall_deadline(monkeypatch, shop):
|
||||
"""一条查询超过它自己的时限（ToolTimeout，也是一种 TimeoutError）：这条查询作罢、原因交回模型让它改，核查
|
||||
照常往下走。不能当成整个核查的时长到了，把还剩的时间全扔掉。"""
|
||||
name, _ = shop
|
||||
real = verify_agent.run_bounded
|
||||
seen = {"n": 0}
|
||||
|
||||
async def stuck_once(call, limit, tool_name):
|
||||
seen["n"] += 1
|
||||
if seen["n"] == 1:
|
||||
call.close()
|
||||
raise ToolTimeout(tool_name, Limit(seconds=30.0, enforce=True, kind="query"))
|
||||
return await real(call, limit, tool_name)
|
||||
|
||||
monkeypatch.setattr(verify_agent, "run_bounded", stuck_once)
|
||||
script = Script(monkeypatch, [q(name, "SELECT COUNT(*) AS n FROM orders", "c1"), q(name, VALID_SUM, "c2")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"})
|
||||
events: list[tuple] = []
|
||||
out = await run(name, timeout_s=600, emit=lambda kind, **data: events.append((kind, data)))
|
||||
assert out.limit is None, "整个核查的时长远没到"
|
||||
assert script.tool_results()["c1"].startswith("查询超过 30 秒没有返回")
|
||||
assert out.result["value"] == 230 and [qd["alias"] for qd in out.queries] == ["Q1"]
|
||||
[error] = [d for k, d in events if k == EventType.TOOL_ERROR]
|
||||
assert error["call_id"] == "c1" and "核查的时长上限" not in error["error"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def priced(monkeypatch):
|
||||
"""每次调用按 0.01 美元估算、按 0.01 美元结算：金额上限的算术好核对。"""
|
||||
monkeypatch.setattr(pricing, "find_pricing", lambda model: {"input": 1.0, "output": 1.0})
|
||||
monkeypatch.setattr(pricing, "estimate_cost", lambda model, inp, out: 0.01)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def by_tokens(monkeypatch):
|
||||
"""按 token 数估算（每个 token 0.000001 美元）：消息越长估得越多，收尾时交给模型的查询结果也要算钱。"""
|
||||
monkeypatch.setattr(pricing, "find_pricing", lambda model: {"input": 1.0, "output": 1.0})
|
||||
monkeypatch.setattr(pricing, "estimate_cost", lambda model, inp, out: (inp + out) / 1_000_000)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def wordy(tmp_path):
|
||||
"""200 条留言，每条一大段中文：查出来交给收尾那次调用时很占 token。返回源名。"""
|
||||
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
||||
db = sqlite3.connect(path)
|
||||
db.execute("CREATE TABLE notes (id INTEGER PRIMARY KEY, note TEXT)")
|
||||
db.executemany("INSERT INTO notes VALUES (?, ?)", [(i, f"第{i}条留言" + "虚构的长段落文字" * 50) for i in range(200)])
|
||||
db.commit()
|
||||
db.close()
|
||||
name = f"notes_{uuid.uuid4().hex[:8]}"
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True)
|
||||
session.add(row)
|
||||
await session.flush()
|
||||
row.schema_cache = await introspect.introspect(row)
|
||||
await session.commit()
|
||||
source_id = row.id
|
||||
yield name
|
||||
await engines.invalidate(source_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cap", [0.003, 0.004, 0.005, 0.02, 0.05])
|
||||
async def test_a_query_never_runs_without_money_left_to_settle_it(monkeypatch, wordy, by_tokens, cap):
|
||||
"""收尾那次调用要把查到的结果交给模型：结果越长越贵。查询之前就要按收尾最贵的样子把钱留好，不然查完了却
|
||||
付不起收尾，查到的结果白查。"""
|
||||
count = {"type": "object", "properties": {"count": {"type": "integer"}}}
|
||||
Script(monkeypatch, [q(wordy, "SELECT id, note FROM notes ORDER BY id", "c1")],
|
||||
{"count": {"alias": "Q1", "row": 0, "column": "id"}, "rationale": "x"}, model_id="priced-model")
|
||||
out = await run(wordy, output_schema=count, max_cost_usd=cap, max_queries=1)
|
||||
assert out.error is None and out.result is not None, "不管上限多紧，都交得出结果"
|
||||
assert out.usage["cost_usd"] <= cap, "收尾那次也在上限以内：钱是查询之前就留好的，不是收尾时硬超出去"
|
||||
if out.queries:
|
||||
assert out.result["count"] == 0, "查了就用上了"
|
||||
else:
|
||||
assert out.limit == "max_cost_usd", "钱不够查询加收尾，就不查"
|
||||
|
||||
|
||||
async def test_a_final_answer_budgeted_in_the_loop_is_not_refused_by_the_daily_cap(monkeypatch, shop, priced):
|
||||
"""查询之后、收尾之前，别的核查把当天的余额用光了（几个核查并行）：收尾那次已经在循环里留过钱，照常交结果，
|
||||
查到的结果不白查。"""
|
||||
name, _ = shop
|
||||
|
||||
async def crowded(kind, **data):
|
||||
if kind == EventType.TOOL_END:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, spend.SPEND_KEY)
|
||||
row.value = {**row.value, "usd": 1.999}
|
||||
await session.commit()
|
||||
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1"), q(name, "SELECT COUNT(*) AS n FROM orders", "c2")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"},
|
||||
model_id="priced-model")
|
||||
out = await run(name, daily_max_usd=2.0, emit=crowded)
|
||||
assert out.limit == "daily_max_usd", "下一轮查询付不起了"
|
||||
assert out.error is None and out.result["value"] == 230
|
||||
|
||||
|
||||
async def test_the_daily_cap_comes_from_the_settings_unless_given(monkeypatch, shop, priced):
|
||||
"""不传每日上限就按设置里的「核查每日金额上限」：节点、质询接口不必各自记得去读；明确传 None 才是不限。"""
|
||||
name, _ = shop
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": spend._today(), "usd": 0.995, "calls": 3}))
|
||||
session.add(Setting(key="verify", value={"daily_max_usd": 1.0}))
|
||||
await session.commit()
|
||||
final = {"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"}
|
||||
script = Script(monkeypatch, [q(name, VALID_SUM, "c1")], final, model_id="priced-model")
|
||||
out = await run(name)
|
||||
assert out.limit == "daily_max_usd" and script.calls == [] and out.result is None
|
||||
assert "$1.00" in out.error
|
||||
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], final, model_id="priced-model")
|
||||
unlimited = await run(name, daily_max_usd=None)
|
||||
assert unlimited.limit is None and unlimited.result["value"] == 230
|
||||
|
||||
|
||||
async def test_the_cost_limit_keeps_room_for_the_final_answer(monkeypatch, shop, priced):
|
||||
name, _ = shop
|
||||
endless = [q(name, f"SELECT COUNT(*) AS n FROM orders WHERE id > {i}", f"c{i}") for i in range(10)]
|
||||
script = Script(monkeypatch, endless,
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "n"}, "verdict": "无法判断", "rationale": "钱用完了。"},
|
||||
model_id="priced-model")
|
||||
out = await run(name, max_cost_usd=0.035, max_queries=None)
|
||||
assert out.limit == "max_cost_usd"
|
||||
assert script.kinds == ["step", "step", "final"], "每一轮都给收尾那次留出钱：两轮之后再查就付不起收尾了"
|
||||
assert out.usage["cost_usd"] == pytest.approx(0.03) and out.usage["cost_usd"] <= 0.035
|
||||
assert out.result["value"] == 4
|
||||
|
||||
|
||||
async def test_the_daily_cap_is_shared_by_every_verification(monkeypatch, shop, priced):
|
||||
name, _ = shop
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": spend._today(), "usd": 1.995, "calls": 9}))
|
||||
await session.commit()
|
||||
script = Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"},
|
||||
model_id="priced-model")
|
||||
out = await run(name, daily_max_usd=2.0)
|
||||
assert out.limit == "daily_max_usd" and script.calls == [] and out.usage["calls"] == 0
|
||||
assert out.result is None and out.error
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(1.995), "没调用就不记账"
|
||||
|
||||
|
||||
async def test_spending_goes_into_the_verify_ledger(monkeypatch, shop, priced):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "见 [[see:Q1]]。"},
|
||||
model_id="priced-model")
|
||||
out = await run(name, daily_max_usd=2.0)
|
||||
day = await spend.daily_spend()
|
||||
assert day["calls"] == out.usage["calls"] == 3 and day["usd"] == pytest.approx(out.usage["cost_usd"])
|
||||
|
||||
|
||||
async def test_an_unpriced_model_makes_the_cost_limits_moot_and_says_so(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "见 [[see:Q1]]。"},
|
||||
model_id="house-model-unpriced")
|
||||
events: list[tuple] = []
|
||||
out = await run(name, max_cost_usd=0.0001, daily_max_usd=0.0001,
|
||||
emit=lambda kind, **data: events.append((kind, data)))
|
||||
assert out.limit is None and out.result["value"] == 230, "估不出金额：金额上限拦不住，也不该拦"
|
||||
assert any("不在价格目录中" in n and "金额上限" in n for n in out.notes)
|
||||
assert any(k == EventType.LOG and d.get("code") == "verify_unpriced" for k, d in events)
|
||||
assert out.usage["unpriced_calls"] == out.usage["calls"] == 3
|
||||
assert (await spend.daily_spend())["unpriced_calls"] == 3
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 数值只认格子
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_numbers_the_model_writes_itself_are_ignored(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": 999, "verdict": "一致", "rationale": "合计是 [[v:Q1.r0.total]]。"})
|
||||
out = await run(name)
|
||||
assert out.result["value"] is None, "直接写的数字不采信，也不兜底成 0"
|
||||
[cell] = out.cells
|
||||
assert cell["status"] == "unresolved" and cell["model_value"] == 999 and "格子引用" in cell["reason"]
|
||||
|
||||
|
||||
async def test_a_value_tucked_into_the_reference_is_ignored_too(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total", "value": 999}, "verdict": "一致", "rationale": "x"})
|
||||
out = await run(name)
|
||||
assert out.result["value"] == 230, "引用里夹带的值不看，按快照取"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ref,reason", [
|
||||
({"alias": "Q7", "row": 0, "column": "total"}, "Q7"),
|
||||
({"alias": "Q1", "row": 5, "column": "total"}, "第 5 行"),
|
||||
({"alias": "Q1", "row": 0, "column": "amount"}, "amount"),
|
||||
])
|
||||
async def test_a_reference_that_points_nowhere_is_unresolved(monkeypatch, shop, ref, reason):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], {"value": ref, "verdict": "一致", "rationale": "x"})
|
||||
out = await run(name)
|
||||
assert out.result["value"] is None
|
||||
assert out.cells[0]["status"] == "unresolved" and reason in out.cells[0]["reason"]
|
||||
|
||||
|
||||
async def test_bare_numbers_in_the_rationale_are_flagged(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致",
|
||||
"rationale": "合计是 230，比原流程多 3 成。"})
|
||||
out = await run(name)
|
||||
flagged = [v["text"] for v in out.rationale["violations"] if v["code"] == "uncited_number"]
|
||||
assert flagged == ["230", "3"] and out.rationale["stats"]["uncited_numbers"] == 2
|
||||
bare = [s for _, u in [(b, u) for b in out.rationale["blocks"] for u in b["units"]] for s in u["segments"]
|
||||
if s.get("issue") == "uncited_number"]
|
||||
assert [s["text"] for s in bare] == ["230", "3"], "界面要在这几个字底下画「无出处」"
|
||||
|
||||
|
||||
async def test_nested_numbers_and_arrays_are_references_too(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
schema = {"type": "object", "properties": {
|
||||
"totals": {"type": "array", "items": {"type": "number"}},
|
||||
"detail": {"type": "object", "properties": {"count": {"type": "integer"}, "note": {"type": "string"}}},
|
||||
}}
|
||||
Script(monkeypatch, [q(name, "SELECT id, amount FROM orders ORDER BY id", "c1")],
|
||||
{"totals": [{"alias": "Q1", "row": 0, "column": "amount"}, {"alias": "Q1", "row": 3, "column": "amount"}],
|
||||
"detail": {"count": {"alias": "Q1", "row": 3, "column": "id"}, "note": "四笔"}, "rationale": "x"})
|
||||
out = await run(name, output_schema=schema)
|
||||
assert out.result["totals"] == [120, 30] and out.result["detail"] == {"count": 4, "note": "四笔"}
|
||||
assert [c["path"] for c in out.cells] == ["totals[0]", "totals[1]", "detail.count"]
|
||||
|
||||
|
||||
#: Pydantic 这类工具生成的结构：可为 null 的数写成 anyOf，嵌套的对象写成 $ref 到 $defs，带说明的引用包一层 allOf，
|
||||
#: Optional[对象] 是 anyOf [$ref, null]
|
||||
PYDANTIC_STYLE = {
|
||||
"type": "object",
|
||||
"$defs": {"Region": {"type": "object", "properties": {"name": {"type": "string"}, "total": {"type": "number"}}}},
|
||||
"properties": {
|
||||
"value": {"anyOf": [{"type": "number"}, {"type": "null"}], "description": "未作废订单的金额合计"},
|
||||
"count": {"type": ["integer", "null"]},
|
||||
"top": {"$ref": "#/$defs/Region"},
|
||||
"backup": {"anyOf": [{"$ref": "#/$defs/Region"}, {"type": "null"}]},
|
||||
"wrapped": {"allOf": [{"$ref": "#/$defs/Region"}], "description": "带说明的引用"},
|
||||
"regions": {"type": "array", "items": {"$ref": "#/$defs/Region"}},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_every_way_of_writing_a_number_becomes_a_reference():
|
||||
"""anyOf、type 列表、$ref、allOf 里的数也换成格子引用：认不出这些写法，模型在那里写的数就绕过了格子这一关。"""
|
||||
shaped = result_schema(PYDANTIC_STYLE)
|
||||
assert "$ref" not in json.dumps(shaped) and "$defs" not in shaped
|
||||
|
||||
def is_ref(node):
|
||||
return node["anyOf"][0]["required"] == ["alias", "row", "column"] and node["anyOf"][1] == {"type": "null"}
|
||||
|
||||
props = shaped["properties"]
|
||||
assert is_ref(props["value"]) and "未作废订单的金额合计" in props["value"]["description"]
|
||||
assert is_ref(props["count"])
|
||||
assert is_ref(props["top"]["properties"]["total"]) and props["top"]["properties"]["name"] == {"type": "string"}
|
||||
assert is_ref(props["wrapped"]["properties"]["total"])
|
||||
assert is_ref(props["regions"]["items"]["properties"]["total"])
|
||||
[region, null] = props["backup"]["anyOf"]
|
||||
assert is_ref(region["properties"]["total"]) and null == {"type": "null"}
|
||||
|
||||
|
||||
async def test_numbers_written_directly_in_any_of_those_shapes_are_ignored(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": 999, "count": 7, "top": {"name": "分区甲", "total": 999}, "backup": {"name": "分区乙", "total": 1},
|
||||
"wrapped": {"name": "分区丙", "total": 2}, "regions": [{"name": "分区丁", "total": 3}], "rationale": "x"})
|
||||
out = await run(name, output_schema=PYDANTIC_STYLE)
|
||||
assert out.result == {"value": None, "count": None, "top": {"name": "分区甲", "total": None},
|
||||
"backup": {"name": "分区乙", "total": None}, "wrapped": {"name": "分区丙", "total": None},
|
||||
"regions": [{"name": "分区丁", "total": None}], "rationale": "x"}
|
||||
assert {c["path"]: c["model_value"] for c in out.cells} == {
|
||||
"value": 999, "count": 7, "top.total": 999, "backup.total": 1, "wrapped.total": 2, "regions[0].total": 3}
|
||||
assert all(c["status"] == "unresolved" and "没有给出格子引用" in c["reason"] for c in out.cells), \
|
||||
"都认成了数值字段（该写格子引用），不是碰巧被当成认不出的结构丢掉"
|
||||
|
||||
|
||||
async def test_references_in_those_shapes_are_resolved(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
ref = {"alias": "Q1", "row": 0, "column": "total"}
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": ref, "count": None, "top": {"name": "分区甲", "total": ref}, "backup": None,
|
||||
"wrapped": {"name": "分区丙", "total": ref}, "regions": [{"name": "分区丁", "total": ref}], "rationale": "x"})
|
||||
out = await run(name, output_schema=PYDANTIC_STYLE)
|
||||
assert out.result["value"] == 230 and out.result["top"]["total"] == 230 and out.result["wrapped"]["total"] == 230
|
||||
assert out.result["regions"] == [{"name": "分区丁", "total": 230}]
|
||||
assert out.result["count"] is None and out.result["backup"] is None, "可以为 null 的照实是 null"
|
||||
|
||||
|
||||
async def test_numbers_under_a_shape_that_is_not_recognised_are_dropped(monkeypatch, shop):
|
||||
"""认不出的结构（没写类型、没写 items 的数组、没写属性的对象）、文本字段里冒出来的数：一律不采信，照记下来。
|
||||
码值（enum）里的数是调用方给的选项，不是查出来的数，照收。"""
|
||||
name, _ = shop
|
||||
schema = {"type": "object", "properties": {
|
||||
"anything": {}, "notes": {"type": "array"}, "label": {"type": "string"}, "level": {"enum": [1, 2, 3]},
|
||||
"extra": {"type": "object"}}}
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"anything": 999, "notes": ["见 Q1", 120], "label": 5, "level": 2,
|
||||
"extra": {"note": "按支付时间算", "share": 0.3}, "rationale": "x"})
|
||||
out = await run(name, output_schema=schema)
|
||||
assert out.result == {"anything": None, "notes": ["见 Q1", None], "label": None, "level": 2,
|
||||
"extra": {"note": "按支付时间算", "share": None}, "rationale": "x"}
|
||||
assert {c["path"]: c["model_value"] for c in out.cells} == {"anything": 999, "notes[1]": 120, "label": 5,
|
||||
"extra.share": 0.3}
|
||||
assert all(c["status"] == "unresolved" for c in out.cells)
|
||||
|
||||
|
||||
async def test_references_written_as_json_text_are_read(monkeypatch, shop):
|
||||
"""有的模型把嵌套的对象、数组写成一段 JSON 文本：照样认（裁判也是这么认的）。还不是数组的照实记下格式不对。"""
|
||||
name, _ = shop
|
||||
schema = {"type": "object", "properties": {
|
||||
"value": {"type": "number"}, "totals": {"type": "array", "items": {"type": "number"}},
|
||||
"broken": {"type": "array", "items": {"type": "number"}}}}
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": '{"alias": "Q1", "row": 0, "column": "total"}',
|
||||
"totals": '[{"alias": "Q1", "row": 0, "column": "total"}]', "broken": "没有这一项", "rationale": "x"})
|
||||
out = await run(name, output_schema=schema)
|
||||
assert out.result == {"value": 230, "totals": [230], "broken": None, "rationale": "x"}
|
||||
broken = next(c for c in out.cells if c["path"] == "broken")
|
||||
assert broken["status"] == "unresolved" and "数组" in broken["reason"]
|
||||
|
||||
|
||||
async def test_a_number_field_pointing_at_a_text_cell_is_unresolved(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, "SELECT paid_at, amount FROM orders ORDER BY id", "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "paid_at"}, "verdict": "一致", "rationale": "x"})
|
||||
out = await run(name)
|
||||
assert out.result["value"] is None, "数值字段只收数：日期文本拿去和原流程的数比，比较会出错"
|
||||
[cell] = out.cells
|
||||
assert cell["status"] == "unresolved" and "不是数值" in cell["reason"] and cell["ref"] == "Q1.r0.paid_at"
|
||||
assert cell["value"] == "2026-09-02"
|
||||
|
||||
|
||||
async def test_a_field_left_out_is_not_a_field_reported_missing(monkeypatch, shop):
|
||||
"""没交这个字段和交了 null（模型说查不到）是两回事：前者是没按结构交，后者才是 missing。"""
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], {"verdict": "无法判断", "rationale": "x"})
|
||||
left_out = await run(name)
|
||||
[cell] = left_out.cells
|
||||
assert cell["status"] == "unresolved" and "没有交" in cell["reason"]
|
||||
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], {"value": None, "verdict": "无法判断", "rationale": "x"})
|
||||
said_null = await run(name)
|
||||
assert said_null.cells == [{"path": "value", "status": "missing"}]
|
||||
|
||||
|
||||
async def test_an_answer_with_none_of_the_fields_is_an_error(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], {"total": 230, "rationale": "合计 230"})
|
||||
out = await run(name)
|
||||
assert out.result is None and out.cells == [] and "没有按结构交结果" in out.error
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 数据版本
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def book(rows: list[list]) -> bytes:
|
||||
from openpyxl import Workbook
|
||||
|
||||
wb = Workbook()
|
||||
ws = wb.active
|
||||
ws.title = "明细"
|
||||
for row in rows:
|
||||
ws.append(row)
|
||||
buf = io.BytesIO()
|
||||
wb.save(buf)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
async def test_pinned_data_versions_are_used(monkeypatch, tmp_path):
|
||||
data = tmp_path / "data"
|
||||
(data / "uploads").mkdir(parents=True)
|
||||
monkeypatch.setattr(settings, "data_dir", data)
|
||||
name = f"upload_{uuid.uuid4().hex[:8]}"
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(id=uuid.uuid4().hex, name=name, kind="sqlite", readonly=True, origin="upload")
|
||||
first = await table_versions.publish_upload(session, row, book([["区域", "数量"], ["分区甲", 1], ["分区乙", 2]]),
|
||||
"客流.xlsx", header_row=1, mixed="reject", raw_mode=False)
|
||||
source_id = row.id
|
||||
async with SessionLocal() as session:
|
||||
row = (await session.execute(select(DataSource).where(DataSource.name == name))).scalar_one()
|
||||
await table_versions.publish_upload(session, row, book([["区域", "数量"], ["分区甲", 60], ["分区乙", 40]]),
|
||||
"客流.xlsx", header_row=1, mixed="reject", raw_mode=False)
|
||||
total = {"type": "object", "properties": {"total": {"type": "number"}}}
|
||||
final = {"total": {"alias": "Q1", "row": 0, "column": "total"}, "rationale": "x"}
|
||||
sql = 'SELECT SUM("数量") AS total FROM "明细"'
|
||||
|
||||
Script(monkeypatch, [q(name, sql, "c1")], final)
|
||||
pinned = await run(name, output_schema=total,
|
||||
data_versions={source_id: {"snapshot": first.snapshot_id, "name": name}})
|
||||
assert pinned.result["total"] == 3, "按运行固定的那一版查"
|
||||
assert artifact_store.load(pinned.queries[0]["artifact"])["data_version"] == first.snapshot_id
|
||||
|
||||
Script(monkeypatch, [q(name, sql, "c1")], final)
|
||||
current = await run(name, output_schema=total, data_versions=Versions.CURRENT)
|
||||
assert current.result["total"] == 100
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 出错不抛
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_a_model_that_is_not_configured_is_reported_not_raised(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
|
||||
async def broken(session, spec):
|
||||
raise ProviderNotConfigured("模型接入「无」没有可用的模型", "请在「设置 → 模型接入」中添加", code="no_model")
|
||||
|
||||
monkeypatch.setattr(verify_agent, "get_chat_model", broken)
|
||||
out = await run(name)
|
||||
assert out.result is None and "核查模型不可用" in out.error and out.usage["calls"] == 0
|
||||
|
||||
|
||||
async def test_an_unknown_source_is_noted(monkeypatch, shop):
|
||||
name, _ = shop
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"})
|
||||
out = await run(name, sources=[name, "no_such_source"])
|
||||
assert out.result["value"] == 230
|
||||
assert any("no_such_source" in n for n in out.notes)
|
||||
|
||||
|
||||
async def test_a_final_answer_in_plain_text_json_is_still_read(monkeypatch, shop):
|
||||
"""结构化输出失败（模型不支持、给的不是合法 JSON）：退回普通调用，从回复里把 JSON 抠出来。"""
|
||||
name, _ = shop
|
||||
script = Script(monkeypatch, [q(name, VALID_SUM, "c1")], "不是 JSON")
|
||||
|
||||
original = MockChatModel._decide
|
||||
|
||||
def decide(model, messages):
|
||||
if not model.response_format and any("交结果" in str(m.content) for m in messages[-1:]):
|
||||
return AIMessage(content='好的：```json\n{"value": {"alias": "Q1", "row": 0, "column": "total"}, '
|
||||
'"verdict": "一致", "rationale": "x"}\n```')
|
||||
return original(model, messages)
|
||||
|
||||
monkeypatch.setattr(MockChatModel, "_decide", decide)
|
||||
out = await run(name)
|
||||
assert out.result["value"] == 230 and script.kinds.count("final") == 1
|
||||
|
||||
|
||||
def plain_reply(monkeypatch, reply):
|
||||
"""普通调用（退回的那一次）的回复：reply(最后一条消息的正文) → 回复正文。其余照剧本。"""
|
||||
original = MockChatModel._decide
|
||||
seen: list[str] = []
|
||||
|
||||
def decide(model, messages):
|
||||
text = str(messages[-1].content)
|
||||
if not model.response_format and "交结果" in text:
|
||||
seen.append(text)
|
||||
return AIMessage(content=reply(text))
|
||||
return original(model, messages)
|
||||
|
||||
monkeypatch.setattr(MockChatModel, "_decide", decide)
|
||||
return seen
|
||||
|
||||
|
||||
async def test_the_plain_fallback_is_told_which_fields_to_fill(monkeypatch, shop):
|
||||
"""结构化输出不成、退回普通调用时，结构不在请求里了：要写进消息，不然模型不知道字段叫什么，交回来的字段
|
||||
对不上，全成了「查不到」。"""
|
||||
name, _ = shop
|
||||
schema = {"type": "object", "properties": {
|
||||
"blind_total_amount": {"type": "number", "description": "盲算的金额合计"},
|
||||
"consistency_verdict": {"type": "string", "enum": ["一致", "不一致"]}}}
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")], "不是 JSON")
|
||||
|
||||
def reply(text):
|
||||
# 模型只能按消息里看得到的字段名交
|
||||
body = {"rationale": "x"}
|
||||
if "blind_total_amount" in text:
|
||||
body["blind_total_amount"] = {"alias": "Q1", "row": 0, "column": "total"}
|
||||
if "consistency_verdict" in text:
|
||||
body["consistency_verdict"] = "一致"
|
||||
return json.dumps(body, ensure_ascii=False)
|
||||
|
||||
seen = plain_reply(monkeypatch, reply)
|
||||
out = await run(name, output_schema=schema)
|
||||
assert len(seen) == 1 and "盲算的金额合计" in seen[0] and "alias" in seen[0]
|
||||
assert out.error is None and out.result == {"blind_total_amount": 230, "consistency_verdict": "一致",
|
||||
"rationale": "x"}
|
||||
|
||||
|
||||
async def test_a_model_without_structured_output_goes_straight_to_the_plain_call(monkeypatch, shop):
|
||||
"""模型不支持结构化输出（建的时候就报错）：直接走普通调用，不算一次调用，也不抛。"""
|
||||
name, _ = shop
|
||||
|
||||
def unsupported(self, schema, **kwargs):
|
||||
raise NotImplementedError("这个模型不支持结构化输出")
|
||||
|
||||
monkeypatch.setattr(MockChatModel, "with_structured_output", unsupported)
|
||||
script = Script(monkeypatch, [q(name, VALID_SUM, "c1")], "不会用到")
|
||||
plain_reply(monkeypatch, lambda text: '{"value": {"alias": "Q1", "row": 0, "column": "total"}, '
|
||||
'"verdict": "一致", "rationale": "x"}')
|
||||
out = await run(name)
|
||||
assert out.error is None and out.result["value"] == 230
|
||||
assert out.usage["calls"] == len(script.calls) + 1 == 3, "两轮查库加一次普通调用，建不起来的那次不算"
|
||||
|
||||
|
||||
async def test_a_failing_ledger_is_reported_not_raised(monkeypatch, shop):
|
||||
"""每日账写不进去（库被锁住）：照实交回，不抛。节点因此降成「核不了」，质询接口不会回 500。"""
|
||||
name, _ = shop
|
||||
|
||||
async def locked(*args, **kwargs):
|
||||
raise RuntimeError("database is locked")
|
||||
|
||||
monkeypatch.setattr(spend, "charge", locked)
|
||||
Script(monkeypatch, [q(name, VALID_SUM, "c1")],
|
||||
{"value": {"alias": "Q1", "row": 0, "column": "total"}, "verdict": "一致", "rationale": "x"})
|
||||
out = await run(name)
|
||||
assert out.result is None and out.error
|
||||
@@ -0,0 +1,205 @@
|
||||
"""核查设置（api/settings.py 的 verify 一组）：双盲复算、反方审稿、审批质询共用。
|
||||
|
||||
- 核查模型（接入、模型）：节点上写了用节点的，没写用设置里的，设置也没写就沿用证据裁判模型的取值链
|
||||
（裁判设置 > 助手模型 > 默认接入），哪一级写了就整组用哪一级，不跨级拼接入和模型
|
||||
- 三项各有开关和上限：双盲复算 {启用, 允许的相对误差, 查询次数, 时长}、反方审稿 {启用, 句数, 查询次数, 时长}、
|
||||
审批质询 {启用, 查询次数, 时长}；另有全局每日金额上限。上限存 null 表示不限，不会被默认值悄悄盖回去
|
||||
- 设置页写进来的值先查后写：有一项不对整个请求回 422，说明写设置页上的叫法，一项都不落库
|
||||
- 核查的每日计数是引擎记的账（engine/spend.py），不从设置接口进出；有一个只读接口给设置页看今天花了多少
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from sqlalchemy import delete
|
||||
|
||||
from app.api.settings import (
|
||||
VERIFY_DEFAULTS,
|
||||
VERIFY_GROUP,
|
||||
VERIFY_SETTING_LABEL,
|
||||
verify_model_spec,
|
||||
verify_settings,
|
||||
)
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Setting
|
||||
from app.engine import spend
|
||||
from app.engine.judge import JUDGE_DEFAULTS, MAX_TOKENS
|
||||
from app.main import app
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def clean():
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key.in_(
|
||||
[VERIFY_GROUP, spend.SPEND_KEY, "judge", "copilot"])))
|
||||
await session.commit()
|
||||
yield
|
||||
|
||||
|
||||
def client() -> AsyncClient:
|
||||
return AsyncClient(transport=ASGITransport(app=app), base_url="http://test")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 默认值与读写
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_the_defaults_are_the_agreed_ones():
|
||||
assert VERIFY_GROUP == "verify"
|
||||
assert VERIFY_DEFAULTS == {
|
||||
"provider": None, "model": None,
|
||||
"dual_check": {"enabled": True, "tolerance": 1e-6, "max_queries": 6, "timeout_s": 120},
|
||||
"challenge": {"enabled": True, "max_claims": 5, "max_queries": 4, "timeout_s": 90},
|
||||
"inquiry": {"enabled": True, "max_queries": 6, "timeout_s": 120},
|
||||
"daily_max_usd": 2.0,
|
||||
}
|
||||
|
||||
|
||||
async def test_settings_expose_the_verify_group_with_defaults():
|
||||
async with client() as c:
|
||||
got = (await c.get("/api/settings")).json()
|
||||
assert got["verify"] == VERIFY_DEFAULTS
|
||||
assert got["judge"] == JUDGE_DEFAULTS, "裁判一组不受影响"
|
||||
assert spend.SPEND_KEY not in got
|
||||
|
||||
|
||||
async def test_a_full_group_round_trips_and_null_stays_unlimited():
|
||||
value = {
|
||||
"provider": "house", "model": "verify-model",
|
||||
"dual_check": {"enabled": False, "tolerance": 0, "max_queries": None, "timeout_s": None},
|
||||
"challenge": {"enabled": True, "max_claims": None, "max_queries": 2, "timeout_s": 30.5},
|
||||
"inquiry": {"enabled": False, "max_queries": 1, "timeout_s": None},
|
||||
"daily_max_usd": None,
|
||||
}
|
||||
async with client() as c:
|
||||
r = await c.put("/api/settings", json={"values": {"verify": value}})
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["verify"] == value
|
||||
assert (await c.get("/api/settings")).json()["verify"] == value
|
||||
async with SessionLocal() as session:
|
||||
assert await verify_settings(session) == value
|
||||
|
||||
|
||||
async def test_a_partial_group_keeps_the_other_defaults_down_to_each_item():
|
||||
async with client() as c:
|
||||
r = await c.put("/api/settings", json={"values": {"verify": {"dual_check": {"enabled": False},
|
||||
"daily_max_usd": 5}}})
|
||||
assert r.status_code == 200, r.text
|
||||
got = (await c.get("/api/settings")).json()["verify"]
|
||||
assert got == {**VERIFY_DEFAULTS, "daily_max_usd": 5,
|
||||
"dual_check": {**VERIFY_DEFAULTS["dual_check"], "enabled": False}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key,value,label", [
|
||||
("daily_max_usd", -1, "daily_max_usd"),
|
||||
("daily_max_usd", "2", "daily_max_usd"),
|
||||
("model", 3, "model"),
|
||||
("dual_check", "on", "dual_check"),
|
||||
("dual_check", {"enabled": "yes"}, "dual_check.enabled"),
|
||||
("dual_check", {"tolerance": -0.1}, "dual_check.tolerance"),
|
||||
("dual_check", {"tolerance": "0.01"}, "dual_check.tolerance"),
|
||||
("dual_check", {"tolerance": None}, "dual_check.tolerance"),
|
||||
("dual_check", {"max_queries": 2.5}, "dual_check.max_queries"),
|
||||
("dual_check", {"max_queries": 0}, "dual_check.max_queries"),
|
||||
("challenge", {"max_claims": True}, "challenge.max_claims"),
|
||||
("challenge", {"timeout_s": -3}, "challenge.timeout_s"),
|
||||
("inquiry", {"timeout_s": 0}, "inquiry.timeout_s"),
|
||||
("inquiry", {"max_claims": 3}, "inquiry.max_claims"),
|
||||
("unknown_key", 1, "unknown_key"),
|
||||
])
|
||||
async def test_bad_values_are_refused_with_the_settings_page_wording(key, value, label):
|
||||
async with client() as c:
|
||||
r = await c.put("/api/settings", json={"values": {"verify": {key: value}, "ui": {"theme": "dark"}}})
|
||||
assert r.status_code == 422, r.text
|
||||
detail = r.json()["detail"]
|
||||
# 报错写设置页上的叫法（「双盲复算 · 允许的相对误差」），不写键名；认不出的键原样点名
|
||||
assert VERIFY_SETTING_LABEL.get(label, label.rpartition(".")[2]) in detail, detail
|
||||
got = (await c.get("/api/settings")).json()
|
||||
assert got["verify"] == VERIFY_DEFAULTS and got["ui"]["theme"] != "dark", "整个请求都没写进去"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["NaN", "Infinity"])
|
||||
async def test_non_finite_numbers_are_refused(raw):
|
||||
async with client() as c:
|
||||
r = await c.put("/api/settings", content='{"values": {"verify": {"dual_check": {"tolerance": %s}}}}' % raw,
|
||||
headers={"content-type": "application/json"})
|
||||
assert r.status_code == 422, r.text
|
||||
assert VERIFY_SETTING_LABEL["dual_check.tolerance"] in r.json()["detail"]
|
||||
|
||||
|
||||
async def test_bad_rows_already_in_the_database_fall_back_to_defaults():
|
||||
"""库里存坏了的值（旧版本、手改）读出来按默认，不让核查拿着负数上限、文字开关去查库。"""
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=VERIFY_GROUP, value={
|
||||
"model": " ", "daily_max_usd": -3,
|
||||
"dual_check": {"enabled": "yes", "tolerance": -1, "max_queries": None},
|
||||
"challenge": "broken", "inquiry": {"timeout_s": 15, "extra": 1}}))
|
||||
await session.commit()
|
||||
got = await verify_settings(session)
|
||||
assert got["model"] is None and got["daily_max_usd"] == VERIFY_DEFAULTS["daily_max_usd"]
|
||||
assert got["dual_check"] == {**VERIFY_DEFAULTS["dual_check"], "max_queries": None}
|
||||
assert got["challenge"] == VERIFY_DEFAULTS["challenge"]
|
||||
assert got["inquiry"] == {**VERIFY_DEFAULTS["inquiry"], "timeout_s": 15}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 每日计数
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_the_daily_counter_does_not_travel_through_the_settings_api():
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": spend._today(), "usd": 1.5, "calls": 3}))
|
||||
await session.commit()
|
||||
async with client() as c:
|
||||
assert spend.SPEND_KEY not in (await c.get("/api/settings")).json()
|
||||
await c.put("/api/settings", json={"values": {spend.SPEND_KEY: {"date": spend._today(), "usd": 0}}})
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, spend.SPEND_KEY)
|
||||
assert row.value["usd"] == 1.5, "设置页整组回写不能把今天的核查花费清零"
|
||||
|
||||
|
||||
async def test_todays_spend_is_readable_next_to_the_daily_cap(monkeypatch):
|
||||
monkeypatch.setattr(spend, "_today", lambda: "2026-10-02")
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": "2026-10-02", "usd": 0.42, "calls": 3,
|
||||
"unpriced_calls": 1}))
|
||||
await session.commit()
|
||||
async with client() as c:
|
||||
got = (await c.get("/api/settings/verify/spend")).json()
|
||||
assert got == {"date": "2026-10-02", "usd": 0.42, "calls": 3, "unpriced_calls": 1, "daily_max_usd": 2.0}
|
||||
await c.put("/api/settings", json={"values": {"verify": {"daily_max_usd": None}}})
|
||||
assert (await c.get("/api/settings/verify/spend")).json()["daily_max_usd"] is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 核查模型的取值链
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_the_verify_model_falls_back_along_the_judge_chain():
|
||||
async with SessionLocal() as session:
|
||||
spec = await verify_model_spec(session)
|
||||
assert (spec.provider, spec.model) == (None, None), "什么都没配：默认接入"
|
||||
assert spec.thinking == "off" and spec.max_tokens == MAX_TOKENS
|
||||
|
||||
session.add(Setting(key="copilot", value={"provider": "", "model": "copilot-model"}))
|
||||
await session.commit()
|
||||
assert (await verify_model_spec(session)).model == "copilot-model", "沿用裁判的取值链：助手模型"
|
||||
|
||||
session.add(Setting(key="judge", value={"provider": "house", "model": "judge-model"}))
|
||||
await session.commit()
|
||||
spec = await verify_model_spec(session)
|
||||
assert (spec.provider, spec.model) == ("house", "judge-model"), "沿用裁判的取值链：证据裁判模型"
|
||||
|
||||
session.add(Setting(key=VERIFY_GROUP, value={"model": "verify-model"}))
|
||||
await session.commit()
|
||||
spec = await verify_model_spec(session)
|
||||
assert (spec.provider, spec.model) == (None, "verify-model"), "核查设置写了模型就整组用它，不和裁判的接入拼"
|
||||
assert spec.thinking == "off" and spec.max_tokens == MAX_TOKENS
|
||||
|
||||
spec = await verify_model_spec(session, {"provider": "other", "model": "node-model", "tolerance": 0.1})
|
||||
assert (spec.provider, spec.model) == ("other", "node-model"), "节点上写了就按节点"
|
||||
spec = await verify_model_spec(session, {"provider": " ", "model": None})
|
||||
assert spec.model == "verify-model", "节点上留空等于没写"
|
||||
@@ -0,0 +1,107 @@
|
||||
"""核查花费账（engine/spend.py）：双盲复算、反方审稿、审批质询的模型调用按本地日期记一本每日账。
|
||||
|
||||
照结论句裁判的每日花费做（judge.py 的 SPEND_KEY、_charge、_settle），但是独立的一行设置：
|
||||
|
||||
- 调用前按估算预留，加上这笔会超过每日上限就不记、返回 None；上限是 None 表示不限
|
||||
- 调用后按实际用量找补差额；跨了天的不找补（预留记在前一天，今天从零算）
|
||||
- 价格目录里没有的模型估不出金额：只计次数（unpriced_calls），不计金额
|
||||
- 和裁判的账互不挤占：核查花掉的钱不算进裁判的每日上限，反过来也一样
|
||||
- 并发预留不会一起越过上限
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import delete
|
||||
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Setting
|
||||
from app.engine import spend
|
||||
from app.engine import judge as judge_mod
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def clean(monkeypatch):
|
||||
monkeypatch.setattr(spend, "_today", lambda: "2026-10-02")
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key.in_([spend.SPEND_KEY, judge_mod.SPEND_KEY])))
|
||||
await session.commit()
|
||||
yield
|
||||
|
||||
|
||||
async def stored() -> dict | None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, spend.SPEND_KEY)
|
||||
return row.value if row is not None else None
|
||||
|
||||
|
||||
async def test_a_call_within_the_cap_is_recorded():
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
assert await spend.charge(0.4, limit=1.0, priced=True) == "2026-10-02"
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.4, "calls": 1, "unpriced_calls": 0}
|
||||
|
||||
|
||||
async def test_a_call_that_would_cross_the_cap_is_refused_and_nothing_is_recorded():
|
||||
assert await spend.charge(0.9, limit=1.0, priced=True)
|
||||
assert await spend.charge(0.2, limit=1.0, priced=True) is None
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.9, "calls": 1, "unpriced_calls": 0}
|
||||
# 正好到上限的还能记：上限是「不超过」
|
||||
assert await spend.charge(0.1, limit=1.0, priced=True) == "2026-10-02"
|
||||
|
||||
|
||||
async def test_no_cap_means_unlimited():
|
||||
assert await spend.charge(50.0, limit=None, priced=True) == "2026-10-02"
|
||||
assert (await spend.daily_spend())["usd"] == 50.0
|
||||
|
||||
|
||||
async def test_unpriced_calls_are_counted_without_money():
|
||||
assert await spend.charge(0.0, limit=0.5, priced=False) == "2026-10-02"
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 1, "unpriced_calls": 1}
|
||||
|
||||
|
||||
async def test_yesterdays_spend_does_not_count_today():
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": "2026-10-01", "usd": 1.9, "calls": 7,
|
||||
"unpriced_calls": 2}))
|
||||
await session.commit()
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
assert await spend.charge(1.5, limit=2.0, priced=True) == "2026-10-02"
|
||||
assert await stored() == {"date": "2026-10-02", "usd": 1.5, "calls": 1, "unpriced_calls": 0}
|
||||
|
||||
|
||||
async def test_settle_corrects_the_reservation_and_never_goes_below_zero():
|
||||
day = await spend.charge(0.5, limit=None, priced=True)
|
||||
await spend.settle(day, -0.2)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.3)
|
||||
await spend.settle(day, 0.15)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.45)
|
||||
await spend.settle(day, -5)
|
||||
assert (await spend.daily_spend())["usd"] == 0.0
|
||||
assert (await spend.daily_spend())["calls"] == 1, "找补只改金额，不改次数"
|
||||
|
||||
|
||||
async def test_settling_after_midnight_does_not_touch_the_new_day(monkeypatch):
|
||||
day = await spend.charge(0.5, limit=None, priced=True)
|
||||
monkeypatch.setattr(spend, "_today", lambda: "2026-10-03")
|
||||
assert await spend.charge(0.2, limit=None, priced=True) == "2026-10-03"
|
||||
await spend.settle(day, -0.4)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.2), "预留记在前一天，今天从零算"
|
||||
|
||||
|
||||
async def test_the_verify_ledger_is_separate_from_the_judges(monkeypatch):
|
||||
monkeypatch.setattr(judge_mod, "_today", lambda: "2026-10-02")
|
||||
assert spend.SPEND_KEY != judge_mod.SPEND_KEY
|
||||
await spend.charge(1.2, limit=2.0, priced=True)
|
||||
assert (await judge_mod.daily_spend())["usd"] == 0.0, "核查花的钱不进裁判的账"
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=judge_mod.SPEND_KEY, value={"date": "2026-10-02", "usd": 1.9, "calls": 4}))
|
||||
await session.commit()
|
||||
assert await spend.charge(0.7, limit=2.0, priced=True) == "2026-10-02", "裁判花的钱也不挤占核查的上限"
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(1.9)
|
||||
|
||||
|
||||
async def test_concurrent_reservations_never_cross_the_cap_together():
|
||||
results = await asyncio.gather(*(spend.charge(0.3, limit=1.0, priced=True) for _ in range(10)))
|
||||
assert sum(1 for r in results if r) == 3
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.9)
|
||||
Reference in New Issue
Block a user