diff --git a/backend/app/api/settings.py b/backend/app/api/settings.py index c0e941ff..a8259092 100644 --- a/backend/app/api/settings.py +++ b/backend/app/api/settings.py @@ -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]: """给设置页顶部展示的运行环境概览。""" diff --git a/backend/app/data/sql_semantics.py b/backend/app/data/sql_semantics.py new file mode 100644 index 00000000..613f5e41 --- /dev/null +++ b/backend/app/data/sql_semantics.py @@ -0,0 +1,1546 @@ +"""SQL 语义：两条 SQL 算的是不是一回事，不是的话差在哪。 + +双盲复算要回答「盲算和原流程的 SQL 差在哪」（A 按支付时间筛、B 按下单时间筛；A 排除了作废的行、B 没有），运行 +对比要回答「这一期和上一期的 SQL 是不是同一个模板」（定时运行每次的日期不同，不能算模板变了），反方审稿要看自己 +的查询和原查询的口径是否一致。三处共用这里： + +- normalize_sql：字面量换成占位符之后的规范文本，加上字面量本身（按规范文本里的顺序）。两条 SQL 只有字面量不同， + 规范文本就相同；关键字大小写、空白、注释、不加引号的名字的大小写都不算不同。 +- facets：把一条查询拆成几个方面——用到的表、连接（类型、关联进来的表、ON 条件）、WHERE 的各个合取式、其中的 + 时间条件（列、运算符、取值）、分组、聚合（函数、参数、是否去重）、HAVING、取哪几行（LIMIT、OFFSET、取前几行 + 时的排序）。别名一律换成真实表名（条件里套着的子查询也是），等式、比较的两边排好顺序：o.paid_at >= '2026-09-01' + 和 '2026-09-01' <= orders.paid_at 是同一个条件。 +- diff_facets：逐个方面比两份 facets，每处不同给一句中文说明。 + +**逐段收集。** CTE、子查询里的条件、聚合和取行也算进来：盲算可能把有效记录条件写在 CTE 里，原流程写在外层，按 +「整条语句里有哪些条件」比才不会把同一个口径报成不同。逗号连接（以及 INNER JOIN 之外又写在 WHERE 里的等值关联） +写在 WHERE 里的列等式算作关联条件，不算筛选。条件里套着的子查询（EXISTS、IN (SELECT …)）既单算一段，也按它自己 +那一段的来源整个写进外层那个条件：别名在子查询里定义，拿外层的来源去认，同一个条件只因别名不同就会被报成不同。 + +**各自起的名字不进比较。** 表的别名换成真实表名；CTE、FROM 里的子查询没有真实表名可换，比较时按它在语句里的位置 +编号（@1、@2…，从外往里、按第一次用到的先后），给人看时照写名字；条件里的子查询（IN、EXISTS、标量子查询）选出来 +的列名不影响取值，比较时去掉。这些名字都是各自起的，双盲复算的盲算更是：按名字比，同一个口径会被说成换了口径。 +编号只看位置，一边多出一个排在前面的 CTE、子查询，后面的编号就错开了：diff_facets 先按编号比，比不上的再按不写 +编号的写法一对一地配，配上的不算不同。仍按名字比的只有 CTE、子查询选出来、外层按名字引用的列（一边叫 total、 +一边叫 sum_amt，算两列）。 + +**和 sqlcheck 一样宁可漏报。** 解析不了（语法错、方言特性、多条语句、不是查询）就带 parse_error 返回，从不抛 +异常。没加表前缀、又有好几张表可能有这一列的，照写列名（给了表结构才去对）。只读语法树，不执行 SQL。 +""" +from __future__ import annotations + +import logging +import re +from collections import deque +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass, fields, is_dataclass +from itertools import zip_longest +from typing import Any + +from app.data import catalog +from app.data.names import name_key +from app.data.sqlcheck import _bare_column, _conjuncts, _parents, _unparen, dialect_of + +logger = logging.getLogger(__name__) + +#: 规范文本里字面量的占位符 +PLACEHOLDER = "?" +#: 太长的 SQL 不做语法分析（normalize_sql 退回逐词规范化，facets 带 parse_error）：解析耗时随长度增长 +MAX_SQL_CHARS = 100_000 +#: 每段 SELECT 在语法树上记段号的 meta 键。条件、表达式要复制一份再改写，复制出来的子查询靠段号找回原来那一段 +#: 算好的来源（别名、外层）：复制品没有父节点，自己找不到外层 +_SEGMENT = "sql_semantics_segment" +#: 派生来源（引用 CTE 的表和 CTE 的定义、FROM / JOIN 里的子查询）在语法树上记比较用编号的 meta 键（见 +#: _Statement._number_derived）。复制出来的条件里套着子查询时，靠它把 CTE 名、子查询的别名换成编号 +_DERIVED = "sql_semantics_derived" +#: 比较用的写法里不写编号时，派生来源一律写成这个（见 Part.anon） +_ANON = "@" + + +# ========================================================================== +# 解析 +# ========================================================================== + + +def _read(dialect: str | None) -> str | None: + """数据源类型（sqlite / postgresql / mysql …）或 sqlglot 的方言名 → 解析用的方言。认不出的按 sqlglot 的通用读法。""" + found = dialect_of(dialect) + if found or not dialect: + return found + from sqlglot.dialects.dialect import Dialect + + name = dialect.strip().lower() + try: + Dialect.get_or_raise(name) + except Exception: # noqa: BLE001 - 不认识的方言名：按通用读法 + return None + return name + + +def _parse_problem(error: Exception) -> str: + """sqlglot 的报错 → 一句给人看的话。原文是英文、还带终端的高亮控制符，只取出错的位置。""" + errors = getattr(error, "errors", None) or [] + first = errors[0] if errors and isinstance(errors[0], dict) else {} + line, col = first.get("line"), first.get("col") + if isinstance(line, int) and isinstance(col, int): + return f"SQL 无法解析（第 {line} 行第 {col} 列附近）" + return "SQL 无法解析" + + +def _parse(sql: Any, read: str | None) -> tuple[Any, str | None]: + """(语法树, None) 或 (None, 为什么分析不了)。只认单条查询。""" + import sqlglot + from sqlglot import exp + + if not isinstance(sql, str) or not sql.strip(): + return None, "SQL 为空" + text = sql.strip() + if len(text) > MAX_SQL_CHARS: + return None, f"SQL 超过 {MAX_SQL_CHARS} 个字符，没有做语法分析" + try: + statements = [s for s in sqlglot.parse(text, read=read) if s is not None] + except Exception as e: # noqa: BLE001 - 语法错、方言特性、sqlglot 不支持的写法 + return None, _parse_problem(e) + if not statements: + return None, "SQL 为空" + if len(statements) > 1: + return None, f"一次只能分析一条 SQL，这里有 {len(statements)} 条" + if not isinstance(statements[0], exp.Query): + return None, "不是查询语句（只分析 SELECT、WITH 开头的查询）" + return statements[0], None + + +# ========================================================================== +# 规范化：字面量换成占位符 +# ========================================================================== + + +@dataclass(frozen=True) +class NormalizedSql: + """规范文本和字面量。text 相同就是同一个模板，literals 是这一次填进去的值（SQL 写法，字符串带引号）。 + + parse_error 不为空时 text 是逐词规范化的结果（关键字转大写、不加引号的名字转小写、字面量换成占位符）： + 同一个模板只改了字面量的两条 SQL 照样得到同一个 text，只是空白、括号之外的写法差异认不出来。 + """ + + text: str + literals: tuple[str, ...] = () + parse_error: str | None = None + + def same_shape(self, other: "NormalizedSql") -> bool: + """两条 SQL 是不是同一个模板（只有字面量不同也算）。""" + return self.text == other.text + + def as_dict(self) -> dict[str, Any]: + out: dict[str, Any] = {"text": self.text, "literals": list(self.literals)} + if self.parse_error: + out["parse_error"] = self.parse_error + return out + + +def _is_literal(token: Any) -> bool: + """字符串（含 N'…'、x'…'、$$…$$ 这些变体）和数字。""" + name = token.token_type.name + return name == "NUMBER" or name.endswith("STRING") + + +def _tokens(text: str, read: str | None) -> list[Any]: + from sqlglot.dialects.dialect import Dialect + + return list(Dialect.get_or_raise(read).tokenize(text)) + + +def normalize_sql(sql: str, dialect: str | None = None) -> NormalizedSql: + """字面量换成占位符后的规范文本 + 字面量。 + + 先按方言解析、重新生成一遍（去掉注释、统一关键字和空白、按方言的规则转换不加引号的名字的大小写），再把生成的 + 文本切成词，字符串和数字换成「?」：字面量按规范文本里出现的顺序排，不受语法树里子句存放顺序的影响。解析 + 不了就直接对原文逐词规范化，带上 parse_error。从不抛异常。 + """ + read = _read(dialect) + root, problem = _parse(sql, read) + if root is not None: + try: + from sqlglot.optimizer.normalize_identifiers import normalize_identifiers + + canonical = normalize_identifiers(root.copy(), dialect=read).sql(dialect=read, comments=False) + pieces: list[str] = [] + literals: list[str] = [] + pos = 0 + for token in _tokens(canonical, read): + if _is_literal(token): + pieces += [canonical[pos:token.start], PLACEHOLDER] + literals.append(canonical[token.start:token.end + 1]) + pos = token.end + 1 + pieces.append(canonical[pos:]) + return NormalizedSql("".join(pieces), tuple(literals)) + except Exception: # noqa: BLE001 - 生成、切词出错：退回逐词规范化，照实标出 + logger.warning("SQL 规范化失败，退回逐词规范化", exc_info=True) + problem = "SQL 的写法无法规范化" + return _by_tokens(sql if isinstance(sql, str) else "", read, problem) + + +def _by_tokens(text: str, read: str | None, problem: str | None) -> NormalizedSql: + """解析不了的 SQL 逐词规范化：关键字转大写，不加引号的名字按 name_key，带引号的名字原样，字面量换成占位符。""" + try: + tokens = _tokens(text, read) + except Exception: # noqa: BLE001 - 引号、括号没配对：只能把空白压成一个 + return NormalizedSql(" ".join(text.split()), (), problem or "SQL 无法解析") + words: list[str] = [] + literals: list[str] = [] + for token in tokens: + raw = text[token.start:token.end + 1] + if _is_literal(token): + words.append(PLACEHOLDER) + literals.append(raw) + elif token.token_type.name == "VAR": + words.append(name_key(token.text)) + elif token.token_type.name == "IDENTIFIER": + words.append(raw) + else: + words.append(token.text.upper()) + return NormalizedSql(" ".join(words), tuple(literals), problem) + + +# ========================================================================== +# 各个方面 +# ========================================================================== + +#: 一列：(表, 列)。表是给人看的真实表名（CTE、子查询写它的名字），认不出是哪张表时为 None +ColumnRef = tuple[str | None, str] + + +@dataclass(frozen=True) +class Part: + """一个条件、一个分组键、一个聚合参数…… + + key 用来比较：别名换成真实表名（CTE、子查询写编号），表名、列名按 name_key，等式和比较的两边排好顺序；loose 同 + key 但不写表名，两边查的表换了一张（orders 换成 orders_archive）时拿它比其余各方面，免得同一个条件只因表名不同 + 被报成不同；anon 同 key，但 CTE、子查询一律写「@」、不写编号（loose 也不写），编号错开时拿它再配一遍（见 + _Compare.unmatched）。text 给人看：好几张表都有的列写「表.列」（真实表名，不是别名），其余只写列名；full 一律 + 写表名，两边的 text 一样、key 却不同时改用它说明。excludes 只在 WHERE 的「排除」写法上有：<>、NOT、NOT IN、 + IS NOT NULL 去掉否定之后排除的是什么（IS NOT NULL 写成「列 为空」）。column_keys 和 columns 一一对应，是各列 + 比较用的写法（实表.列；CTE、子查询的列写成「@.列」）：两边给 CTE、子查询起的名字不同、编号错开，也认得出两个 + 条件说的是同一列，配成一对说。 + """ + + key: str + text: str + columns: tuple[ColumnRef, ...] = () + excludes: str | None = None + loose: str = "" + full: str = "" + column_keys: tuple[str, ...] = () + anon: str = "" + + @staticmethod + def plain(key: str, text: str) -> "Part": + """没有列的一项（表名、COUNT(*) 的「*」）：比较用的写法都是 key，给人看的都是 text。""" + return Part(key=key, text=text, loose=key, full=text, anon=key) + + +@dataclass(frozen=True) +class JoinFacet: + """一次连接。kind：inner / left / right / full / cross / semi / anti / natural。table 是关联进来的那一侧，other + 是另一侧：关联条件里出现的别的表（给人看的真实表名）；没有条件的（CROSS JOIN）是在它之前出现的各个来源。on 是 + 关联条件的合取式排好序后用 AND 连起来的写法，没有条件是 None。""" + + kind: str + table: Part + on: Part | None = None + other: tuple[str, ...] = () + + +@dataclass(frozen=True) +class TimePredicate: + """WHERE 里的一个时间条件。op：= / <> / > / >= / < / <= / between / in。value 是另一侧给人看的写法： + BETWEEN 是「低 AND 高」，IN 是「(…)」。""" + + column: Part + op: str + value: str + condition: Part + + +@dataclass(frozen=True) +class Aggregate: + """一个聚合。func 是函数名的小写（sum / count / avg / min / max …）；arg 是参数，COUNT(*)、COUNT(1) 的 arg + 是「*」，去重聚合是 DISTINCT 里面的表达式；text 是整个聚合给人看的写法。窗口函数不算。""" + + func: str + arg: Part + distinct: bool + text: str + + +@dataclass(frozen=True) +class LimitFacet: + """一段查询取哪几行。count、offset 是行数和跳过行数的写法（没写为 None，OFFSET 0 也算没写）；order 是这一段的 + 排序键 (Part, 是否降序)：只取前几行时，按什么排决定了取到的是哪几行。scope 是哪一段：None 是最外层，CTE 写它的 + 名字，FROM 里的子查询写它的别名（没有别名写「子查询」）。条件里套着的子查询不单算：它整个写在那个条件里比。""" + + count: str | None + offset: str | None = None + order: tuple[tuple[Part, bool], ...] = () + scope: str | None = None + + +@dataclass(frozen=True) +class Facets: + """一条查询的各个方面。每一项都按在 SQL 里出现的先后排、按 key 去重；parse_error 不为空时其余都是空的。 + + limit 是最外层 LIMIT 的行数（写法）；limits 是各段取哪几行（最外层在前，见 LimitFacet），只有写了 LIMIT 或 + OFFSET 的段才有。 + """ + + dialect: str | None = None + tables: tuple[Part, ...] = () + joins: tuple[JoinFacet, ...] = () + where: tuple[Part, ...] = () + time: tuple[TimePredicate, ...] = () + group_by: tuple[Part, ...] = () + aggregates: tuple[Aggregate, ...] = () + having: tuple[Part, ...] = () + limit: str | None = None + limits: tuple[LimitFacet, ...] = () + parse_error: str | None = None + + def as_dict(self) -> dict[str, Any]: + return _jsonable(self) + + +def _jsonable(value: Any) -> Any: + if is_dataclass(value) and not isinstance(value, type): + return {f.name: _jsonable(getattr(value, f.name)) for f in fields(value)} + if isinstance(value, (list, tuple)): + return [_jsonable(v) for v in value] + return value + + +def facets(sql: str, dialect: str | None = None, *, schema: Mapping[str, Any] | None = None) -> Facets: + """把一条查询拆成各个方面。从不抛异常：解析不了、或者分析时出了意外，返回带 parse_error 的 Facets。 + + schema 是数据源的 schema_cache 形状（{"tables": {表名: {"columns": [{"name": …}]}}}），可以不给。给了的话， + 几张表连在一起、列又没加表前缀时按它认列是哪张表的；不给就照写列名，比较时可能把同一列当成两列。 + """ + read = _read(dialect) + root, problem = _parse(sql, read) + if root is None: + return Facets(dialect=read, parse_error=problem) + try: + return _Statement(root, read, schema).facets() + except Exception: # noqa: BLE001 - 分析本身的缺陷不能连累调用方：当作分析不了，记下来修 + logger.warning("SQL 语义分析失败", exc_info=True) + return Facets(dialect=read, parse_error="SQL 的结构无法分析") + + +# ---- 语法树上的小工具 + +#: 比较：(运算符的写法, 两边对调之后的类型) +_COMPARE: dict[str, tuple[str, str]] = {"EQ": ("=", "EQ"), "NEQ": ("<>", "NEQ"), "GT": (">", "LT"), + "GTE": (">=", "LTE"), "LT": ("<", "GT"), "LTE": ("<=", "GTE")} +#: 连接类型的写法（diff_facets 里用） +_JOIN_WORDS = {"inner": "INNER JOIN", "left": "LEFT JOIN", "right": "RIGHT JOIN", "full": "FULL JOIN", + "cross": "CROSS JOIN", "semi": "SEMI JOIN", "anti": "ANTI JOIN", "natural": "NATURAL JOIN"} +#: 日期、时间函数（sqlglot 把各方言的写法归成这些类型：strftime、DATE_FORMAT 都是 TimeToStr） +_DATE_FUNCS = ("CurrentDate", "CurrentDatetime", "CurrentTimestamp", "CurrentTime", "Date", "Datetime", "DateAdd", + "DateSub", "DateDiff", "DateTrunc", "TimestampTrunc", "DatetimeTrunc", "TimeToStr", "StrToDate", + "StrToTime", "TsOrDsToDate", "TsOrDsToTimestamp", "TsOrDsAdd", "DateStrToDate", "TimeStrToTime", + "UnixToTime", "Interval", "Extract", "DateFromParts", "LastDay", "Year", "Month", "Day", "Week", + "Quarter", "DayOfWeek") +#: sqlglot 不认识、原样当普通函数的日期函数名 +_DATE_NAMES = frozenset({"date", "datetime", "strftime", "julianday", "now", "today", "to_date", "to_char", + "date_format", "unixepoch", "curdate", "sysdate", "getdate"}) +#: 像日期、时间的字符串：2026-09-01、2026-09、2026/9/1、2026-09-01 08:00:00 +_DATE_TEXT = re.compile(r"^\s*\d{4}[-/]\d{1,2}(?:[-/]\d{1,2})?(?:[ T]\d{1,2}:\d{2}(?::\d{2}(?:\.\d+)?)?)?\s*$") +#: 中文列名里像日期、时间的写法（英文列名按 catalog.is_date_column） +_CN_TIME = re.compile(r"(?:日期|时间|时刻|年月|月份)$|^(?:年|月|日|周|日期|时间)$") +_PLAIN_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _nested(node: Any, stop: Any) -> bool: + """node 在 stop 底下的某个子查询里（不属于 stop 这一段）。""" + from sqlglot import exp + + return any(isinstance(p, (exp.Select, exp.Subquery)) for p in _parents(node, stop=stop)) + + +def _own_columns(node: Any) -> list[Any]: + """node 里属于这一段的列（不进嵌套的子查询），不含 t.*。""" + from sqlglot import exp + + if node is None: + return [] + found = [node] if isinstance(node, exp.Column) else \ + [c for c in node.find_all(exp.Column) if not _nested(c, node)] + return [c for c in found if c.name and not c.is_star] + + +def _inner_queries(query: Any, skip: Iterable[Any] = ()) -> list[Any]: + """query 里直接套着的查询（不再往里找），按出现的先后：条件、SELECT 列表里的子查询，UNION 的各支，子查询里面 + 那段。WITH 里的 CTE 定义不算（在用到它的地方才算），skip 里的节点整个跳过。""" + from sqlglot import exp + + skipped = {id(n) for n in skip} + found = [] + stack = list(reversed(list(query.iter_expressions()))) + while stack: + node = stack.pop() + if id(node) in skipped or isinstance(node, exp.With): + continue + if isinstance(node, exp.Query): + found.append(node) + continue + stack.extend(reversed(list(node.iter_expressions()))) + return found + + +def _named_output(select: Any) -> bool: + """这段 SELECT 选出来的列会被别处按名字引用：它（或者它所在的 UNION）是 FROM / JOIN 里的子查询、CTE 的定义。 + IN、EXISTS、标量子查询只看取值，选出来的列叫什么不影响口径。""" + from sqlglot import exp + + node = select + while isinstance(node.parent, (exp.Subquery, exp.SetOperation)): + node = node.parent + return isinstance(node.parent, (exp.From, exp.Join, exp.CTE, exp.Lateral)) + + +def _alias(item: Any) -> str: + """FROM / JOIN 项的别名（没有别名的实表、CTE 是它的名字）；表函数这类认不出的来源是空串。LATERAL 的别名 + 挂在 Lateral 节点上，不在里面的子查询上。""" + from sqlglot import exp + + return item.alias_or_name if isinstance(item, (exp.Table, exp.Subquery, exp.Lateral)) else "" + + +def _derived_body(item: Any) -> Any: + """FROM / JOIN 项是子查询（含 LATERAL 子查询）时，里面那段查询；别的（实表、CTE、表函数）返回 None。""" + from sqlglot import exp + + if isinstance(item, exp.Lateral): + item = item.this + return item.this if isinstance(item, exp.Subquery) else None + + +def _own_nodes(select: Any) -> Iterator[Any]: + """一段 SELECT 自己的节点：不进 CTE、子查询（它们是别的段）。""" + from sqlglot import exp + + stack = [select] + while stack: + node = stack.pop() + yield node + children = [c for c in node.iter_expressions() + if not isinstance(c, (exp.Select, exp.Subquery, exp.CTE, exp.With))] + stack.extend(reversed(children)) + + +def _datey(node: Any) -> bool: + """node 里有日期函数、转成日期的 CAST，或者像日期的字符串。""" + from sqlglot import exp + + kinds = tuple(k for k in (getattr(exp, n, None) for n in _DATE_FUNCS) if isinstance(k, type)) + for n in node.walk(): + if isinstance(n, kinds): + return True + if isinstance(n, exp.Anonymous) and str(n.name).lower() in _DATE_NAMES: + return True + if isinstance(n, exp.Cast) and n.to is not None and any(w in n.to.this.name for w in ("DATE", "TIME")): + return True + if isinstance(n, exp.Literal) and n.is_string and _DATE_TEXT.match(str(n.this)): + return True + return False + + +def _temporal_name(name: str) -> bool: + """列名像日期、时间：英文按 catalog.is_date_column 的命名规则（paid_at、order_date），中文看结尾（支付日期）。""" + return catalog.is_date_column(name, None) or bool(_CN_TIME.search(name)) + + +def _join_kind(join: Any) -> str: + side = str(join.args.get("side") or "").lower() + kind = str(join.args.get("kind") or "").lower() + if side in ("left", "right", "full"): + return side + if join.args.get("method"): + return "natural" + if kind in ("semi", "anti"): + return kind + if kind == "cross" or not (join.args.get("on") is not None or join.args.get("using")): + return "cross" + return "inner" + + +def _dedupe(items: Iterable[Any], key: Any) -> tuple[Any, ...]: + seen: dict[Any, Any] = {} + for item in items: + seen.setdefault(key(item), item) + return tuple(seen.values()) + + +class _Sources(dict): + """一段 SELECT 的来源 {别名的 name_key: _Source}。outer 是外层那段的来源：相关子查询里引用外层的别名（o.id） + 从这里找。""" + + outer: "_Sources | None" = None + + +@dataclass +class _Source: + """一段 SELECT 的一个来源。kind：table（实表）/ derived（CTE、子查询、LATERAL 子查询）/ other（表函数、认不出的 + 别名）。""" + + kind: str + name: str # 给人看：表名（有库名就带上）、CTE 名、子查询（含 LATERAL）的别名 + key: str # 比较用：实表是 name_key，派生来源是编号（见 _Statement._number_derived） + idents: tuple[Any, ...] = () # 表名各段的 Identifier：给人看的写法照原样写回去 + query: Any = None # derived：里面那段查询 + + +class _Statement: + """一条语句的分析：CTE、实表、每段 SELECT 的来源，以及把表达式写成 key / text 的规则。""" + + def __init__(self, root: Any, read: str | None, schema: Mapping[str, Any] | None) -> None: + from sqlglot import exp + + self.root = root + self.read = read + #: 各段 SELECT（原树上的节点），下标就是记在 meta 里的段号 + self.segments: list[Any] = [] + for select in root.find_all(exp.Select, bfs=False): + select.meta[_SEGMENT] = len(self.segments) + self.segments.append(select) + self.ctes: dict[str, Any] = {} + #: CTE 的定义（exp.CTE 节点），和 ctes 同一个键：编号也记在定义上，定义写进条件时靠它换掉名字 + self._cte_defs: dict[str, Any] = {} + for cte in root.find_all(exp.CTE): + if cte.alias: + self.ctes.setdefault(name_key(cte.alias), cte.this) + self._cte_defs.setdefault(name_key(cte.alias), cte) + self._number_derived() + tables: dict[str, _Source] = {} + for table in root.find_all(exp.Table, bfs=False): + if isinstance(table.this, exp.Identifier) and table.name and not self._is_cte(table): + source = self._table(table) + tables.setdefault(source.key, source) + self.tables = list(tables.values()) + #: 整条语句只用到一张实表：给人看的写法一律不写表名 + self.single = len(self.tables) <= 1 + self.columns: dict[str, set[str]] = {} + for name, meta in ((schema or {}).get("tables") or {}).items(): + if isinstance(meta, Mapping): + self.columns[name_key(str(name))] = {name_key(str(c.get("name"))) for c in meta.get("columns") or [] + if isinstance(c, Mapping) and c.get("name")} + self._memo: dict[int, _Sources] = {} + self._shared: set[str] | None = None + + # ---- 来源 + + def _is_cte(self, table: Any) -> bool: + return not table.text("db") and not table.text("catalog") and name_key(table.name) in self.ctes + + @staticmethod + def _items(select: Any) -> list[Any]: + """一段 SELECT 的 FROM、JOIN 项，按出现的先后。""" + first = select.args.get("from_") + return ([first.this] if first is not None else []) + [j.this for j in select.args.get("joins") or []] + + def _number_derived(self) -> None: + """给派生来源（引用 CTE 的表、FROM / JOIN 里的子查询和 LATERAL 子查询）编比较用的 key：@1、@2…，记在节点的 + meta 上。 + + CTE 名、子查询的别名是各自起的（双盲复算的盲算更是），拿它作 key，同一个口径只因起的名字不同就被报成不同； + 拿里面那段查询的写法作 key 也不行：里面差一个条件，外层按它的列写的条件就全跟着变，同一处不同报好几遍，还会 + 出现两边各说一句「只保留 total > 100 的行」的自相矛盾。所以只按位置编：从最外层往里一层一层地编，同一层按 + 出现的先后；CTE 在第一次用到的地方编号，在哪儿用到都是这一个编号。写成 CTE 还是 FROM 里的子查询、CTE 定义 + 的先后都不影响编号，里层多一段、少一段也不挪动外层的编号。 + """ + from sqlglot import exp + + ctes: dict[str, str] = {} + count = 0 + + def visit(select: Any) -> list[Any]: + """给这一段还没编号的 FROM / JOIN 项编号，返回要接着往里看的查询（子查询、第一次用到的 CTE 的定义）。""" + nonlocal count + inner = [] + for item in self._items(select): + body = _derived_body(item) + if body is not None: + if _DERIVED not in item.meta: + count += 1 + item.meta[_DERIVED] = f"@{count}" + inner.append(body) + elif isinstance(item, exp.Table) and isinstance(item.this, exp.Identifier) and item.name \ + and self._is_cte(item): + name = name_key(item.name) + if name not in ctes: + count += 1 + ctes[name] = f"@{count}" + self._cte_defs[name].meta[_DERIVED] = ctes[name] + inner.append(self.ctes[name]) + item.meta[_DERIVED] = ctes[name] + return inner + + queue: deque[Any] = deque([self.root]) + while queue: + query = queue.popleft() + items: list[Any] = [] + if isinstance(query, exp.Select): + queue.extend(visit(query)) + items = self._items(query) + queue.extend(_inner_queries(query, skip=items)) + # 没走到的（没用到的 CTE 里的、表函数参数里的）按出现的先后补上：每个派生来源都得有编号 + for select in self.root.find_all(exp.Select, bfs=False): + visit(select) + + @staticmethod + def _table(table: Any) -> _Source: + from sqlglot import exp + + idents = tuple(i for i in (table.args.get("catalog"), table.args.get("db"), table.this) + if isinstance(i, exp.Identifier)) + return _Source("table", ".".join(i.name for i in idents), ".".join(name_key(i.name) for i in idents), + tuple(i.copy() for i in idents)) + + def sources(self, select: Any) -> _Sources: + """{别名的 name_key: 来源}。没有别名的实表按表名；认不出的来源按位置编号。""" + from sqlglot import exp + + if id(select) in self._memo: + return self._memo[id(select)] + out = _Sources() + for i, node in enumerate(self._items(select)): + alias = _alias(node) + slot = name_key(alias) if alias else f"#{i}" + if isinstance(node, exp.Table) and isinstance(node.this, exp.Identifier) and node.name: + if self._is_cte(node): + out[slot] = _Source("derived", node.name, node.meta.get(_DERIVED, slot), (node.this.copy(),), + self.ctes[name_key(node.name)]) + else: + out[slot] = self._table(node) + elif (body := _derived_body(node)) is not None: + out[slot] = _Source("derived", alias or "子查询", node.meta.get(_DERIVED, slot), + (exp.to_identifier(alias),) if alias else (), body) + else: + out[slot] = _Source("other", alias, slot) + self._memo[id(select)] = out + parent = select.parent + while parent is not None and not isinstance(parent, exp.Select): + parent = parent.parent + out.outer = self.sources(parent) if parent is not None else None + return out + + def _exposes(self, source: _Source, column: str) -> bool: + """这个来源可能有这一列：实表看表结构（没给表结构就当可能有），CTE、子查询看它选出来的列。""" + from sqlglot import exp + + if source.kind == "table": + return source.key not in self.columns or name_key(column) in self.columns[source.key] + if source.kind == "derived" and isinstance(source.query, exp.Select): + return any(isinstance(s, exp.Star) or (isinstance(s, exp.Column) and s.is_star) + or name_key(s.alias_or_name) == name_key(column) for s in source.query.selects) + return True + + def locate(self, column: Any, sources: dict[str, _Source]) -> tuple[str | None, _Source | None]: + """列 → (来源的别名键, 来源)。别名在这一段找不到的，到外层找（相关子查询引用外层），再找不到按「别的来源」 + 照写（别名写错了）；没加前缀、又有不止一个来源可能有它的，返回 (None, None)。""" + qualifier = column.table + if qualifier: + slot = name_key(qualifier) + if slot in sources: + return slot, sources[slot] + outer = getattr(sources, "outer", None) + while outer is not None: + if slot in outer: + return None, outer[slot] # 外层的来源：不算这一段的关联 + outer = outer.outer + return None, _Source("other", qualifier, slot, (column.args["table"].copy(),)) + hits = [(slot, s) for slot, s in sources.items() if self._exposes(s, column.name)] + return hits[0] if len(hits) == 1 else (None, None) + + def through(self, source: _Source | None, column: str, depth: int = 0) -> tuple[_Source | None, str]: + """CTE、子查询只是把一张表的列原样透出来（SELECT *，或者直接选出这一列）时，换成那张表的列： + 原流程直接查 orders、盲算先在 CTE 里筛一遍再查，算的是同一列。""" + from sqlglot import exp + + if source is None or source.kind != "derived" or depth > 3 or not isinstance(source.query, exp.Select): + return source, column + inner = self.sources(source.query) + if len(inner) != 1: + return source, column + only = next(iter(inner.values())) + for item in source.query.selects: + if isinstance(item, exp.Star) or (isinstance(item, exp.Column) and item.is_star): + return self.through(only, column, depth + 1) + target = _unparen(item.this if isinstance(item, exp.Alias) else item) + if name_key(item.alias_or_name) == name_key(column) and isinstance(target, exp.Column): + return self.through(only, target.name, depth + 1) + return source, column + + def resolve(self, column: Any, sources: dict[str, _Source]) -> tuple[_Source | None, str]: + return self.through(self.locate(column, sources)[1], column.name) + + def shared(self) -> set[str]: + """好几张表都有的列名（name_key）：这条语句里出现时挂在不止一张表上（orders.order_id、refunds.order_id）， + 或者表结构里不止一张用到的表有它。给人看时只有这些列写表名，其余只写列名，读起来和原 SQL 一样短。""" + from sqlglot import exp + + if self._shared is None: + owners: dict[str, set[str]] = {} + for select in self.root.find_all(exp.Select, bfs=False): + sources = self.sources(select) + for node in _own_nodes(select): + if isinstance(node, exp.Column) and node.name and not node.is_star: + source, name = self.resolve(node, sources) + if source is not None: + owners.setdefault(name_key(name), set()).add(source.key) + used = [t.key for t in self.tables] + self._shared = {name for name, keys in owners.items() + if len(keys) > 1 or sum(1 for t in used if name in self.columns.get(t, ())) > 1} + return self._shared + + # ---- 写成 key / text + + @staticmethod + def _owner(node: Any) -> Any: + """node 往上最近的一段 SELECT（带段号的，node 本身是也算）；没有就是 None：node 在调用方那一段里。""" + from sqlglot import exp + + while node is not None: + if isinstance(node, exp.Select) and node.meta.get(_SEGMENT) is not None: + return node + node = node.parent + return None + + def scope(self, node: Any, default: dict[str, _Source]) -> dict[str, _Source]: + """node 所在那一段的来源。条件里套着的子查询（EXISTS、IN (SELECT …)）是另一段：别名在那一段里定义，相关 + 子查询还要到外层找。复制出来的子查询带着段号，按段号找回原树上那一段算好的来源；不在子查询里就是 default。""" + owner = self._owner(node) + return default if owner is None else self.sources(self.segments[owner.meta[_SEGMENT]]) + + def render(self, node: Any, sources: dict[str, _Source], mode: str = "text") -> str: + """五种写法（见 Part）：key、anon、loose 比较用，名字按 name_key，key、anon 写真实表名、loose 不写（CTE、子查询 + 只有 key 写编号）；text、full 给人看，text 只在好几张表都有这一列时写真实表名，full 一律写（关联条件说的就是 + 两张表各出哪一列，用 full）。 + + node 里套着的子查询逐段改写：每一列按它所在那一段的来源认，子查询里的表去掉别名（列已经写成了真实表名）。 + 子查询里拿两张表的列互相比的（相关子查询），给人看时和关联条件一样写表名。比较用的写法里各自起的名字也去掉 + （见 _drop_names）。 + """ + from sqlglot import exp + + copy = node.copy() + groups: dict[int, tuple[Any, list[Any]]] = {} + for col in copy.find_all(exp.Column): + if col.name and not col.is_star: + owner = self._owner(col) + groups.setdefault(id(owner), (owner, []))[1].append(col) + for owner, cols in groups.values(): + scope = sources if owner is None else self.sources(self.segments[owner.meta[_SEGMENT]]) + found = [(col, *self.resolve(col, scope)) for col in cols] + here = mode + if owner is not None and mode == "text" and len({s.key for _, s, _ in found if s is not None}) > 1: + here = "full" + for col, source, name in found: + self._write(col, source, name, here) + for table in copy.find_all(exp.Table): + if isinstance(table.this, exp.Identifier) and self._owner(table) is not None: + table.set("alias", None) + if mode in ("key", "loose", "anon"): + for arg in ("catalog", "db", "this"): + ident = table.args.get(arg) + if isinstance(ident, exp.Identifier): + table.set(arg, exp.to_identifier(name_key(ident.name))) + if mode in ("key", "loose", "anon"): + self._drop_names(copy, numbered=mode == "key") + return copy.sql(dialect=self.read, comments=False) + + @staticmethod + def _drop_names(copy: Any, *, numbered: bool) -> None: + """比较用的写法里抹掉套着的子查询里各自起的名字（就地改）：引用 CTE 的表、FROM 里的子查询（含 LATERAL）、 + 子查询里自带的 CTE 定义写编号（numbered 为假时写「@」；引用它们的列已经照这样写了）；只看取值的子查询（IN、 + EXISTS、标量子查询）去掉选出来的列名。FROM 里的子查询、CTE 定义选出来的列名留着：外层按这个名字引用它们， + 去掉了就分不出外层用的是哪一列。""" + from sqlglot import exp + + for node in copy.find_all(exp.Table, exp.Subquery, exp.Lateral, exp.CTE): + derived = node.meta.get(_DERIVED) + if not derived: + continue + mark = exp.to_identifier(derived if numbered else _ANON) + if isinstance(node, exp.Table): + for arg in ("catalog", "db", "alias"): + node.set(arg, None) + node.set("this", mark) + elif isinstance(node.args.get("alias"), exp.TableAlias): + node.args["alias"].set("this", mark) + else: + node.set("alias", exp.TableAlias(this=mark)) + for select in copy.find_all(exp.Select): + if not _named_output(select): + select.set("expressions", [e.this if isinstance(e, exp.Alias) else e for e in select.expressions]) + + def _write(self, col: Any, source: _Source | None, name: str, mode: str) -> None: + """按 render 的写法改写一列（就地改）。""" + from sqlglot import exp + + quoted = isinstance(col.this, exp.Identifier) and bool(col.this.args.get("quoted")) + if mode in ("key", "loose", "anon"): + owner = None if source is None or mode == "loose" else \ + _ANON if mode == "anon" and source.kind == "derived" else source.key + col.set("this", exp.to_identifier(name_key(name))) + col.set("table", exp.to_identifier(owner) if owner else None) + col.set("db", None) + col.set("catalog", None) + return + col.set("this", exp.Identifier(this=name, quoted=quoted or not _PLAIN_NAME.match(name))) + if source is None or (source.kind != "other" and mode != "full" + and (self.single or name_key(name) not in self.shared())): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) + else: + *heads, last = source.idents or (exp.to_identifier(source.name),) + col.set("table", last.copy()) + col.set("db", heads[-1].copy() if heads else None) + col.set("catalog", heads[-2].copy() if len(heads) > 1 else None) + + def canonical(self, node: Any, sources: dict[str, _Source]) -> Any: + """比较的两边排好顺序（列在左、字面量在右；两边都有列的按 _rank 排），IN 的取值排好序。返回新的树。 + 套着的子查询里的比较按它那一段的来源排：拿外层的来源去认子查询的别名，排出来的先后就取决于别名怎么起。""" + from sqlglot import exp + + def fix(n: Any) -> Any: + found = _COMPARE.get(type(n).__name__) + if found is not None and isinstance(n, exp.Binary): + left, right = n.this, n.expression + lc, rc = bool(_own_columns(left)), bool(_own_columns(right)) + here = self.scope(n, sources) if lc and rc else sources + if (rc and not lc) or (lc and rc and self._rank(left, here) > self._rank(right, here)): + return getattr(exp, found[1])(this=right, expression=left) + elif isinstance(n, exp.In) and n.expressions and not n.args.get("query") \ + and not any(_own_columns(e) for e in n.expressions): + n.set("expressions", sorted(n.expressions, key=lambda e: e.sql(dialect=self.read))) + return n + + return node.transform(fix) + + def _rank(self, node: Any, sources: dict[str, _Source]) -> tuple[str, str]: + """比较的两边谁排前面：先按不写编号的写法，分不出（同名的列）再按编号。只按编号排的话，编号错开的两条语句里 + 同一个比较会排成两个方向，按不写编号的写法也配不上了。""" + return self.render(node, sources, "anon"), self.render(node, sources, "key") + + def part(self, node: Any, sources: dict[str, _Source], *, excludes: bool = False) -> Part: + refs: dict[tuple[str | None, str], str] = {} + for col in _own_columns(node): + source, name = self.resolve(col, sources) + owner = "" if source is None else _ANON if source.kind == "derived" else source.key + refs.setdefault((source.name if source is not None else None, name), f"{owner}.{name_key(name)}") + full = self.render(node, sources, "full") + # 拿两张表的列互相比的（相关子查询里的 o.id = r.order_id）和关联条件一样，给人看也写表名 + across = len({table for table, _ in refs}) > 1 + return Part(key=self.render(node, sources, "key"), text=full if across else self.render(node, sources), + columns=tuple(refs), excludes=self._excludes(node, sources) if excludes else None, + loose=self.render(node, sources, "loose"), full=full, column_keys=tuple(refs.values()), + anon=self.render(node, sources, "anon")) + + def _excludes(self, node: Any, sources: dict[str, _Source]) -> str | None: + """排除写法去掉否定之后排除的是什么：a <> 1 → a = 1；NOT IN、NOT LIKE、NOT BETWEEN 同理；IS NOT NULL → 「a 为空」。 + 各方言把否定放的位置不一样（NOT 包一层，或者节点上带 negate），两种都认。""" + from sqlglot import exp + + n = _unparen(node) + if isinstance(n, exp.NEQ): + return self.render(exp.EQ(this=n.this.copy(), expression=n.expression.copy()), sources) + inner = None + if isinstance(n, exp.Not): + inner = _unparen(n.this).copy() + elif n.args.get("negate"): + inner = n.copy() + inner.set("negate", None) + if inner is None: + return None + if isinstance(inner, exp.Is) and isinstance(_unparen(inner.expression), exp.Null): + return f"{self.render(inner.this, sources)} 为空" + return self.render(inner, sources) + + # ---- 各个方面 + + def facets(self) -> Facets: + from sqlglot import exp + + joins: list[JoinFacet] = [] + where: list[Part] = [] + time: list[TimePredicate] = [] + group: list[Part] = [] + aggregates: list[Aggregate] = [] + having: list[Part] = [] + for select in self.root.find_all(exp.Select, bfs=False): + sources = self.sources(select) + drafts = self._join_drafts(select, sources) + clause = select.args.get("where") + for cond in _conjuncts(clause.this) if clause is not None else []: + canon = self.canonical(cond, sources) + if self._joins_on_where(canon, sources, drafts): + continue + where.append(self.part(canon, sources, excludes=True)) + if found := self._time(canon, sources): + time.append(found) + joins += [self._join(d, sources) for d in drafts] + group += self._group(select, sources) + aggregates += self._aggregates(select, sources) + having += self._having(select, sources) + return Facets( + dialect=self.read, + tables=tuple(Part.plain(s.key, s.name) for s in self.tables), + joins=_dedupe(joins, lambda j: (j.kind, j.table.key, j.on.key if j.on else None)), + where=_dedupe(where, lambda p: p.key), + time=_dedupe(time, lambda t: t.condition.key), + group_by=_dedupe(group, lambda p: p.key), + aggregates=_dedupe(aggregates, lambda a: (a.func, a.arg.key, a.distinct)), + having=_dedupe(having, lambda p: p.key), + limit=self._limit(), + limits=tuple(self._limits()), + ) + + def _join_drafts(self, select: Any, sources: dict[str, _Source]) -> list[dict[str, Any]]: + from sqlglot import exp + + first = select.args.get("from_") + first_alias = _alias(first.this) if first is not None else "" + drafts: list[dict[str, Any]] = [] + for i, join in enumerate(select.args.get("joins") or []): + alias = _alias(join.this) + slot = name_key(alias) if alias else f"#{i + 1}" + conds: list[Any] = [] + if join.args.get("on") is not None: + conds = [self.canonical(c, sources) for c in _conjuncts(join.args["on"])] + elif join.args.get("using") and first_alias and alias: + # USING (id)：按「FROM 的那张表.id = 这张表.id」写，和 ON 的写法放在一起比 + conds = [self.canonical(exp.EQ(this=exp.column(ident.name, table=first_alias), + expression=exp.column(ident.name, table=alias)), sources) + for ident in join.args["using"]] + drafts.append({"kind": _join_kind(join), "slot": slot, "source": sources.get(slot), "conds": conds, + "before": [s.name for key, s in list(sources.items())[:i + 1] if key != slot]}) + return drafts + + def _joins_on_where(self, cond: Any, sources: dict[str, _Source], drafts: list[dict[str, Any]]) -> bool: + """WHERE 里两张表的列等式（逗号连接、或者 INNER JOIN 之外又写的等值关联）并进那次连接的条件，不算筛选。 + LEFT JOIN 这类外连接不并：写在 WHERE 里会把外连接变成内连接，和写在 ON 里不是一回事。""" + from sqlglot import exp + + n = _unparen(cond) + if not isinstance(n, exp.EQ): + return False + left, right = _unparen(n.this), _unparen(n.expression) + if not (isinstance(left, exp.Column) and isinstance(right, exp.Column)): + return False + a, b = self.locate(left, sources)[0], self.locate(right, sources)[0] + if a is None or b is None or a == b: + return False + for draft in reversed(drafts): + if draft["slot"] in (a, b) and draft["kind"] in ("inner", "cross"): + draft["conds"].append(cond) + draft["kind"] = "inner" + return True + return False + + def _join(self, draft: dict[str, Any], sources: dict[str, _Source]) -> JoinFacet: + source = draft["source"] + table = Part.plain("?", "?") if source is None else \ + Part(key=source.key, text=source.name, loose=source.key, full=source.name, + anon=_ANON if source.kind == "derived" else source.key) + forms = sorted((self.render(c, sources, "key"), self.render(c, sources, "loose"), + self.render(c, sources, "full"), self.render(c, sources, "anon")) for c in draft["conds"]) + on = Part(key=" AND ".join(f[0] for f in forms), text=" AND ".join(f[2] for f in forms), + loose=" AND ".join(sorted(f[1] for f in forms)), full=" AND ".join(f[2] for f in forms), + anon=" AND ".join(sorted(f[3] for f in forms))) if forms else None + mentioned = {table for c in draft["conds"] for table, _ in self.part(c, sources).columns} + other = [name for name in dict.fromkeys(t for t in mentioned if t) if name != table.text] if draft["conds"] \ + else draft["before"] + return JoinFacet(kind=draft["kind"], table=table, on=on, other=tuple(sorted(other))) + + def _time(self, cond: Any, sources: dict[str, _Source]) -> TimePredicate | None: + """时间条件：一边恰好一列（可以套着日期函数、CAST），另一边没有列；列名像日期，或者用了日期函数、像日期的字符串。""" + from sqlglot import exp + + n = _unparen(cond) + if n.args.get("negate"): + return None + if isinstance(n, exp.Between): + subject, values, op = n.this, [n.args.get("low"), n.args.get("high")], "between" + elif isinstance(n, exp.In) and n.expressions and not n.args.get("query"): + subject, values, op = n.this, list(n.expressions), "in" + elif type(n).__name__ in _COMPARE and isinstance(n, exp.Binary): + subject, values, op = n.this, [n.expression], _COMPARE[type(n).__name__][0] + else: + return None + if subject is None or any(v is None or _own_columns(v) for v in values): + return None + column = _bare_column(subject) + cols = [column] if column is not None else _own_columns(subject) + if len(cols) != 1 or any(isinstance(x, (exp.AggFunc, exp.Select)) for x in subject.walk()): + return None + if not (_temporal_name(cols[0].name) or _datey(subject) or any(_datey(v) for v in values)): + return None + shown = [self.render(v, sources) for v in values] + value = f"{shown[0]} AND {shown[1]}" if op == "between" else f"({', '.join(shown)})" if op == "in" \ + else shown[0] + return TimePredicate(column=self.part(cols[0], sources), op=op, value=value, + condition=self.part(cond, sources, excludes=True)) + + @staticmethod + def _aliases(query: Any) -> dict[str, Any]: + """SELECT 里的别名 {name_key: 别名指的表达式}。""" + from sqlglot import exp + + return {name_key(s.alias): s.this for s in query.selects if isinstance(s, exp.Alias) and s.alias} + + @staticmethod + def _target(item: Any, selects: list[Any], aliases: dict[str, Any]) -> Any: + """GROUP BY、ORDER BY 里的一项：写序号（1）的换成 SELECT 里对应的表达式，写别名的换成别名指的表达式 + （MySQL、SQLite 的读法），其余原样。""" + from sqlglot import exp + + target = _unparen(item) + if isinstance(target, exp.Literal) and not target.is_string and str(target.this).isdigit(): + index = int(target.this) - 1 + if 0 <= index < len(selects): + return selects[index].this if isinstance(selects[index], exp.Alias) else selects[index] + elif isinstance(target, exp.Column) and not target.table and name_key(target.name) in aliases: + return aliases[name_key(target.name)] + return target + + def _group(self, select: Any, sources: dict[str, _Source]) -> list[Part]: + """分组键。GROUP BY 1 换成 SELECT 里对应的表达式，GROUP BY 别名换成别名指的表达式（MySQL、SQLite 的读法）。""" + group = select.args.get("group") + if group is None: + return [] + selects, aliases = select.selects, self._aliases(select) + return [self.part(self.canonical(self._target(item, selects, aliases), sources), sources) + for item in group.expressions] + + def _having(self, select: Any, sources: dict[str, _Source]) -> list[Part]: + """HAVING 的各个合取式。引用 SELECT 别名的列（HAVING s > 100，MySQL、SQLite 的写法）换成别名指的表达式， + 和 GROUP BY 别名一个读法：不换的话，写别名和写表达式的同一个条件会被报成两个不同的条件。""" + from sqlglot import exp + + clause = select.args.get("having") + if clause is None: + return [] + aliases = self._aliases(select) + out = [] + for cond in _conjuncts(clause.this): + copy = cond.copy() + for col in _own_columns(copy): + target = aliases.get(name_key(col.name)) if not col.table else None + # 聚合函数里的名字指的是表里的列，不是别名：SUM(amount) AS amount 之后的 SUM(amount) 不能再换 + if target is None or any(isinstance(p, exp.AggFunc) for p in _parents(col, stop=None)): + continue + if col is copy: + copy = target.copy() + break + col.replace(target.copy()) + out.append(self.part(self.canonical(copy, sources), sources)) + return out + + def _aggregates(self, select: Any, sources: dict[str, _Source]) -> list[Aggregate]: + from sqlglot import exp + + out = [] + for node in _own_nodes(select): + if not isinstance(node, exp.AggFunc) or isinstance(node.parent, exp.Window): + continue + try: + func = str(node.sql_name()).lower() + except Exception: # noqa: BLE001 - 没有固定函数名的聚合：按类型名 + func = node.key + arg = node.this + distinct = isinstance(arg, exp.Distinct) + if distinct: + items = list(arg.expressions) + arg = items[0] if len(items) == 1 else exp.Tuple(expressions=[i.copy() for i in items]) + if arg is None or isinstance(arg, exp.Star) or \ + (func == "count" and not distinct and isinstance(_unparen(arg), exp.Literal)): + part = Part.plain("*", "*") # COUNT(*)、COUNT(1)：数的是行数 + else: + part = self.part(self.canonical(arg, sources), sources) + out.append(Aggregate(func=func, arg=part, distinct=distinct, text=self.render(node, sources))) + return out + + def _limit(self) -> str | None: + from sqlglot import exp + + node = self.root.args.get("limit") + if node is None: + return None + value = node.expression if isinstance(node, exp.Limit) else node.args.get("count") + return value.sql(dialect=self.read, comments=False) if value is not None else None + + def _limits(self) -> list[LimitFacet]: + """各段取哪几行：最外层在前，其余按出现的先后。条件里套着的子查询（WHERE、HAVING、ON 里的）不单算： + 它整个写在那个条件的 key 里，单算的话同一处不同会报两遍。""" + from sqlglot import exp + + out = [] + for query in self.root.find_all(exp.Select, exp.SetOperation, bfs=False): + if query is not self.root and self._in_condition(query): + continue + if (found := self._limit_of(query)) is not None: + out.append(found) + return out + + def _in_condition(self, node: Any) -> bool: + from sqlglot import exp + + while node is not None and node is not self.root: + if isinstance(node, (exp.Where, exp.Having)) or \ + (isinstance(node.parent, exp.Join) and node.arg_key == "on"): + return True + node = node.parent + return False + + def _limit_of(self, query: Any) -> LimitFacet | None: + """一段的 LIMIT / OFFSET，连同这一段的排序；两样都没写返回 None。""" + from sqlglot import exp + + def written(node: Any) -> str | None: + return node.sql(dialect=self.read, comments=False) if node is not None else None + + limit, offset = query.args.get("limit"), query.args.get("offset") + count = written(limit.expression if isinstance(limit, exp.Limit) else limit.args.get("count")) \ + if limit is not None else None + if offset is None and isinstance(limit, exp.Limit): + offset = limit.args.get("offset") # 有的方言把 OFFSET 挂在 LIMIT 底下 + skip = written(offset.expression if isinstance(offset, exp.Offset) else offset) + if skip is not None and skip.strip() == "0": + skip = None # OFFSET 0 等于没写 + if limit is None and skip is None: + return None + return LimitFacet(count=count, offset=skip, order=self._order(query), scope=self._scope_name(query)) + + def _order(self, query: Any) -> tuple[tuple[Part, bool], ...]: + """这一段的排序键 (Part, 是否降序)。ORDER BY 1、ORDER BY 别名换成 SELECT 里的表达式（UNION 按第一支的）。""" + from sqlglot import exp + + order = query.args.get("order") + if order is None: + return () + select = query if isinstance(query, exp.Select) else None + sources = self.sources(select) if select is not None else _Sources() + selects = list(query.selects) + aliases = self._aliases(query) + out = [] + for item in order.expressions: + target = self._target(item.this if isinstance(item, exp.Ordered) else item, selects, aliases) + desc = bool(item.args.get("desc")) if isinstance(item, exp.Ordered) else False + out.append((self.part(self.canonical(target, sources), sources), desc)) + return tuple(out) + + def _scope_name(self, query: Any) -> str | None: + from sqlglot import exp + + if query is self.root: + return None + parent = query.parent + if isinstance(parent, (exp.CTE, exp.Subquery)) and parent.alias: + return parent.alias + return "子查询" + + +# ========================================================================== +# 比较 +# ========================================================================== + + +@dataclass(frozen=True) +class FacetDiff: + """一处不同。aspect：time（时间列）/ time_range（同一时间列、范围不同）/ filter / aggregate / join / tables / + group_by / having / limit（取哪几行：行数、跳过的行数、取前几行时的排序）/ parse。a、b 是两边各自的写法（给人看， + 没有就是 None），text 是一句中文说明。 + detail 只在 time 上有：两边的时间列 [{table, column}]，调用方据此对照数据目录（业务日期）。""" + + aspect: str + a: Any + b: Any + text: str + detail: dict[str, Any] | None = None + + def as_dict(self) -> dict[str, Any]: + out = {"aspect": self.aspect, "a": self.a, "b": self.b, "text": self.text} + if self.detail is not None: + out["detail"] = self.detail + return out + + +_CJK_PUNCT = frozenset("，。；：、！？「」『』（）《》【】…—") +#: 聚合函数的说法 +_VERBS = {"sum": "求和", "count": "计数", "avg": "求平均", "min": "取最小值", "max": "取最大值"} + + +def _wide(ch: str) -> bool: + return "⺀" <= ch <= "鿿" or "豈" <= ch <= "﫿" or "＀" <= ch <= "￯" + + +def _say(*parts: str) -> str: + """拼一句话：汉字和英文、数字、SQL 之间加一个空格，挨着全角标点的不加，汉字和汉字之间不加。 + + 两边的叫法由调用方定（A / B，或者「原流程」「盲算」），所以空格只能拼的时候按字判断。 + """ + out = "" + for part in parts: + if not part: + continue + if out and not out[-1].isspace() and not part[0].isspace(): + a, b = out[-1], part[0] + if not (a in _CJK_PUNCT or b in _CJK_PUNCT or (_wide(a) and _wide(b))): + out += " " + out += part + return out + + +def catalog_labels(notes: Mapping[str, Mapping[str, Any]] | None) -> dict[str, str]: + """数据目录（{表名: notes}）里列的中文名 → diff_facets 的 names：{「表.列」: 中文名}；只在一张表里有、或者各表 + 叫法一致的列另记 {列: 中文名}。被驳回的项不算。""" + out: dict[str, str] = {} + bare: dict[str, tuple[str, set[str]]] = {} + for table, raw in (notes or {}).items(): + visible = catalog.visible_notes(raw) if isinstance(raw, Mapping) else {} + columns = visible.get("columns") if isinstance(visible.get("columns"), Mapping) else {} + for column, items in columns.items(): + item = items.get("label") if isinstance(items, Mapping) else None + label = item.get("value") if isinstance(item, Mapping) else None + if isinstance(label, str) and label.strip(): + out[f"{table}.{column}"] = label.strip() + bare.setdefault(name_key(str(column)), (str(column), set()))[1].add(label.strip()) + for column, labels in bare.values(): + if len(labels) == 1: + out.setdefault(column, next(iter(labels))) + return out + + +def _labeler(names: Mapping[str, str] | None) -> Any: + index = {name_key(str(k)): str(v) for k, v in (names or {}).items() if isinstance(v, str) and v.strip()} + + def label(ref: ColumnRef) -> str | None: + table, column = ref + return (index.get(name_key(f"{table}.{column}")) if table else None) or index.get(name_key(column)) + + return label + + +def _single_column(part: Part, *, loose: bool) -> str | None: + """只涉及一列的条件：那一列（Part.column_keys：实表.列，CTE、子查询的列写「@.列」，两边起的名字不同、编号错开 + 也配得上；loose 时只看列名），用来把两边对同一列的不同写法配成一对。""" + if len(part.columns) != 1 or len(part.column_keys) != 1: + return None + return name_key(part.columns[0][1]) if loose else part.column_keys[0] + + +def _verb(g: Aggregate) -> str: + return ("去重" if g.distinct else "") + _VERBS.get(g.func, f"用 {g.func.upper()} 汇总") + + +def _rows(limit: str | None) -> str: + if limit is None: + return "不限行数" + return f"只取前 {limit} 行" if limit.isdigit() else f"按 LIMIT {limit} 取行" + + +def _skip(offset: str | None) -> str: + if offset is None: + return "不跳过行" + return f"跳过前 {offset} 行" if offset.isdigit() else f"按 OFFSET {offset} 跳过行" + + +class _Compare: + """一次比较：两边的叫法、列的中文名，以及按哪种写法比。 + + 两边各有一张只在自己这边出现、又不是关联进来的表（orders 换成了 orders_archive）时，其余各方面按不写表名的 + loose 比：表换了这件事由一条「tables」说清，同一个条件、同一个聚合不再因为表名不同被重复报一遍。 + + 每个方面都先按 k 比，比不上的再按 ak（CTE、子查询不写编号）一对一地配一遍（见 unmatched）。 + """ + + def __init__(self, a: Facets, b: Facets, labels: tuple[str, str], names: Mapping[str, str] | None) -> None: + self.a, self.b = a, b + self.la, self.lb = labels + self.label = _labeler(names) + joined = {j.table.key for j in (*a.joins, *b.joins)} + self.ta = [t for t in a.tables if t.key not in {u.key for u in b.tables} and t.key not in joined] + self.tb = [t for t in b.tables if t.key not in {u.key for u in a.tables} and t.key not in joined] + self.swapped = len(self.ta) == 1 and len(self.tb) == 1 + + def k(self, part: Part | None) -> str | None: + if part is None: + return None + return part.loose if self.swapped else part.key + + def ak(self, part: Part | None) -> str | None: + """同 k，但 CTE、子查询不写编号（见 unmatched）。""" + if part is None: + return None + return part.loose if self.swapped else (part.anon or part.key) + + @staticmethod + def unmatched(xs: Iterable[Any], ys: Iterable[Any], key: Any, anon: Any) -> tuple[list[Any], list[Any]]: + """两边各自独有的项，各按原来的先后。先按 key 比；剩下的再按 anon（CTE、子查询不写编号）一对一地配，配上的 + 也不算不同。编号只看位置，一边多出一个排在前面的 CTE、子查询，后面的编号就错开了：同一个条件不能因此被说成 + 两边各有一个。编号对得上时 key 就分得清是哪个 CTE 的列，anon 只收拾 key 比不上的；两个 CTE 算出同名的列、 + 编号又错开时可能配错一对，宁可漏报。""" + xs, ys = list(xs), list(ys) + kx, ky = {key(x) for x in xs}, {key(y) for y in ys} + only_x = [x for x in xs if key(x) not in ky] + only_y = [y for y in ys if key(y) not in kx] + paired_x: set[int] = set() + paired_y: set[int] = set() + for i, x in enumerate(only_x): + j = next((j for j, y in enumerate(only_y) if j not in paired_y and anon(y) == anon(x)), None) + if j is not None: + paired_x.add(i) + paired_y.add(j) + return [x for i, x in enumerate(only_x) if i not in paired_x], \ + [y for j, y in enumerate(only_y) if j not in paired_y] + + @staticmethod + def shown(xs: list[str], ys: list[str], full_xs: list[str], full_ys: list[str]) -> tuple[str, str]: + """两边给人看的写法；读起来一样（各自只写列名）却不是同一个东西时，改成一律写表名的写法。""" + wa, wb = "、".join(xs), "、".join(ys) + return ("、".join(full_xs), "、".join(full_ys)) if wa == wb else (wa, wb) + + def named(self, parts: list[Part], *, full: bool = False) -> str: + """时间列：有中文名的写「支付时间 paid_at」，几列用顿号隔开。""" + out = [] + for p in parts: + text = p.full if full else p.text + name = self.label(p.columns[0]) if p.columns else None + out.append(_say(name, text) if name else text) + return "、".join(out) + + def one_sided(self, side: str, other: str, part: Part, *, group: bool = False) -> str: + if group: + return _say(side, "分组后只保留", part.text, "的组，", other, "没有这个条件") + if part.excludes: + return _say(side, "排除了", part.excludes, "的行，", other, "没有") + return _say(side, "只保留", part.text, "的行，", other, "没有这个条件") + + # ---- 各个方面 + + def time(self) -> list[FacetDiff]: + a, b, la, lb = self.a, self.b, self.la, self.lb + ca = list(_dedupe((t.column for t in a.time), self.k)) + cb = list(_dedupe((t.column for t in b.time), self.k)) + if any(self.unmatched(ca, cb, self.k, self.ak)): + wa, wb = self.named(ca), self.named(cb) + if wa == wb: + wa, wb = self.named(ca, full=True), self.named(cb, full=True) + if ca and cb: + text = _say(la, "按", wa, "筛时间，", lb, "按", wb) + elif ca: + text = _say(la, "按", wa, "筛时间，", lb, "没有按时间筛选") + else: + text = _say(la, "没有按时间筛选，", lb, "按", wb, "筛时间") + detail = {side: [{"table": p.columns[0][0], "column": p.columns[0][1]} for p in parts if p.columns] + for side, parts in (("a", ca), ("b", cb))} + return [FacetDiff("time", "、".join(p.text for p in ca) or None, "、".join(p.text for p in cb) or None, + text, detail)] + if any(self.unmatched(a.time, b.time, lambda t: self.k(t.condition), lambda t: self.ak(t.condition))): + ra = " AND ".join(t.condition.text for t in a.time) + rb = " AND ".join(t.condition.text for t in b.time) + return [FacetDiff("time_range", ra, rb, _say(la, "的时间范围是", ra, "，", lb, "是", rb))] + return [] + + def filters(self) -> list[FacetDiff]: + timed = {self.k(t.condition) for t in (*self.a.time, *self.b.time)} + fa = [p for p in self.a.where if self.k(p) not in timed] + fb = [p for p in self.b.where if self.k(p) not in timed] + only_a, only_b = self.unmatched(fa, fb, self.k, self.ak) + out: list[FacetDiff] = [] + for p in list(only_a): + column = _single_column(p, loose=self.swapped) + mine = [q for q in only_a if column and _single_column(q, loose=self.swapped) == column] + theirs = [q for q in only_b if column and _single_column(q, loose=self.swapped) == column] + if len(mine) == 1 and len(theirs) == 1: + q = theirs[0] + wa, wb = self.shown([p.text], [q.text], [p.full], [q.full]) + out.append(FacetDiff("filter", p.text, q.text, _say(self.la, "按", wa, "筛选，", self.lb, "按", wb, "筛选"))) + only_a.remove(p) + only_b.remove(q) + out += [FacetDiff("filter", p.text, None, self.one_sided(self.la, self.lb, p)) for p in only_a] + out += [FacetDiff("filter", None, p.text, self.one_sided(self.lb, self.la, p)) for p in only_b] + return out + + def aggregates(self) -> list[FacetDiff]: + la, lb = self.la, self.lb + + def key(g: Aggregate) -> tuple[str, str | None, bool]: + return g.func, self.k(g.arg), g.distinct + + def anon(g: Aggregate) -> tuple[str, str | None, bool]: + return g.func, self.ak(g.arg), g.distinct + + def pair(xs: list[Aggregate], ys: list[Aggregate], same: Any) -> list[tuple[Aggregate, Aggregate]]: + pairs = [] + for x in list(xs): + y = next((y for y in ys if same(x, y)), None) + if y is not None: + pairs.append((x, y)) + xs.remove(x) + ys.remove(y) + return pairs + + ga, gb = self.unmatched(self.a.aggregates, self.b.aggregates, key, anon) + # 先配参数相同的（同一列：函数或去重不同），再配函数相同的（同一种算法：列不同）；各剩一个的也配成一对 + pairs = pair(ga, gb, lambda x, y: self.k(x.arg) == self.k(y.arg)) + pairs += pair(ga, gb, lambda x, y: x.func == y.func) + if len(ga) == 1 and len(gb) == 1: + pairs.append((ga.pop(), gb.pop())) + out = [] + for x, y in pairs: + same_arg = self.k(x.arg) == self.k(y.arg) + wa, wb = self.shown([x.arg.text], [y.arg.text], [x.arg.full], [y.arg.full]) if not same_arg \ + else (x.arg.text, y.arg.text) + if same_arg and x.func == y.func: + # 只差去重：COUNT(DISTINCT order_id) 和 COUNT(order_id) + text = _say(la, "对", wa, _verb(x) + "，", lb, "没去重") if x.distinct \ + else _say(la, "对", wa, _verb(x) + "时没去重，", lb, _verb(y)) + elif same_arg and x.arg.key != "*": + # 同一列、算法不同：后半句不再重复列名 + text = _say(la, "对", wa, _verb(x) + "，", lb, _verb(y)) + else: + what_a = ["数的是行数"] if x.func == "count" and x.arg.key == "*" else ["对", wa, _verb(x)] + what_b = ["数的是行数"] if y.func == "count" and y.arg.key == "*" else ["对", wb, _verb(y)] + text = _say(la, *what_a, "，", lb, *what_b) + out.append(FacetDiff("aggregate", x.text, y.text, text)) + out += [FacetDiff("aggregate", g.text, None, _say(la, "多算了", g.text, "，", lb, "没有")) for g in ga] + out += [FacetDiff("aggregate", None, g.text, _say(lb, "多算了", g.text, "，", la, "没有")) for g in gb] + return out + + def joins(self) -> list[FacetDiff]: + """连接：先把完全一样的（类型、关联的表、条件）两边配掉，剩下的按关联的是哪张表配成一对、说类型或条件哪里 + 不同，再剩下的是只有一边关联了。直接按表配的话，编号错开时会把两个不相干的 CTE 配成一对，说成「关联条件 + 不同」；先配完全一样的，同一张表关联了两次、先后不同也不算不同。""" + la, lb = self.la, self.lb + rest_a, rest_b = self.unmatched(self.a.joins, self.b.joins, + lambda j: (j.kind, j.table.key, self.k(j.on)), + lambda j: (j.kind, j.table.anon, self.ak(j.on))) + pairs: list[tuple[JoinFacet, JoinFacet]] = [] + for table in (lambda j: j.table.key, lambda j: j.table.anon): + for x in list(rest_a): + y = next((y for y in rest_b if table(y) == table(x)), None) + if y is not None: + pairs.append((x, y)) + rest_a.remove(x) + rest_b.remove(y) + out: list[FacetDiff] = [] + for x, y in pairs: + if x.kind != y.kind: + kx, ky = _JOIN_WORDS.get(x.kind, x.kind), _JOIN_WORDS.get(y.kind, y.kind) + out.append(FacetDiff("join", kx, ky, _say(la, "用", kx, "关联", x.table.text, "表，", lb, "用", ky))) + if self.k(x.on) != self.k(y.on) and self.ak(x.on) != self.ak(y.on): + xa, ya = (x.on.text if x.on else "（无条件）"), (y.on.text if y.on else "（无条件）") + out.append(FacetDiff("join", xa, ya, _say(la, "关联", x.table.text, "表的条件是", xa, "，", lb, "是", ya))) + out += [FacetDiff("join", x.table.text, None, _say(la, "关联了", x.table.text, "表，", lb, "没有")) for x in rest_a] + out += [FacetDiff("join", None, y.table.text, _say(lb, "关联了", y.table.text, "表，", la, "没有")) for y in rest_b] + # 关联之外的表（FROM 的表、子查询里的表）只在一边出现的 + if self.swapped: + ta, tb = self.ta[0], self.tb[0] + return out + [FacetDiff("tables", ta.text, tb.text, _say(la, "查的是", ta.text, "表，", lb, "查的是", tb.text, + "表"))] + out += [FacetDiff("tables", t.text, None, _say(la, "用到了", t.text, "表，", lb, "没有")) for t in self.ta] + out += [FacetDiff("tables", None, t.text, _say(lb, "用到了", t.text, "表，", la, "没有")) for t in self.tb] + return out + + def group(self) -> list[FacetDiff]: + a, b = self.a.group_by, self.b.group_by + if not any(self.unmatched(a, b, self.k, self.ak)): + return [] + wa, wb = self.shown([p.text for p in a], [p.text for p in b], [p.full for p in a], [p.full for p in b]) + if wa and wb: + text = _say(self.la, "按", wa, "分组，", self.lb, "按", wb, "分组") + elif wa: + text = _say(self.la, "按", wa, "分组，", self.lb, "没有分组") + else: + text = _say(self.la, "没有分组，", self.lb, "按", wb, "分组") + return [FacetDiff("group_by", wa or None, wb or None, text)] + + def having(self) -> list[FacetDiff]: + only_a, only_b = self.unmatched(self.a.having, self.b.having, self.k, self.ak) + return [FacetDiff("having", p.text, None, self.one_sided(self.la, self.lb, p, group=True)) for p in only_a] + \ + [FacetDiff("having", None, p.text, self.one_sided(self.lb, self.la, p, group=True)) for p in only_b] + + def limit(self) -> list[FacetDiff]: + """取哪几行：最外层两边比；CTE、子查询里的按出现的先后一一配对（两边给它起的名字可以不同），整组一样就不报。""" + def outer(f: Facets) -> LimitFacet | None: + return next((x for x in f.limits if x.scope is None), None) + + out = self.rows(outer(self.a), outer(self.b)) + ia = [x for x in self.a.limits if x.scope is not None] + ib = [x for x in self.b.limits if x.scope is not None] + if [self.lk(x) for x in ia] != [self.lk(x) for x in ib] and \ + [self.lk(x, anon=True) for x in ia] != [self.lk(x, anon=True) for x in ib]: + for x, y in zip_longest(ia, ib): + out += self.rows(x, y) + return out + + def lk(self, x: LimitFacet | None, *, anon: bool = False) -> tuple[Any, ...] | None: + """一段取哪几行的比较用写法：行数、跳过的行数、排序键。anon 时排序键按 ak（见 unmatched）。""" + if x is None: + return None + key = self.ak if anon else self.k + return x.count, x.offset, tuple((key(p), desc) for p, desc in x.order) + + def sorting(self, x: LimitFacet, *, full: bool = False) -> str: + if not x.order: + return "没有排序" + keys = "、".join(_say(p.full if full else p.text, "降序" if desc else "升序") for p, desc in x.order) + return _say("按", keys, "排") + + def rows(self, x: LimitFacet | None, y: LimitFacet | None) -> list[FacetDiff]: + """一对段（一边可以没有）取的行有什么不同：行数、跳过的行数，两边都限了行时再比排序。""" + if self.lk(x) == self.lk(y) or self.lk(x, anon=True) == self.lk(y, anon=True): + return [] + sa, sb = (x.scope if x else None), (y.scope if y else None) + la = _say(self.la, "在", sa, "里") if sa else self.la + lb = _say(self.lb, "在", sb, "里") if sb and sb != sa else self.lb + out = [] + ca, cb = (x.count if x else None), (y.count if y else None) + if ca != cb: + out.append(FacetDiff("limit", ca, cb, _say(la, _rows(ca), "，", lb, _rows(cb)))) + oa, ob = (x.offset if x else None), (y.offset if y else None) + if oa != ob: + out.append(FacetDiff("limit", oa, ob, _say(la, _skip(oa), "，", lb, _skip(ob)))) + if x is not None and y is not None and self.lk(x)[2] != self.lk(y)[2] \ + and self.lk(x, anon=True)[2] != self.lk(y, anon=True)[2]: + wa, wb = self.sorting(x), self.sorting(y) + if wa == wb: + wa, wb = self.sorting(x, full=True), self.sorting(y, full=True) + out.append(FacetDiff("limit", wa, wb, _say(la, "取行前", wa, "，", lb, wb))) + return out + + +def diff_facets(a: Facets, b: Facets, labels: tuple[str, str] = ("A", "B"), *, + names: Mapping[str, str] | None = None) -> list[FacetDiff]: + """逐个方面比两份 facets，每处不同一条，按对结果影响大致从大到小排：时间、筛选、聚合、连接和表、分组、HAVING、 + 取哪几行。完全一致返回空列表。 + + labels 是两边的叫法（缺省 A / B，双盲复算可以传「原流程」「盲算」）；names 是列的中文名（{「表.列」或「列」: + 中文名}，可以用 catalog_labels 从数据目录取），给了的话时间列写成「支付时间 paid_at」。任一边解析不了就只返回 + 一条 parse，不逐项比较：拿半份结构去比，只会报出一堆不存在的不同。 + """ + la, lb = labels + if a.parse_error or b.parse_error: + if a.parse_error and b.parse_error: + text = _say(la, "：", a.parse_error, "；", lb, "：", b.parse_error, "；没有逐项比较") + else: + side, error = (la, a.parse_error) if a.parse_error else (lb, b.parse_error) + text = _say(side, "：", error or "", "，没有逐项比较") + return [FacetDiff("parse", a.parse_error, b.parse_error, text)] + cmp = _Compare(a, b, labels, names) + return cmp.time() + cmp.filters() + cmp.aggregates() + cmp.joins() + cmp.group() + cmp.having() + cmp.limit() + + +__all__ = ["MAX_SQL_CHARS", "PLACEHOLDER", "Aggregate", "ColumnRef", "FacetDiff", "Facets", "JoinFacet", "LimitFacet", + "NormalizedSql", "Part", "TimePredicate", "catalog_labels", "diff_facets", "facets", "normalize_sql"] diff --git a/backend/app/engine/spend.py b/backend/app/engine/spend.py new file mode 100644 index 00000000..0efa52a6 --- /dev/null +++ b/backend/app/engine/spend.py @@ -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"] diff --git a/backend/app/engine/verify_agent.py b/backend/app/engine/verify_agent.py new file mode 100644 index 00000000..ad0a39d2 --- /dev/null +++ b/backend/app/engine/verify_agent.py @@ -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"] diff --git a/backend/app/tools/datasource.py b/backend/app/tools/datasource.py index 575fc847..88aa8b03 100644 --- a/backend/app/tools/datasource.py +++ b/backend/app/tools/datasource.py @@ -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}, ) diff --git a/backend/tests/test_datasource_force_readonly.py b/backend/tests/test_datasource_force_readonly.py new file mode 100644 index 00000000..856a5cb3 --- /dev/null +++ b/backend/tests/test_datasource_force_readonly.py @@ -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)] diff --git a/backend/tests/test_sql_semantics.py b/backend/tests/test_sql_semantics.py new file mode 100644 index 00000000..1e62c3ab --- /dev/null +++ b/backend/tests/test_sql_semantics.py @@ -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) == [] diff --git a/backend/tests/test_verify_agent.py b/backend/tests/test_verify_agent.py new file mode 100644 index 00000000..0a8c04d0 --- /dev/null +++ b/backend/tests/test_verify_agent.py @@ -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 diff --git a/backend/tests/test_verify_settings.py b/backend/tests/test_verify_settings.py new file mode 100644 index 00000000..b94fd3fb --- /dev/null +++ b/backend/tests/test_verify_settings.py @@ -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", "节点上留空等于没写" diff --git a/backend/tests/test_verify_spend.py b/backend/tests/test_verify_spend.py new file mode 100644 index 00000000..54311657 --- /dev/null +++ b/backend/tests/test_verify_spend.py @@ -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)