From b459c3bb588ebff566552c5060f73b180db207b9 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:06:03 +0800 Subject: [PATCH 1/9] =?UTF-8?q?=E6=A0=B8=E6=9F=A5=E8=8A=B1=E8=B4=B9?= =?UTF-8?q?=E8=B4=A6=EF=BC=9Aengine/spend.py=20=E7=85=A7=E7=BB=93=E8=AE=BA?= =?UTF-8?q?=E5=8F=A5=E8=A3=81=E5=88=A4=E7=9A=84=E6=AF=8F=E6=97=A5=E8=AE=A1?= =?UTF-8?q?=E6=95=B0=E5=81=9A=E4=B8=80=E6=9C=AC=E6=A0=B8=E6=9F=A5=E7=94=A8?= =?UTF-8?q?=E7=9A=84=E6=AF=8F=E6=97=A5=E8=B4=A6=EF=BC=88=E8=AE=BE=E7=BD=AE?= =?UTF-8?q?=E8=A1=A8=20verify=5Fspend=20=E4=B8=80=E8=A1=8C=EF=BC=8C?= =?UTF-8?q?=E6=8C=89=E6=9C=AC=E5=9C=B0=E6=97=A5=E6=9C=9F=E6=BB=9A=E5=8A=A8?= =?UTF-8?q?=EF=BC=89=E3=80=82=E8=B0=83=E7=94=A8=E5=89=8D=E6=8C=89=E4=BC=B0?= =?UTF-8?q?=E7=AE=97=E9=A2=84=E7=95=99=EF=BC=8C=E5=8A=A0=E4=B8=8A=E8=BF=99?= =?UTF-8?q?=E7=AC=94=E4=BC=9A=E8=B6=85=E8=BF=87=E6=AF=8F=E6=97=A5=E4=B8=8A?= =?UTF-8?q?=E9=99=90=E5=B0=B1=E4=B8=8D=E8=AE=B0=E3=80=81=E8=BF=94=E5=9B=9E?= =?UTF-8?q?=20None=EF=BC=8C=E4=B8=8A=E9=99=90=E4=B8=BA=20None=20=E8=A1=A8?= =?UTF-8?q?=E7=A4=BA=E4=B8=8D=E9=99=90=EF=BC=9B=E8=B0=83=E7=94=A8=E5=90=8E?= =?UTF-8?q?=E6=8C=89=E5=AE=9E=E9=99=85=E7=94=A8=E9=87=8F=E6=89=BE=E8=A1=A5?= =?UTF-8?q?=EF=BC=8C=E8=B7=A8=E4=BA=86=E5=A4=A9=E4=B8=8D=E6=89=BE=E8=A1=A5?= =?UTF-8?q?=EF=BC=9B=E4=BB=B7=E6=A0=BC=E7=9B=AE=E5=BD=95=E9=87=8C=E6=B2=A1?= =?UTF-8?q?=E6=9C=89=E7=9A=84=E6=A8=A1=E5=9E=8B=E5=8F=AA=E8=AE=A1=E6=AC=A1?= =?UTF-8?q?=E6=95=B0=E3=80=82=E5=92=8C=E8=A3=81=E5=88=A4=E5=88=86=E5=BC=80?= =?UTF-8?q?=E8=AE=B0=EF=BC=9A=E6=A0=B8=E6=9F=A5=E4=B8=80=E6=AC=A1=E8=A6=81?= =?UTF-8?q?=E6=9F=A5=E5=A5=BD=E5=87=A0=E8=BD=AE=E5=BA=93=E3=80=81=E8=8A=B1?= =?UTF-8?q?=E5=BE=97=E5=A4=9A=EF=BC=8C=E8=AE=B0=E5=9C=A8=E5=90=8C=E4=B8=80?= =?UTF-8?q?=E6=9C=AC=E8=B4=A6=E4=B8=8A=E4=BC=9A=E6=8A=8A=E5=BD=93=E5=A4=A9?= =?UTF-8?q?=E7=9A=84=E8=A3=81=E5=88=A4=E6=8C=A4=E5=88=B0=E5=88=A4=E4=B8=8D?= =?UTF-8?q?=E4=BA=86=E3=80=81=E8=AE=A9=E6=8A=A5=E5=91=8A=E5=87=BA=E5=85=B7?= =?UTF-8?q?=E6=97=A0=E6=95=85=E9=99=8D=E6=A1=A3=E3=80=82=E8=AF=BB=E6=94=B9?= =?UTF-8?q?=E5=86=99=E6=8C=89=E4=BA=8B=E4=BB=B6=E5=BE=AA=E7=8E=AF=E5=90=84?= =?UTF-8?q?=E6=8C=81=E4=B8=80=E6=8A=8A=E9=94=81=EF=BC=8C=E5=B9=B6=E5=8F=91?= =?UTF-8?q?=E9=A2=84=E7=95=99=E4=B8=8D=E4=BC=9A=E4=B8=80=E8=B5=B7=E8=B6=8A?= =?UTF-8?q?=E8=BF=87=E4=B8=8A=E9=99=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/engine/spend.py | 98 ++++++++++++++++++++++++++ backend/tests/test_verify_spend.py | 107 +++++++++++++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 backend/app/engine/spend.py create mode 100644 backend/tests/test_verify_spend.py 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/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) From aa828021d7d1ac5dd496fefc30ca493760b5c36d Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:06:03 +0800 Subject: [PATCH 2/9] =?UTF-8?q?=E6=A0=B8=E6=9F=A5=E8=AE=BE=E7=BD=AE?= =?UTF-8?q?=EF=BC=9A=E8=AE=BE=E7=BD=AE=E9=87=8C=E5=8A=A0=20verify=20?= =?UTF-8?q?=E4=B8=80=E7=BB=84=EF=BC=8C=E5=8F=8C=E7=9B=B2=E5=A4=8D=E7=AE=97?= =?UTF-8?q?=E3=80=81=E5=8F=8D=E6=96=B9=E5=AE=A1=E7=A8=BF=E3=80=81=E5=AE=A1?= =?UTF-8?q?=E6=89=B9=E8=B4=A8=E8=AF=A2=E5=85=B1=E7=94=A8=E3=80=82=E6=A0=B8?= =?UTF-8?q?=E6=9F=A5=E6=A8=A1=E5=9E=8B=EF=BC=88=E6=8E=A5=E5=85=A5=E3=80=81?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=EF=BC=89=E7=A9=BA=E7=9D=80=E6=B2=BF=E7=94=A8?= =?UTF-8?q?=E8=AF=81=E6=8D=AE=E8=A3=81=E5=88=A4=E6=A8=A1=E5=9E=8B=E7=9A=84?= =?UTF-8?q?=E5=8F=96=E5=80=BC=E9=93=BE=EF=BC=9B=E5=8F=8C=E7=9B=B2=E5=A4=8D?= =?UTF-8?q?=E7=AE=97=20{=E5=90=AF=E7=94=A8,=20=E5=85=81=E8=AE=B8=E7=9A=84?= =?UTF-8?q?=E7=9B=B8=E5=AF=B9=E8=AF=AF=E5=B7=AE=E9=BB=98=E8=AE=A4=201e-6,?= =?UTF-8?q?=20=E6=9F=A5=E8=AF=A2=E6=AC=A1=E6=95=B0,=20=E6=97=B6=E9=95=BF}?= =?UTF-8?q?=E3=80=81=E5=8F=8D=E6=96=B9=E5=AE=A1=E7=A8=BF=20{=E5=90=AF?= =?UTF-8?q?=E7=94=A8,=20=E6=AF=8F=E4=BB=BD=E6=8A=A5=E5=91=8A=E6=9C=80?= =?UTF-8?q?=E5=A4=9A=E5=AE=A1=E5=87=A0=E5=8F=A5,=20=E6=AF=8F=E5=8F=A5?= =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E6=AC=A1=E6=95=B0,=20=E6=97=B6=E9=95=BF}?= =?UTF-8?q?=E3=80=81=E5=AE=A1=E6=89=B9=E8=B4=A8=E8=AF=A2=20{=E5=90=AF?= =?UTF-8?q?=E7=94=A8,=20=E6=AF=8F=E6=AC=A1=E6=9F=A5=E8=AF=A2=E6=AC=A1?= =?UTF-8?q?=E6=95=B0,=20=E6=97=B6=E9=95=BF}=EF=BC=8C=E5=8F=A6=E6=9C=89?= =?UTF-8?q?=E6=A0=B8=E6=9F=A5=E6=AF=8F=E6=97=A5=E9=87=91=E9=A2=9D=E4=B8=8A?= =?UTF-8?q?=E9=99=90=EF=BC=8C=E4=B8=8A=E9=99=90=E5=AD=98=20null=20?= =?UTF-8?q?=E8=A1=A8=E7=A4=BA=E4=B8=8D=E9=99=90=E3=80=82=E5=86=99=E6=B3=95?= =?UTF-8?q?=E7=85=A7=E8=A3=81=E5=88=A4=E4=B8=80=E7=BB=84=EF=BC=9APUT=20?= =?UTF-8?q?=E5=85=88=E6=9F=A5=E5=90=8E=E5=86=99=EF=BC=8C=E6=9C=89=E4=B8=80?= =?UTF-8?q?=E9=A1=B9=E4=B8=8D=E5=AF=B9=E6=95=B4=E4=B8=AA=E8=AF=B7=E6=B1=82?= =?UTF-8?q?=E5=9B=9E=20422=E3=80=81=E8=AF=B4=E6=98=8E=E5=86=99=E8=AE=BE?= =?UTF-8?q?=E7=BD=AE=E9=A1=B5=E4=B8=8A=E7=9A=84=E5=8F=AB=E6=B3=95=EF=BC=9B?= =?UTF-8?q?=E5=BA=93=E9=87=8C=E5=AD=98=E5=9D=8F=E7=9A=84=E9=A1=B9=E8=AF=BB?= =?UTF-8?q?=E5=87=BA=E6=9D=A5=E6=8C=89=E9=BB=98=E8=AE=A4=EF=BC=8C=E4=B8=89?= =?UTF-8?q?=E9=A1=B9=E5=90=84=E8=87=AA=E9=82=A3=E4=B8=80=E7=BB=84=E9=80=90?= =?UTF-8?q?=E9=A1=B9=E8=A1=A5=E3=80=82verify=5Fsettings(session)=20?= =?UTF-8?q?=E7=BB=99=E8=8A=82=E7=82=B9=E5=92=8C=E8=B4=A8=E8=AF=A2=E6=8E=A5?= =?UTF-8?q?=E5=8F=A3=E8=AF=BB=EF=BC=8Cverify=5Fmodel=5Fspec(session,=20nod?= =?UTF-8?q?e=5Fconfig)=20=E6=8C=89=E8=8A=82=E7=82=B9=20=E2=86=92=20?= =?UTF-8?q?=E8=AE=BE=E7=BD=AE=20=E2=86=92=20judge=5Fmodel=5Fspec=20?= =?UTF-8?q?=E5=8F=96=E6=A8=A1=E5=9E=8B=EF=BC=8C=E6=95=B4=E7=BB=84=E5=8F=96?= =?UTF-8?q?=E4=B8=80=E7=BA=A7=E3=80=81=E4=B8=8D=E8=B7=A8=E7=BA=A7=E6=8B=BC?= =?UTF-8?q?=E3=80=82=E6=A0=B8=E6=9F=A5=E7=9A=84=E6=AF=8F=E6=97=A5=E8=AE=A1?= =?UTF-8?q?=E6=95=B0=E4=B8=8D=E4=BB=8E=E8=AE=BE=E7=BD=AE=E6=8E=A5=E5=8F=A3?= =?UTF-8?q?=E8=BF=9B=E5=87=BA=EF=BC=8C=E5=8F=A6=E5=8A=A0=E5=8F=AA=E8=AF=BB?= =?UTF-8?q?=E7=9A=84=20GET=20/api/settings/verify/spend=20=E7=BB=99?= =?UTF-8?q?=E8=AE=BE=E7=BD=AE=E9=A1=B5=E7=9C=8B=E4=BB=8A=E5=A4=A9=E8=8A=B1?= =?UTF-8?q?=E4=BA=86=E5=A4=9A=E5=B0=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/api/settings.py | 179 +++++++++++++++++++++- backend/tests/test_verify_settings.py | 205 ++++++++++++++++++++++++++ 2 files changed, 377 insertions(+), 7 deletions(-) create mode 100644 backend/tests/test_verify_settings.py 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/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", "节点上留空等于没写" From 031aae9343156be8e65f7ec95d7c59ef14dadfc7 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:22:51 +0800 Subject: [PATCH 3/9] =?UTF-8?q?SQL=20=E8=AF=AD=E4=B9=89=EF=BC=9Adata/sql?= =?UTF-8?q?=5Fsemantics.py=20=E7=BB=99=E5=8F=8C=E7=9B=B2=E5=A4=8D=E7=AE=97?= =?UTF-8?q?=E3=80=81=E8=BF=90=E8=A1=8C=E5=AF=B9=E6=AF=94=E3=80=81=E5=8F=8D?= =?UTF-8?q?=E6=96=B9=E5=AE=A1=E7=A8=BF=E5=85=B1=E7=94=A8=E3=80=82normalize?= =?UTF-8?q?=5Fsql=20=E6=8C=89=E6=96=B9=E8=A8=80=E8=A7=A3=E6=9E=90=E3=80=81?= =?UTF-8?q?=E9=87=8D=E6=96=B0=E7=94=9F=E6=88=90=EF=BC=88=E5=8E=BB=E6=B3=A8?= =?UTF-8?q?=E9=87=8A=E3=80=81=E7=BB=9F=E4=B8=80=E5=85=B3=E9=94=AE=E5=AD=97?= =?UTF-8?q?=E5=92=8C=E7=A9=BA=E7=99=BD=E3=80=81=E6=8C=89=E6=96=B9=E8=A8=80?= =?UTF-8?q?=E8=BD=AC=E6=8D=A2=E4=B8=8D=E5=8A=A0=E5=BC=95=E5=8F=B7=E7=9A=84?= =?UTF-8?q?=E5=90=8D=E5=AD=97=EF=BC=89=EF=BC=8C=E5=86=8D=E6=8A=8A=E5=AD=97?= =?UTF-8?q?=E7=AC=A6=E4=B8=B2=E5=92=8C=E6=95=B0=E5=AD=97=E6=8D=A2=E6=88=90?= =?UTF-8?q?=E5=8D=A0=E4=BD=8D=E7=AC=A6=EF=BC=9A=E5=8F=AA=E6=94=B9=E4=BA=86?= =?UTF-8?q?=E5=AD=97=E9=9D=A2=E9=87=8F=E7=9A=84=E4=B8=A4=E6=9D=A1=20SQL=20?= =?UTF-8?q?=E8=A7=84=E8=8C=83=E6=96=87=E6=9C=AC=E7=9B=B8=E5=90=8C=EF=BC=88?= =?UTF-8?q?=E5=AE=9A=E6=97=B6=E8=BF=90=E8=A1=8C=E6=AF=8F=E6=9C=9F=E6=97=A5?= =?UTF-8?q?=E6=9C=9F=E4=B8=8D=E5=90=8C=EF=BC=8C=E4=B8=8D=E7=AE=97=E6=A8=A1?= =?UTF-8?q?=E6=9D=BF=E5=8F=98=E4=BA=86=EF=BC=89=EF=BC=8C=E5=AD=97=E9=9D=A2?= =?UTF-8?q?=E9=87=8F=E6=8C=89=E8=A7=84=E8=8C=83=E6=96=87=E6=9C=AC=E9=87=8C?= =?UTF-8?q?=E7=9A=84=E5=85=88=E5=90=8E=E5=88=97=E5=87=BA=EF=BC=9B=E8=A7=A3?= =?UTF-8?q?=E6=9E=90=E4=B8=8D=E4=BA=86=E9=80=80=E5=9B=9E=E9=80=90=E8=AF=8D?= =?UTF-8?q?=E8=A7=84=E8=8C=83=E5=8C=96=EF=BC=8C=E7=85=A7=E6=A0=B7=E8=AE=A4?= =?UTF-8?q?=E5=BE=97=E5=87=BA=E5=8F=AA=E6=94=B9=E4=BA=86=E5=AD=97=E9=9D=A2?= =?UTF-8?q?=E9=87=8F=E3=80=82facets=20=E6=8A=8A=E6=9F=A5=E8=AF=A2=E6=8B=86?= =?UTF-8?q?=E6=88=90=E8=A1=A8=E3=80=81=E8=BF=9E=E6=8E=A5=EF=BC=88=E7=B1=BB?= =?UTF-8?q?=E5=9E=8B=20/=20=E5=85=B3=E8=81=94=E8=BF=9B=E6=9D=A5=E7=9A=84?= =?UTF-8?q?=E8=A1=A8=20/=20ON=20=E6=9D=A1=E4=BB=B6=EF=BC=89=E3=80=81WHERE?= =?UTF-8?q?=20=E5=90=88=E5=8F=96=E5=BC=8F=E3=80=81=E6=97=B6=E9=97=B4?= =?UTF-8?q?=E6=9D=A1=E4=BB=B6=EF=BC=88=E5=88=97=20/=20=E8=BF=90=E7=AE=97?= =?UTF-8?q?=E7=AC=A6=20/=20=E5=8F=96=E5=80=BC=EF=BC=89=E3=80=81=E5=88=86?= =?UTF-8?q?=E7=BB=84=E3=80=81=E8=81=9A=E5=90=88=EF=BC=88=E5=87=BD=E6=95=B0?= =?UTF-8?q?=20/=20=E5=8F=82=E6=95=B0=20/=20=E5=8E=BB=E9=87=8D=EF=BC=89?= =?UTF-8?q?=E3=80=81HAVING=E3=80=81LIMIT=EF=BC=9A=E5=88=AB=E5=90=8D?= =?UTF-8?q?=E6=8D=A2=E6=88=90=E7=9C=9F=E5=AE=9E=E8=A1=A8=E5=90=8D=EF=BC=8C?= =?UTF-8?q?=E6=AF=94=E8=BE=83=E4=B8=A4=E8=BE=B9=E6=8E=92=E5=A5=BD=E9=A1=BA?= =?UTF-8?q?=E5=BA=8F=EF=BC=8CGROUP=20BY=201=20=E5=92=8C=E5=88=AB=E5=90=8D?= =?UTF-8?q?=E6=8D=A2=E6=88=90=E8=A1=A8=E8=BE=BE=E5=BC=8F=EF=BC=8CCTE?= =?UTF-8?q?=E3=80=81=E5=AD=90=E6=9F=A5=E8=AF=A2=E9=80=90=E6=AE=B5=E6=94=B6?= =?UTF-8?q?=E9=9B=86=EF=BC=8C=E5=8F=AA=E9=80=8F=E5=87=BA=E4=B8=80=E5=BC=A0?= =?UTF-8?q?=E8=A1=A8=E7=9A=84=20CTE=20=E6=8D=A2=E5=9B=9E=E9=82=A3=E5=BC=A0?= =?UTF-8?q?=E8=A1=A8=E7=9A=84=E5=88=97=EF=BC=8C=E9=80=97=E5=8F=B7=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5=E5=86=99=E5=9C=A8=20WHERE=20=E9=87=8C=E7=9A=84?= =?UTF-8?q?=E7=AD=89=E5=BC=8F=E7=AE=97=E5=85=B3=E8=81=94=E6=9D=A1=E4=BB=B6?= =?UTF-8?q?=EF=BC=8C=E7=9B=B8=E5=85=B3=E5=AD=90=E6=9F=A5=E8=AF=A2=E5=88=B0?= =?UTF-8?q?=E5=A4=96=E5=B1=82=E6=89=BE=E5=88=AB=E5=90=8D=EF=BC=9B=E8=A7=A3?= =?UTF-8?q?=E6=9E=90=E4=B8=8D=E4=BA=86=E5=B8=A6=20parse=5Ferror=EF=BC=8C?= =?UTF-8?q?=E4=BB=8E=E4=B8=8D=E6=8A=9B=E3=80=82diff=5Ffacets=20=E6=AF=8F?= =?UTF-8?q?=E5=A4=84=E4=B8=8D=E5=90=8C=E7=BB=99=E4=B8=80=E5=8F=A5=E4=B8=AD?= =?UTF-8?q?=E6=96=87=EF=BC=88=E3=80=8CA=20=E6=8C=89=E6=94=AF=E4=BB=98?= =?UTF-8?q?=E6=97=B6=E9=97=B4=20paid=5Fat=20=E7=AD=9B=E6=97=B6=E9=97=B4?= =?UTF-8?q?=EF=BC=8CB=20=E6=8C=89=E4=B8=8B=E5=8D=95=E6=97=B6=E9=97=B4=20cr?= =?UTF-8?q?eated=5Fat=E3=80=8D=E3=80=8CA=20=E6=8E=92=E9=99=A4=E4=BA=86=20i?= =?UTF-8?q?s=5Fdeleted=20=3D=201=20=E7=9A=84=E8=A1=8C=EF=BC=8CB=20?= =?UTF-8?q?=E6=B2=A1=E6=9C=89=E3=80=8D=E3=80=8CA=20=E5=AF=B9=20order=5Fid?= =?UTF-8?q?=20=E5=8E=BB=E9=87=8D=E8=AE=A1=E6=95=B0=EF=BC=8CB=20=E6=B2=A1?= =?UTF-8?q?=E5=8E=BB=E9=87=8D=E3=80=8D=EF=BC=89=EF=BC=8C=E4=B8=A4=E8=BE=B9?= =?UTF-8?q?=E5=8F=AB=E6=B3=95=E5=8F=AF=E4=BB=A5=E6=8D=A2=E6=88=90=E3=80=8C?= =?UTF-8?q?=E5=8E=9F=E6=B5=81=E7=A8=8B=E3=80=8D=E3=80=8C=E7=9B=B2=E7=AE=97?= =?UTF-8?q?=E3=80=8D=EF=BC=8C=E5=88=97=E7=9A=84=E4=B8=AD=E6=96=87=E5=90=8D?= =?UTF-8?q?=E5=8F=AF=E4=BB=A5=E7=94=A8=20catalog=5Flabels=20=E4=BB=8E?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E7=9B=AE=E5=BD=95=E5=8F=96=EF=BC=9B=E5=8F=AA?= =?UTF-8?q?=E6=8D=A2=E4=BA=86=E4=B8=80=E5=BC=A0=E8=A1=A8=E6=97=B6=E8=AF=B4?= =?UTF-8?q?=E4=B8=80=E6=AC=A1=E3=80=8C=E6=9F=A5=E7=9A=84=E6=98=AF=E5=93=AA?= =?UTF-8?q?=E5=BC=A0=E8=A1=A8=E3=80=8D=EF=BC=8C=E5=85=B6=E4=BD=99=E6=8C=89?= =?UTF-8?q?=E4=B8=8D=E5=86=99=E8=A1=A8=E5=90=8D=E7=9A=84=E5=86=99=E6=B3=95?= =?UTF-8?q?=E6=AF=94=E3=80=82=E5=A4=8D=E7=94=A8=20sqlcheck=20=E7=9A=84=20d?= =?UTF-8?q?ialect=5Fof=E3=80=81=5Fconjuncts=E3=80=81=5Fbare=5Fcolumn?= =?UTF-8?q?=E3=80=81=5Funparen?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/data/sql_semantics.py | 1122 +++++++++++++++++++++++++++ backend/tests/test_sql_semantics.py | 418 ++++++++++ 2 files changed, 1540 insertions(+) create mode 100644 backend/app/data/sql_semantics.py create mode 100644 backend/tests/test_sql_semantics.py diff --git a/backend/app/data/sql_semantics.py b/backend/app/data/sql_semantics.py new file mode 100644 index 00000000..cc07568c --- /dev/null +++ b/backend/app/data/sql_semantics.py @@ -0,0 +1,1122 @@ +"""SQL 语义:两条 SQL 算的是不是一回事,不是的话差在哪。 + +双盲复算要回答「盲算和原流程的 SQL 差在哪」(A 按支付时间筛、B 按下单时间筛;A 排除了作废的行、B 没有),运行 +对比要回答「这一期和上一期的 SQL 是不是同一个模板」(定时运行每次的日期不同,不能算模板变了),反方审稿要看自己 +的查询和原查询的口径是否一致。三处共用这里: + +- normalize_sql:字面量换成占位符之后的规范文本,加上字面量本身(按规范文本里的顺序)。两条 SQL 只有字面量不同, + 规范文本就相同;关键字大小写、空白、注释、不加引号的名字的大小写都不算不同。 +- facets:把一条查询拆成几个方面——用到的表、连接(类型、关联进来的表、ON 条件)、WHERE 的各个合取式、其中的 + 时间条件(列、运算符、取值)、分组、聚合(函数、参数、是否去重)、HAVING、LIMIT。别名一律换成真实表名, + 等式、比较的两边排好顺序:o.paid_at >= '2026-09-01' 和 '2026-09-01' <= orders.paid_at 是同一个条件。 +- diff_facets:逐个方面比两份 facets,每处不同给一句中文说明。 + +**逐段收集。** CTE、子查询里的条件和聚合也算进来:盲算可能把有效记录条件写在 CTE 里,原流程写在外层,按 +「整条语句里有哪些条件」比才不会把同一个口径报成不同。逗号连接(以及 INNER JOIN 之外又写在 WHERE 里的等值关联) +写在 WHERE 里的列等式算作关联条件,不算筛选。 + +**和 sqlcheck 一样宁可漏报。** 解析不了(语法错、方言特性、多条语句、不是查询)就带 parse_error 返回,从不抛 +异常。没加表前缀、又有好几张表可能有这一列的,照写列名(给了表结构才去对)。只读语法树,不执行 SQL。 +""" +from __future__ import annotations + +import logging +import re +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass, fields, is_dataclass +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 + + +# ========================================================================== +# 解析 +# ========================================================================== + + +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 用来比较:别名换成真实表名,表名、列名按 name_key,等式和比较的两边排好顺序;loose 同 key 但不写表名,两边 + 查的表换了一张(orders 换成 orders_archive)时拿它比其余各方面,免得同一个条件只因表名不同被报成不同。 + text 给人看:好几张表都有的列写「表.列」(真实表名,不是别名),其余只写列名;full 一律写表名,两边的 text + 一样、key 却不同时改用它说明。excludes 只在 WHERE 的「排除」写法上有:<>、NOT、NOT IN、IS NOT NULL 去掉 + 否定之后排除的是什么(IS NOT NULL 写成「列 为空」)。 + """ + + key: str + text: str + columns: tuple[ColumnRef, ...] = () + excludes: str | None = None + loose: str = "" + full: str = "" + + @staticmethod + def plain(key: str, text: str) -> "Part": + """没有列的一项(表名、COUNT(*) 的「*」):四种写法相同。""" + return Part(key=key, text=text, loose=key, full=text) + + +@dataclass(frozen=True) +class JoinFacet: + """一次连接。kind:inner / left / right / full / cross / semi / anti / natural。on 是关联条件的合取式排好序后 + 用 AND 连起来的写法,没有条件是 None。""" + + kind: str + table: Part + on: Part | None = None + + +@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 Facets: + """一条查询的各个方面。每一项都按在 SQL 里出现的先后排、按 key 去重;parse_error 不为空时其余都是空的。""" + + 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 + 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 _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、子查询)/ other(表函数、认不出的别名)。""" + + kind: str + name: str # 给人看:表名(有库名就带上)、CTE 名、子查询别名 + key: str # 比较用:name_key + 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 + self.ctes: dict[str, Any] = {} + for cte in root.find_all(exp.CTE): + if cte.alias: + self.ctes.setdefault(name_key(cte.alias), cte.this) + 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 _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)] + items = ([select.args["from_"].this] if select.args.get("from_") is not None else []) + \ + [j.this for j in select.args.get("joins") or []] + out = _Sources() + for i, node in enumerate(items): + alias = node.alias_or_name if isinstance(node, (exp.Table, exp.Subquery)) else "" + 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, name_key(node.name), (node.this.copy(),), + self.ctes[name_key(node.name)]) + else: + out[slot] = self._table(node) + elif isinstance(node, exp.Subquery): + out[slot] = _Source("derived", alias or "子查询", slot, (exp.to_identifier(alias),) if alias else (), + node.this) + 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 + + def render(self, node: Any, sources: dict[str, _Source], mode: str = "text") -> str: + """四种写法(见 Part):key、loose 比较用,名字按 name_key,key 写真实表名、loose 不写;text、full 给人看, + text 只在好几张表都有这一列时写真实表名,full 一律写(关联条件说的就是两张表各出哪一列,用 full)。""" + from sqlglot import exp + + copy = node.copy() + for col in _own_columns(copy): + source, name = self.resolve(col, sources) + quoted = isinstance(col.this, exp.Identifier) and bool(col.this.args.get("quoted")) + if mode in ("key", "loose"): + col.set("this", exp.to_identifier(name_key(name))) + col.set("table", exp.to_identifier(source.key) if source is not None and mode == "key" else None) + col.set("db", None) + col.set("catalog", None) + continue + 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) + return copy.sql(dialect=self.read, comments=False) + + def canonical(self, node: Any, sources: dict[str, _Source]) -> Any: + """比较的两边排好顺序(列在左、字面量在右;两边都有列的按 key 排),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)) + if (rc and not lc) or (lc and rc and self.render(left, sources, "key") + > self.render(right, sources, "key")): + 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 part(self, node: Any, sources: dict[str, _Source], *, excludes: bool = False) -> Part: + refs: dict[tuple[str | None, str], None] = {} + for col in _own_columns(node): + source, name = self.resolve(col, sources) + refs.setdefault((source.name if source is not None else None, 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) + + 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) + clause = select.args.get("having") + having += [self.part(self.canonical(c, sources), sources) + for c in (_conjuncts(clause.this) if clause is not None else [])] + 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(), + ) + + 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 = first.this.alias_or_name if first is not None and \ + isinstance(first.this, (exp.Table, exp.Subquery)) else "" + drafts: list[dict[str, Any]] = [] + for i, join in enumerate(select.args.get("joins") or []): + target = join.this + alias = target.alias_or_name if isinstance(target, (exp.Table, exp.Subquery)) else "" + 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}) + 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(source.key, source.name) if source is not None else Part.plain("?", "?") + triples = sorted((self.render(c, sources, "key"), self.render(c, sources, "loose"), + self.render(c, sources, "full")) for c in draft["conds"]) + on = Part(key=" AND ".join(t[0] for t in triples), text=" AND ".join(t[2] for t in triples), + loose=" AND ".join(sorted(t[1] for t in triples)), full=" AND ".join(t[2] for t in triples)) \ + if triples else None + return JoinFacet(kind=draft["kind"], table=table, on=on) + + 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)) + + def _group(self, select: Any, sources: dict[str, _Source]) -> list[Part]: + """分组键。GROUP BY 1 换成 SELECT 里对应的表达式,GROUP BY 别名换成别名指的表达式(MySQL、SQLite 的读法)。""" + from sqlglot import exp + + group = select.args.get("group") + if group is None: + return [] + selects = select.selects + by_alias = {name_key(s.alias): s.this for s in selects if isinstance(s, exp.Alias) and s.alias} + out = [] + for item in group.expressions: + 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): + target = 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 by_alias: + target = by_alias[name_key(target.name)] + out.append(self.part(self.canonical(target, 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 + + +# ========================================================================== +# 比较 +# ========================================================================== + + +@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: + """只涉及一列的条件:那一列(表.列;loose 时只看列名),用来把两边对同一列的不同写法配成一对。""" + if len(part.columns) != 1: + return None + table, column = part.columns[0] + return name_key(column) if loose else f"{name_key(table or '')}.{name_key(column)}" + + +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} 取行" + + +class _Compare: + """一次比较:两边的叫法、列的中文名,以及按哪种写法比。 + + 两边各有一张只在自己这边出现、又不是关联进来的表(orders 换成了 orders_archive)时,其余各方面按不写表名的 + loose 比:表换了这件事由一条「tables」说清,同一个条件、同一个聚合不再因为表名不同被重复报一遍。 + """ + + 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 + + @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 {self.k(p) for p in ca} != {self.k(p) for p in cb}: + 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 {self.k(t.condition) for t in a.time} != {self.k(t.condition) for t in b.time}: + 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 = [p for p in fa if self.k(p) not in {self.k(q) for q in fb}] + only_b = [p for p in fb if self.k(p) not in {self.k(q) for q in fa}] + 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 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 = [g for g in self.a.aggregates if key(g) not in {key(h) for h in self.b.aggregates}] + gb = [g for g in self.b.aggregates if key(g) not in {key(h) for h in self.a.aggregates}] + # 先配参数相同的(同一列:函数或去重不同),再配函数相同的(同一种算法:列不同);各剩一个的也配成一对 + 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]: + la, lb = self.la, self.lb + + def index(f: Facets) -> dict[str, list[JoinFacet]]: + out: dict[str, list[JoinFacet]] = {} + for j in f.joins: + out.setdefault(j.table.key, []).append(j) + return out + + ja, jb = index(self.a), index(self.b) + out: list[FacetDiff] = [] + for key in dict.fromkeys([*ja, *jb]): + xs, ys = ja.get(key, []), jb.get(key, []) + for x, y in zip(xs, ys): + 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): + 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 xs[len(ys):]] + out += [FacetDiff("join", None, y.table.text, _say(lb, "关联了", y.table.text, "表,", la, "没有")) + for y in ys[len(xs):]] + # 关联之外的表(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 {self.k(p) for p in a} == {self.k(p) for p in b}: + 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]: + ka, kb = {self.k(p) for p in self.a.having}, {self.k(p) for p in self.b.having} + return [FacetDiff("having", p.text, None, self.one_sided(self.la, self.lb, p, group=True)) + for p in self.a.having if self.k(p) not in kb] + \ + [FacetDiff("having", None, p.text, self.one_sided(self.lb, self.la, p, group=True)) + for p in self.b.having if self.k(p) not in ka] + + def limit(self) -> list[FacetDiff]: + a, b = self.a.limit, self.b.limit + return [FacetDiff("limit", a, b, _say(self.la, _rows(a), ",", self.lb, _rows(b)))] if a != b else [] + + +def diff_facets(a: Facets, b: Facets, labels: tuple[str, str] = ("A", "B"), *, + names: Mapping[str, str] | None = None) -> list[FacetDiff]: + """逐个方面比两份 facets,每处不同一条,按对结果影响大致从大到小排:时间、筛选、聚合、连接和表、分组、HAVING、 + LIMIT。完全一致返回空列表。 + + 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", + "NormalizedSql", "Part", "TimePredicate", "catalog_labels", "diff_facets", "facets", "normalize_sql"] diff --git a/backend/tests/test_sql_semantics.py b/backend/tests/test_sql_semantics.py new file mode 100644 index 00000000..ee3cd7e1 --- /dev/null +++ b/backend/tests/test_sql_semantics.py @@ -0,0 +1,418 @@ +"""SQL 语义(data/sql_semantics.py):两条 SQL 算的是不是一回事,不是的话差在哪。 + +- normalize_sql:字面量换成占位符后的规范文本 + 字面量。只有字面量不同的两条 SQL 规范文本相同(定时运行每期的日期 + 不同,不能算模板变了);关键字大小写、空白、注释也不算不同;解析不了的退回逐词规范化,照样认得出只改了字面量 +- facets:表、连接(类型 / 关联进来的表 / ON 条件)、WHERE 合取式、时间条件(列 / 运算符 / 取值)、分组、聚合 + (函数 / 参数 / 去重)、HAVING、LIMIT。别名换成真实表名、比较的两边排好顺序;解析不了带 parse_error,不抛 +- diff_facets:每处不同一句中文:「A 按支付时间 paid_at 筛时间,B 按下单时间 created_at」「A 排除了 + is_deleted = 1 的行,B 没有」「A 对 order_id 去重计数,B 没去重」……写法不同、口径相同的不报 + +夹具全是虚构的表:订单 orders、退款 refunds、门店 stores,以及一张中文表「订单」。 +""" +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" + # 别名换成真实表名,两边排好顺序,几个条件按 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 [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" + + +@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] + assert facets("SELECT orders.id FROM orders WHERE NOT EXISTS " + "(SELECT 1 FROM refunds WHERE orders.id = refunds.order_id)").where[1].key == f.where[1].key + + +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 = '虚构市'"), +]) +def test_the_same_meaning_written_differently_is_no_difference(a, b): + assert texts(a, b) == [] + + +@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) == [] From e46ff55fd4b4dd826707c8f13f37c6f5b62e0eb4 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:26:04 +0800 Subject: [PATCH 4/9] =?UTF-8?q?=E6=95=B0=E6=8D=AE=E6=BA=90=E6=9F=A5?= =?UTF-8?q?=E8=AF=A2=E5=B7=A5=E5=85=B7=E5=8A=A0=20force=5Freadonly?= =?UTF-8?q?=EF=BC=9A=E6=A0=B8=E6=9F=A5=E7=9A=84=20agent=20=E5=8F=A6?= =?UTF-8?q?=E8=B5=B7=E4=B8=80=E8=B7=AF=E3=80=81=E6=B2=A1=E6=9C=89=E5=AE=A1?= =?UTF-8?q?=E6=89=B9=E5=85=B3=E5=8D=A1=E5=85=9C=E5=BA=95=EF=BC=8C=E8=A6=81?= =?UTF-8?q?=E6=9F=A5=E7=9A=84=E6=BA=90=E5=8D=B4=E5=8F=AF=E8=83=BD=E5=8F=AF?= =?UTF-8?q?=E5=86=99=E3=80=82=E5=BC=80=E4=BA=86=E5=B0=B1=E4=B8=8D=E7=AE=A1?= =?UTF-8?q?=E6=BA=90=E5=8F=AF=E4=B8=8D=E5=8F=AF=E5=86=99=EF=BC=8C=E4=B8=80?= =?UTF-8?q?=E5=BE=8B=E6=8C=89=E5=8F=AA=E8=AF=BB=E5=AE=88=E5=8D=AB=E6=A3=80?= =?UTF-8?q?=E6=9F=A5=EF=BC=9A=E5=85=88=E8=BF=87=E4=B8=80=E9=81=8D=E5=92=8C?= =?UTF-8?q?=20run=5Fquery=20=E7=9B=B8=E5=90=8C=E7=9A=84=E6=A3=80=E6=9F=A5?= =?UTF-8?q?=EF=BC=88=E5=A4=9A=E6=9D=A1=E8=AF=AD=E5=8F=A5=E3=80=81=E7=BB=93?= =?UTF-8?q?=E6=9E=84=E5=8F=98=E6=9B=B4=E7=85=A7=E5=8E=9F=E6=9D=A5=E7=9A=84?= =?UTF-8?q?=E8=AF=B4=E6=B3=95=E6=8B=92=EF=BC=89=EF=BC=8C=E5=86=8D=E6=8C=89?= =?UTF-8?q?=E5=8F=AA=E8=AF=BB=E5=88=A4=EF=BC=8C=E5=86=99=E8=AF=AD=E5=8F=A5?= =?UTF-8?q?=E4=BA=A4=E5=9B=9E=E3=80=8C=E8=BF=99=E6=AC=A1=E6=9F=A5=E8=AF=A2?= =?UTF-8?q?=E5=8F=AA=E5=85=81=E8=AE=B8=E8=AF=BB=E5=8F=96=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E3=80=8D=E3=80=81=E4=B8=8D=E6=89=A7=E8=A1=8C=EF=BC=8C=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=8E=9F=E6=A0=B7=EF=BC=9B=E6=9F=A5=E8=AF=A2=E7=85=A7?= =?UTF-8?q?=E5=B8=B8=E8=90=BD=E6=9F=A5=E8=AF=A2=E5=BF=AB=E7=85=A7=E3=80=82?= =?UTF-8?q?=E5=B7=A5=E5=85=B7=E6=8F=8F=E8=BF=B0=E8=B7=9F=E7=9D=80=E5=86=99?= =?UTF-8?q?=E3=80=8C=E5=8F=AA=E8=AF=BB=E3=80=8D=EF=BC=8C=E5=AE=A1=E6=89=B9?= =?UTF-8?q?=E5=88=A4=E5=AE=9A=E4=B9=9F=E4=B8=8D=E5=86=8D=E6=8A=8A=E5=86=99?= =?UTF-8?q?=E8=AF=AD=E5=8F=A5=E5=BD=93=E5=8D=B1=E9=99=A9=E8=B0=83=E7=94=A8?= =?UTF-8?q?=E3=80=82=E5=8F=AA=E9=9D=A0=E5=AE=88=E5=8D=AB=E8=BF=99=E4=B8=80?= =?UTF-8?q?=E9=81=93=EF=BC=9A=E8=BF=9E=E6=8E=A5=E6=B1=A0=E6=8C=89=E6=BA=90?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E3=80=81=E4=B8=8D=E5=88=86=E5=8F=AA=E8=AF=BB?= =?UTF-8?q?=E4=B8=8E=E5=90=A6=EF=BC=8C=E4=B8=8D=E8=83=BD=E4=B8=BA=E6=AD=A4?= =?UTF-8?q?=E5=8F=A6=E8=B5=B7=E4=B8=80=E4=B8=AA=E5=8F=AA=E8=AF=BB=E6=B1=A0?= =?UTF-8?q?=EF=BC=88=E7=BC=93=E5=AD=98=E9=94=AE=E7=9B=B8=E5=90=8C=EF=BC=8C?= =?UTF-8?q?=E4=BC=9A=E6=8A=8A=E8=BF=99=E4=B8=AA=E6=BA=90=E5=B9=B3=E6=97=B6?= =?UTF-8?q?=E7=94=A8=E7=9A=84=E5=8F=AF=E5=86=99=E8=BF=9E=E6=8E=A5=E6=8D=A2?= =?UTF-8?q?=E6=88=90=E5=8F=AA=E8=AF=BB=EF=BC=89=E3=80=82build=5Fdatasource?= =?UTF-8?q?=5Ftools=20=E9=80=8F=E4=BC=A0=E8=BF=99=E4=B8=AA=E5=8F=82?= =?UTF-8?q?=E6=95=B0=EF=BC=8C=E4=B8=8D=E4=BC=A0=E6=97=B6=E4=B8=80=E5=88=87?= =?UTF-8?q?=E7=85=A7=E6=97=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/tools/datasource.py | 41 +++++-- .../tests/test_datasource_force_readonly.py | 106 ++++++++++++++++++ 2 files changed, 138 insertions(+), 9 deletions(-) create mode 100644 backend/tests/test_datasource_force_readonly.py 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)] From 1efa5f7b20665ace6d0a37d94f483be6d44c2c28 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:34:17 +0800 Subject: [PATCH 5/9] =?UTF-8?q?=E6=A0=B8=E6=9F=A5=E7=94=A8=E7=9A=84?= =?UTF-8?q?=E5=8F=AA=E8=AF=BB=20agent=EF=BC=9Aengine/verify=5Fagent.py=20?= =?UTF-8?q?=E7=9A=84=20run=5Fverify=5Fagent=20=E7=BB=99=E5=8F=8C=E7=9B=B2?= =?UTF-8?q?=E5=A4=8D=E7=AE=97=E3=80=81=E5=8F=8D=E6=96=B9=E5=AE=A1=E7=A8=BF?= =?UTF-8?q?=E3=80=81=E5=AE=A1=E6=89=B9=E8=B4=A8=E8=AF=A2=E5=85=B1=E7=94=A8?= =?UTF-8?q?=E3=80=82=E5=B7=A5=E5=85=B7=E5=8F=AA=E5=BB=BA=E7=BB=99=E5=AE=9A?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E6=BA=90=E7=9A=84=20db=5Fschema=20/=20db=5Fq?= =?UTF-8?q?uery=EF=BC=8C=E4=B8=80=E5=BE=8B=20force=5Freadonly=EF=BC=9B?= =?UTF-8?q?=E8=87=AA=E5=B7=B1=E7=AE=A1=E5=B7=A5=E5=85=B7=E8=B0=83=E7=94=A8?= =?UTF-8?q?=E5=BE=AA=E7=8E=AF=EF=BC=88=E4=B8=8D=E4=BE=9D=E8=B5=96=20NodeCo?= =?UTF-8?q?ntext=20/=20@task=EF=BC=89=EF=BC=8C=E6=9F=A5=E8=AF=A2=E6=AC=A1?= =?UTF-8?q?=E6=95=B0=E3=80=81=E6=97=B6=E9=95=BF=E3=80=81=E9=87=91=E9=A2=9D?= =?UTF-8?q?=EF=BC=88=E8=BF=99=E4=B8=80=E6=AC=A1=E3=80=81=E5=85=A8=E5=B1=80?= =?UTF-8?q?=E6=AF=8F=E6=97=A5=EF=BC=89=E4=BB=BB=E4=BD=95=E4=B8=80=E9=A1=B9?= =?UTF-8?q?=E5=88=B0=E9=A1=B6=E5=B0=B1=E5=81=9C=E6=AD=A2=E6=9F=A5=E8=AF=A2?= =?UTF-8?q?=E3=80=81=E6=8C=89=E5=B7=B2=E6=9C=89=E7=BB=93=E6=9E=9C=E6=94=B6?= =?UTF-8?q?=E5=B0=BE=EF=BC=8C=E6=94=B6=E5=B0=BE=E9=82=A3=E6=AC=A1=E6=97=B6?= =?UTF-8?q?=E9=95=BF=E5=88=B0=E9=A1=B6=E5=8F=A6=E7=BB=99=E5=AE=BD=E9=99=90?= =?UTF-8?q?=EF=BC=8C=E9=87=91=E9=A2=9D=E6=AF=8F=E4=B8=80=E8=BD=AE=E9=83=BD?= =?UTF-8?q?=E7=BB=99=E5=AE=83=E7=95=99=E5=87=BA=E9=92=B1=EF=BC=9B=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E4=B8=8D=E5=9C=A8=E4=BB=B7=E6=A0=BC=E7=9B=AE=E5=BD=95?= =?UTF-8?q?=E9=87=8C=E6=97=B6=E9=87=91=E9=A2=9D=E4=B8=8A=E9=99=90=E6=97=A0?= =?UTF-8?q?=E6=95=88=EF=BC=8C=E7=85=A7=E8=A3=81=E5=88=A4=E7=9A=84=E8=AF=B4?= =?UTF-8?q?=E6=B3=95=E8=AE=B0=E4=B8=80=E6=9D=A1=E8=AF=B4=E6=98=8E=EF=BC=9B?= =?UTF-8?q?=E6=AF=8F=E6=AC=A1=E8=B0=83=E7=94=A8=E8=AE=B0=E8=BF=9B=E6=A0=B8?= =?UTF-8?q?=E6=9F=A5=E7=9A=84=E6=AF=8F=E6=97=A5=E8=B4=A6=E3=80=82=E6=88=90?= =?UTF-8?q?=E5=8A=9F=E7=9A=84=E6=9F=A5=E8=AF=A2=E6=8C=89=E6=9C=AC=20agent?= =?UTF-8?q?=20=E8=87=AA=E5=B7=B1=E7=9A=84=E9=A1=BA=E5=BA=8F=E7=BC=96=20Q1?= =?UTF-8?q?=E3=80=81Q2=E2=80=A6=EF=BC=8C=E7=BB=99=E4=BA=86=20emit=20?= =?UTF-8?q?=E6=89=8D=E5=8F=91=20llm=20/=20tool=20=E4=BA=8B=E4=BB=B6?= =?UTF-8?q?=EF=BC=88role:=20verify=EF=BC=8C=E6=9F=A5=E8=AF=A2=E7=9A=84=20t?= =?UTF-8?q?ool.end=20=E5=B8=A6=20query=5Fartifact=EF=BC=8C=E9=9A=8F?= =?UTF-8?q?=E8=BF=90=E8=A1=8C=E5=B0=81=E5=AD=98=E3=80=81=E8=83=BD=E4=B8=8B?= =?UTF-8?q?=E9=92=BB=EF=BC=89=EF=BC=8C=E7=BB=9D=E4=B8=8D=E5=BE=80=E8=AF=81?= =?UTF-8?q?=E6=8D=AE=E5=8F=B0=E8=B4=A6=E5=86=99=20kind:=20query=EF=BC=8C?= =?UTF-8?q?=E5=85=8D=E5=BE=97=E8=BF=9B=E4=BA=86=E6=8A=A5=E5=91=8A=E7=9A=84?= =?UTF-8?q?=E5=86=99=E4=BD=9C=E7=9B=AE=E5=BD=95=E3=80=82=E6=9C=80=E5=90=8E?= =?UTF-8?q?=E6=8C=89=E8=B0=83=E7=94=A8=E6=96=B9=E7=9A=84=E7=BB=93=E6=9E=84?= =?UTF-8?q?=E4=BA=A4=E7=BB=93=E6=9E=9C=EF=BC=9A=E6=95=B0=E5=80=BC=E5=AD=97?= =?UTF-8?q?=E6=AE=B5=E6=8D=A2=E6=88=90=E6=A0=BC=E5=AD=90=E5=BC=95=E7=94=A8?= =?UTF-8?q?=20{alias,=20row,=20column}=EF=BC=8C=E7=94=B1=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E4=BB=8E=E5=BF=AB=E7=85=A7=E5=8F=96=E5=80=BC=EF=BC=8C=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E8=87=AA=E5=B7=B1=E5=86=99=E7=9A=84=E6=95=B0=E5=AD=97?= =?UTF-8?q?=E4=B8=8D=E9=87=87=E4=BF=A1=E3=80=81=E5=8F=96=E4=B8=8D=E5=88=B0?= =?UTF-8?q?=E6=98=AF=20None=EF=BC=9B=E7=90=86=E7=94=B1=E9=87=8C=E7=9A=84?= =?UTF-8?q?=E6=95=B0=E5=AD=97=E5=86=99=E6=88=90=20[[v:Qn.rX.=E5=88=97]]?= =?UTF-8?q?=EF=BC=8C=E7=94=A8=E6=9C=AC=20agent=20=E8=87=AA=E5=B7=B1?= =?UTF-8?q?=E7=9A=84=E5=8F=B0=E8=B4=A6=20build=5Fcatalog=20+=20compose=5Fd?= =?UTF-8?q?oc=20=E7=BB=84=E8=A3=85=EF=BC=8C=E8=A3=B8=E6=95=B0=E5=AD=97?= =?UTF-8?q?=E4=BC=9A=E8=A2=AB=E6=A0=87=E5=87=BA=EF=BC=9B=E7=BB=93=E6=9E=84?= =?UTF-8?q?=E5=8C=96=E8=BE=93=E5=87=BA=E5=A4=B1=E8=B4=A5=E9=80=80=E5=9B=9E?= =?UTF-8?q?=E6=99=AE=E9=80=9A=E8=B0=83=E7=94=A8=E6=8A=A0=20JSON=E3=80=82?= =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E8=A1=A8=E6=A0=BC=E6=8C=89=E8=B0=83=E7=94=A8?= =?UTF-8?q?=E6=96=B9=E7=BB=99=E7=9A=84=E6=95=B0=E6=8D=AE=E7=89=88=E6=9C=AC?= =?UTF-8?q?=E6=9F=A5=E3=80=82=E6=B0=B8=E8=BF=9C=E4=B8=8D=E6=8A=9B=EF=BC=8C?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E6=B2=A1=E9=85=8D=E5=A5=BD=E3=80=81=E8=B0=83?= =?UTF-8?q?=E7=94=A8=E5=A4=B1=E8=B4=A5=E3=80=81=E8=A7=A6=E9=A1=B6=E9=83=BD?= =?UTF-8?q?=E8=90=BD=E5=9C=A8=E8=BF=94=E5=9B=9E=E5=80=BC=E9=87=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/engine/verify_agent.py | 659 +++++++++++++++++++++++++++++ backend/tests/test_verify_agent.py | 483 +++++++++++++++++++++ 2 files changed, 1142 insertions(+) create mode 100644 backend/app/engine/verify_agent.py create mode 100644 backend/tests/test_verify_agent.py diff --git a/backend/app/engine/verify_agent.py b/backend/app/engine/verify_agent.py new file mode 100644 index 00000000..eabc69ce --- /dev/null +++ b/backend/app/engine/verify_agent.py @@ -0,0 +1,659 @@ +"""核查用的只读 agent:查库,交回有出处的结果。双盲复算、反方审稿、审批质询共用这一个。 + +三处要的是同一件事:另起一个模型,只给它几个数据源的只读查询工具,让它自己写 SQL 去核对一个数、一句结论、一个 +提问,最后交回结构化的结果——而且结果里的数要能核对。 + +**只读。** 工具只建给定数据源的 db_schema / db_query,建的时候 force_readonly:可写源上的写语句也被守卫拒绝 +(tools/datasource.py)。这一路没有审批关卡,不能指望人在半路拦下一条 DELETE。 + +**数从快照取。** 每次成功的查询是一份查询快照(工具自己落的 put_json),按这个 agent 自己的顺序编 Q1、Q2…。 +最后交结果时,调用方结构里的数值字段一律改成格子引用 {alias, row, column},由代码按引用从快照里取值:模型自己 +写的数字不采信,取不到就是 None,不兜底成 0。理由文字里的数字要写成 [[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 的宽限),所以金额上限每一轮都给它留出钱。模型不在价格目录里 +时估不出金额,金额上限对它无效,照裁判的说法记一条说明。每次调用都记进核查的每日账(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 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 +from app.engine.judge import _json_in, _tokens, _usd +from app.engine.nodes.llm import _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 +#: 结果里自动加上的理由字段(调用方的结构里已经有同名的字符串字段就沿用它) +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] + + +@dataclass +class VerifyOutcome: + """一次核查的结果。 + + result:调用方结构的形状,数值字段是从快照取到的值(取不到为 None),另有 rationale(理由渲染后的正文); + 没交成结果时为 None,error 说明为什么。cells:每个数值字段的出处 {path, status, ref, alias, row, column, + artifact, eid, value},status 是 resolved / missing(模型说查不到)/ 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_number(schema: Any) -> bool: + return bool(_types(schema) & {"number", "integer"}) + + +def result_schema(schema: dict[str, Any]) -> dict[str, Any]: + """交给模型的结构:调用方结构里的数值字段(number / integer,嵌套的对象、数组里的也算)一律换成可为 null 的 + 格子引用,另加一个理由字段(已有同名的字符串字段就沿用)。其余字段原样。""" + + def rewrite(node: Any) -> Any: + if not isinstance(node, dict): + return node + if _is_number(node): + out: dict[str, Any] = {"anyOf": [_CELL_REF, {"type": "null"}]} + if node.get("description"): + out["description"] = f"{node['description']}(写格子引用,查不到写 null)" + return out + if _is_object(node): + props = node.get("properties") or {} + return {**node, "type": "object", "properties": {k: rewrite(v) for k, v in props.items()}} + if "array" in _types(node) and isinstance(node.get("items"), dict): + return {**node, "items": rewrite(node["items"])} + return node + + root = rewrite(schema if _is_object(schema) 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} + + +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, schema: dict[str, Any]) -> Any: + """一个数值字段:按引用取值,记下出处,返回值(取不到为 None)。""" + if value is None: + self.cells.append({"path": path, "status": "missing"}) + return None + if not isinstance(value, dict): + # 模型直接写了数:不采信。照记下来,界面上能看到它想填什么 + self.cells.append({"path": path, "status": "unresolved", "model_value": value, + "reason": "模型直接写了数值,没有给出格子引用,不采信"}) + 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, schema, column_kind(content, name)) + self.cells.append({"path": path, "status": "resolved", "ref": f"{alias}.r{row}.{name}", "alias": alias, + "row": row, "column": name, "artifact": query["artifact"], + "eid": cell_eid(query["artifact"], row, name), "value": truth}) + return truth + + def walk(self, schema: Any, value: Any, path: str = "") -> Any: + """按调用方的结构走一遍模型交的结果:数值字段换成从快照取到的值,其余原样。""" + if _is_number(schema): + return self.take(value, path, schema) + if _is_object(schema): + given = value if isinstance(value, dict) else {} + return {k: self.walk(sub, given.get(k), f"{path}.{k}" if path else k) + for k, sub in (schema.get("properties") or {}).items()} + if "array" in _types(schema) and isinstance(schema.get("items"), dict): + if not isinstance(value, list): + return None + return [self.walk(schema["items"], v, f"{path}[{i}]") for i, v in enumerate(value)] + return value + + +# -------------------------------------------------------------------------- +# 花费 +# -------------------------------------------------------------------------- + + +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) -> float: + return pricing.estimate_cost(self.model_id, _tokens(text), out) if self.priced else 0.0 + + async def reserve(self, estimate: float, *, keep: float = 0.0) -> tuple[str | None, str | None]: + """(记在哪天, None) 或 (None, 拦下的上限)。keep 是要额外留出的钱(收尾那次调用):只用来判断,不预留。""" + 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" + cap = self.daily_max_usd if self.priced else None + 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, stopped: str | None + ) -> list[BaseMessage]: + """收尾那次调用的消息:任务、查到的结果(每次查询列出行号)、模型最后说的话,以及交结果的规则。""" + parts = [f"任务:\n{task.strip()}"] + if stopped: + parts.append(f"已达上限({stopped}),不能再查询,按已经查到的结果交结果。") + parts.append(f"你查到的结果:\n{_digest(queries) if queries else '(没有成功的查询)'}") + parts.append(f"你最后的说明(只作参考,数以查询结果为准):\n{said.strip() or '(没有)'}") + parts.append(VERIFY_FINAL_PROMPT) + return [SystemMessage(content=system), HumanMessage(content="\n\n".join(parts))] + + +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 + + +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 = None, + 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 对象)。数值字段会换成格子引用交给模型、由代码取值。 + - 上限: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 时上传源一律报错,不拿当前版本顶上)。 + """ + started = time.monotonic() + deadline = started + timeout_s if timeout_s is not None else None + usage: dict[str, Any] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cost_usd": 0.0, "calls": 0, + "unpriced_calls": 0} + notes: list[str] = [] + queries: list[dict[str, Any]] = [] + ledger: list[dict[str, Any]] = [] + 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() + + def outcome(**fields: Any) -> VerifyOutcome: + 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 queries] + base: dict[str, Any] = {"result": None, "cells": [], "queries": shown, "rationale": None, "usage": usage, + "notes": notes, "duration_ms": int((time.monotonic() - started) * 1000)} + return VerifyOutcome(**{**base, **fields}) + + # ---- 模型 + 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})") + 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) + + 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 + 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=budget.estimate(_text_of(messages), _OUT_FINAL)) + if refused: + stopped = refused + break + 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(model=model_id, 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 + schema = args_model_of(tool) + if schema is not None: + try: + args, _ = prepare_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" + + words = None + if stopped: + words = _limit_words(stopped, max_queries=max_queries, timeout_s=timeout_s, max_cost_usd=max_cost_usd, + daily_max_usd=daily_max_usd, steps=steps) + await _announce(emit, EventType.LOG, level="warn", code="verify_limit", limit=stopped, + message=f"核查已达上限({words}),已停止查询,按已有结果收尾") + + # ---- 收尾:按结构交结果 + schema = result_schema(output_schema) + final = _final_messages(head, task, queries, said, words) + estimate = budget.estimate(_text_of(final), _OUT_FINAL) + day, refused = await budget.reserve(estimate) + if refused: + limit_words = _limit_words(refused, max_queries=max_queries, timeout_s=timeout_s, max_cost_usd=max_cost_usd, + daily_max_usd=daily_max_usd, steps=steps) + return outcome(model=model_id, limit=stopped or refused, answered=answered, + error=f"已达{limit_words},没能交结果") + remaining = left() + grace = None if remaining is None else max(remaining, SETTLE_GRACE_S) + finish = None if grace is None else time.monotonic() + grace + try: + async with SessionLocal() as session: + settler, _ = await get_chat_model(session, spec.model_copy(update={"thinking": "off"})) + except Exception as e: # noqa: BLE001 + await spend.settle(day or "", -estimate) + return outcome(model=model_id, limit=stopped, answered=answered, error=f"核查模型不可用({describe_exception(e)})") + + def until_finish() -> float | None: + return None if finish is None else finish - time.monotonic() + + got, raw, failure = await ask(final, settler.with_structured_output(schema, include_raw=True).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 抠出来 + remaining = until_finish() + if remaining is None or remaining > 0: + estimate = budget.estimate(_text_of(final), _OUT_FINAL) + day, refused = await budget.reserve(estimate) + if not refused: + got, raw, failure = await ask(final, 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(model=model_id, limit=stopped, answered=answered, error=why) + + # ---- 取值、组装理由 + cells = _Cells(queries) + result = cells.walk(output_schema, parsed) if _is_object(output_schema) 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(model=model_id, 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 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 ToolTimeout as e: + 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 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", "VerifyOutcome", "result_schema", "run_verify_agent"] diff --git a/backend/tests/test_verify_agent.py b/backend/tests/test_verify_agent.py new file mode 100644 index 00000000..242317dc --- /dev/null +++ b/backend/tests/test_verify_agent.py @@ -0,0 +1,483 @@ +"""核查用的只读 agent(engine/verify_agent.py):查库、交回有出处的结果。双盲复算、反方审稿、审批质询共用。 + +- 工具只有给定数据源的 db_schema / db_query,建的时候 force_readonly:可写源上的写语句也被拒绝,数据不动 +- 每次成功的查询是一份查询快照,按这个 agent 自己的顺序编 Q1、Q2…;节点里调用时发 tool.end(带 query_artifact、 + role: verify),但不往证据台账里写 query:核查的查询进了报告的写作目录,写作者就会把它们当成可以引用的证据 +- 最后按调用方的结构交结果:数值字段一律是格子引用 {alias, row, column},由代码从快照取值,模型自己写的数字 + 不采信;理由里的数字要写成 [[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.verify_agent import 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 == spend.SPEND_KEY)) + await session.commit() + yield + + +@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, "收尾那次也没赶上:照实说没交成结果" + + +@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) + + +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"] + + +# -------------------------------------------------------------------------- +# 数据版本 +# -------------------------------------------------------------------------- + + +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 From f1c068d6a99307deefa3f24c20248cdd73d0e879 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 00:42:02 +0800 Subject: [PATCH 6/9] =?UTF-8?q?=E6=A0=B8=E6=9F=A5=E5=B7=A5=E5=85=B7?= =?UTF-8?q?=E7=AE=B1=E6=94=B6=E5=B0=BE=EF=BC=9A=E8=BF=9E=E6=8E=A5=E7=9A=84?= =?UTF-8?q?=E6=96=B9=E9=9D=A2=E8=A1=A5=E4=B8=8A=E5=8F=A6=E4=B8=80=E4=BE=A7?= =?UTF-8?q?=EF=BC=88=E5=85=B3=E8=81=94=E6=9D=A1=E4=BB=B6=E9=87=8C=E5=87=BA?= =?UTF-8?q?=E7=8E=B0=E7=9A=84=E5=88=AB=E7=9A=84=E8=A1=A8=EF=BC=8CCROSS=20J?= =?UTF-8?q?OIN=20=E5=86=99=E5=AE=83=E4=B9=8B=E5=89=8D=E7=9A=84=E6=9D=A5?= =?UTF-8?q?=E6=BA=90=EF=BC=89=EF=BC=8C=E5=8F=8C=E7=9B=B2=E5=A4=8D=E7=AE=97?= =?UTF-8?q?=E5=AF=B9=E7=85=A7=E6=95=B0=E6=8D=AE=E7=9B=AE=E5=BD=95=E9=87=8C?= =?UTF-8?q?=E7=9A=84=E5=85=B3=E7=B3=BB=E6=97=B6=E7=94=A8=E5=BE=97=E4=B8=8A?= =?UTF-8?q?=EF=BC=9B=E5=8F=AA=E8=AF=BB=20agent=20=E6=94=B6=E5=B0=BE?= =?UTF-8?q?=E5=89=8D=E5=85=88=E6=8B=BF=E5=88=B0=E6=94=B6=E5=B0=BE=E7=94=A8?= =?UTF-8?q?=E7=9A=84=E6=A8=A1=E5=9E=8B=E5=86=8D=E9=A2=84=E7=95=99=E9=87=91?= =?UTF-8?q?=E9=A2=9D=EF=BC=8C=E6=A8=A1=E5=9E=8B=E6=8B=BF=E4=B8=8D=E5=88=B0?= =?UTF-8?q?=E6=97=B6=E4=B8=8D=E5=9C=A8=E6=AF=8F=E6=97=A5=E8=B4=A6=E4=B8=8A?= =?UTF-8?q?=E7=99=BD=E8=AE=B0=E4=B8=80=E6=AC=A1=E8=B0=83=E7=94=A8=EF=BC=9B?= =?UTF-8?q?=E6=94=B6=E5=B0=BE=E6=B6=88=E6=81=AF=E7=9A=84=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E6=94=B9=E5=90=8D=20limit=5Fwords=EF=BC=8C=E8=AF=B4=E6=B8=85?= =?UTF-8?q?=E6=98=AF=E5=93=AA=E4=B8=AA=E4=B8=8A=E9=99=90=E7=9A=84=E8=AF=B4?= =?UTF-8?q?=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/data/sql_semantics.py | 14 ++++++++++---- backend/app/engine/verify_agent.py | 22 +++++++++++----------- backend/tests/test_sql_semantics.py | 5 ++++- backend/tests/test_verify_agent.py | 3 ++- 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/backend/app/data/sql_semantics.py b/backend/app/data/sql_semantics.py index cc07568c..d9830c25 100644 --- a/backend/app/data/sql_semantics.py +++ b/backend/app/data/sql_semantics.py @@ -217,12 +217,14 @@ class Part: @dataclass(frozen=True) class JoinFacet: - """一次连接。kind:inner / left / right / full / cross / semi / anti / natural。on 是关联条件的合取式排好序后 - 用 AND 连起来的写法,没有条件是 None。""" + """一次连接。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) @@ -688,7 +690,8 @@ class _Statement: 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}) + 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: @@ -720,7 +723,10 @@ class _Statement: on = Part(key=" AND ".join(t[0] for t in triples), text=" AND ".join(t[2] for t in triples), loose=" AND ".join(sorted(t[1] for t in triples)), full=" AND ".join(t[2] for t in triples)) \ if triples else None - return JoinFacet(kind=draft["kind"], table=table, on=on) + 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),另一边没有列;列名像日期,或者用了日期函数、像日期的字符串。""" diff --git a/backend/app/engine/verify_agent.py b/backend/app/engine/verify_agent.py index eabc69ce..4294f043 100644 --- a/backend/app/engine/verify_agent.py +++ b/backend/app/engine/verify_agent.py @@ -323,12 +323,13 @@ def _text_of(messages: list[BaseMessage]) -> str: for m in messages) -def _final_messages(system: str, task: str, queries: list[dict[str, Any]], said: str, stopped: str | None +def _final_messages(system: str, task: str, queries: list[dict[str, Any]], said: str, limit_words: str | None ) -> list[BaseMessage]: - """收尾那次调用的消息:任务、查到的结果(每次查询列出行号)、模型最后说的话,以及交结果的规则。""" + """收尾那次调用的消息:任务、查到的结果(每次查询列出行号)、模型最后说的话,以及交结果的规则。 + limit_words:触顶时是哪个上限(「最多查询 6 次」),告诉模型为什么不能再查。""" parts = [f"任务:\n{task.strip()}"] - if stopped: - parts.append(f"已达上限({stopped}),不能再查询,按已经查到的结果交结果。") + if limit_words: + parts.append(f"已达上限({limit_words}),不能再查询,按已经查到的结果交结果。") parts.append(f"你查到的结果:\n{_digest(queries) if queries else '(没有成功的查询)'}") parts.append(f"你最后的说明(只作参考,数以查询结果为准):\n{said.strip() or '(没有)'}") parts.append(VERIFY_FINAL_PROMPT) @@ -541,7 +542,12 @@ async def run_verify_agent( 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(model=model_id, limit=stopped, answered=answered, error=f"核查模型不可用({describe_exception(e)})") schema = result_schema(output_schema) final = _final_messages(head, task, queries, said, words) estimate = budget.estimate(_text_of(final), _OUT_FINAL) @@ -554,12 +560,6 @@ async def run_verify_agent( remaining = left() grace = None if remaining is None else max(remaining, SETTLE_GRACE_S) finish = None if grace is None else time.monotonic() + grace - try: - async with SessionLocal() as session: - settler, _ = await get_chat_model(session, spec.model_copy(update={"thinking": "off"})) - except Exception as e: # noqa: BLE001 - await spend.settle(day or "", -estimate) - return outcome(model=model_id, limit=stopped, answered=answered, error=f"核查模型不可用({describe_exception(e)})") def until_finish() -> float | None: return None if finish is None else finish - time.monotonic() diff --git a/backend/tests/test_sql_semantics.py b/backend/tests/test_sql_semantics.py index ee3cd7e1..690fd351 100644 --- a/backend/tests/test_sql_semantics.py +++ b/backend/tests/test_sql_semantics.py @@ -116,7 +116,7 @@ def test_tables_are_the_real_tables_without_cte_names(): 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" + 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' " @@ -128,6 +128,9 @@ 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 = '虚构市'"] diff --git a/backend/tests/test_verify_agent.py b/backend/tests/test_verify_agent.py index 242317dc..7e72083f 100644 --- a/backend/tests/test_verify_agent.py +++ b/backend/tests/test_verify_agent.py @@ -212,7 +212,8 @@ async def test_writes_are_refused_on_a_writable_source(monkeypatch, 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]]。"}) + {"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 被拒绝:") From a427b53874c243a35200db3acffa98dad7fff1cc Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 01:46:55 +0800 Subject: [PATCH 7/9] =?UTF-8?q?SQL=20=E8=AF=AD=E4=B9=89=E5=A4=8D=E6=A0=B8?= =?UTF-8?q?=E4=BF=AE=E6=AD=A3=EF=BC=9A=E5=90=8C=E4=B8=80=E4=B8=AA=E5=8F=A3?= =?UTF-8?q?=E5=BE=84=E4=B8=8D=E5=86=8D=E5=9B=A0=E4=B8=BA=E5=AD=90=E6=9F=A5?= =?UTF-8?q?=E8=AF=A2=E9=87=8C=E7=9A=84=E5=88=AB=E5=90=8D=E3=80=81HAVING=20?= =?UTF-8?q?=E9=87=8C=E7=9A=84=E5=88=AB=E5=90=8D=E8=A2=AB=E6=8A=A5=E6=88=90?= =?UTF-8?q?=E4=B8=8D=E5=90=8C=EF=BC=9B=E5=8F=96=E5=93=AA=E5=87=A0=E8=A1=8C?= =?UTF-8?q?=E9=80=90=E6=AE=B5=E6=AF=94=EF=BC=8C=E8=BF=9E=E5=90=8C=20OFFSET?= =?UTF-8?q?=20=E5=92=8C=E5=8F=96=E5=89=8D=E5=87=A0=E8=A1=8C=E6=97=B6?= =?UTF-8?q?=E7=9A=84=E6=8E=92=E5=BA=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 条件里套着的子查询(EXISTS、IN (SELECT …))原来按外层的来源整个照抄进外层那个条件,别名原样留着:带别名和不带别名的同一条 NOT EXISTS 被报成两处筛选不同,双盲复算就会说盲算换了口径。现在每段 SELECT 在语法树上记段号(复制出来的子查询靠它找回原树上那一段算好的来源,相关子查询照样到外层找),改写时每一列按它所在那一段认,子查询里的表去掉别名;比较两边排序时也按那一段的来源排。HAVING 引用 SELECT 别名(HAVING s > 100)换成别名指的表达式,和 GROUP BY 别名一个读法,聚合函数里的同名列不换。 取哪几行原来只看最外层 LIMIT 的行数:CTE 里 LIMIT 10 和 LIMIT 100、OFFSET 0 和 OFFSET 10、取前 10 行时降序和升序都报成没有不同,取到的却是不同的行。Facets 加 limits,逐段收集(最外层、CTE、FROM 里的子查询;条件里的子查询已经整个写在条件里,不单算,免得报两遍)行数、跳过的行数,以及这一段的排序(ORDER BY 1、别名换成 SELECT 里的表达式);OFFSET 0 算没写,排序只在限了行的段里比。原来的 limit(最外层行数)照旧保留。 --- backend/app/data/sql_semantics.py | 325 +++++++++++++++++++++++----- backend/tests/test_sql_semantics.py | 74 ++++++- 2 files changed, 346 insertions(+), 53 deletions(-) diff --git a/backend/app/data/sql_semantics.py b/backend/app/data/sql_semantics.py index d9830c25..0d7a2ad8 100644 --- a/backend/app/data/sql_semantics.py +++ b/backend/app/data/sql_semantics.py @@ -7,13 +7,15 @@ - normalize_sql:字面量换成占位符之后的规范文本,加上字面量本身(按规范文本里的顺序)。两条 SQL 只有字面量不同, 规范文本就相同;关键字大小写、空白、注释、不加引号的名字的大小写都不算不同。 - facets:把一条查询拆成几个方面——用到的表、连接(类型、关联进来的表、ON 条件)、WHERE 的各个合取式、其中的 - 时间条件(列、运算符、取值)、分组、聚合(函数、参数、是否去重)、HAVING、LIMIT。别名一律换成真实表名, - 等式、比较的两边排好顺序:o.paid_at >= '2026-09-01' 和 '2026-09-01' <= orders.paid_at 是同一个条件。 + 时间条件(列、运算符、取值)、分组、聚合(函数、参数、是否去重)、HAVING、取哪几行(LIMIT、OFFSET、取前几行 + 时的排序)。别名一律换成真实表名(条件里套着的子查询也是),等式、比较的两边排好顺序:o.paid_at >= '2026-09-01' + 和 '2026-09-01' <= orders.paid_at 是同一个条件。 - diff_facets:逐个方面比两份 facets,每处不同给一句中文说明。 -**逐段收集。** CTE、子查询里的条件和聚合也算进来:盲算可能把有效记录条件写在 CTE 里,原流程写在外层,按 +**逐段收集。** CTE、子查询里的条件、聚合和取行也算进来:盲算可能把有效记录条件写在 CTE 里,原流程写在外层,按 「整条语句里有哪些条件」比才不会把同一个口径报成不同。逗号连接(以及 INNER JOIN 之外又写在 WHERE 里的等值关联) -写在 WHERE 里的列等式算作关联条件,不算筛选。 +写在 WHERE 里的列等式算作关联条件,不算筛选。条件里套着的子查询(EXISTS、IN (SELECT …))既单算一段,也按它自己 +那一段的来源整个写进外层那个条件:别名在子查询里定义,拿外层的来源去认,同一个条件只因别名不同就会被报成不同。 **和 sqlcheck 一样宁可漏报。** 解析不了(语法错、方言特性、多条语句、不是查询)就带 parse_error 返回,从不抛 异常。没加表前缀、又有好几张表可能有这一列的,照写列名(给了表结构才去对)。只读语法树,不执行 SQL。 @@ -24,6 +26,7 @@ import logging import re 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 @@ -36,6 +39,9 @@ logger = logging.getLogger(__name__) PLACEHOLDER = "?" #: 太长的 SQL 不做语法分析(normalize_sql 退回逐词规范化,facets 带 parse_error):解析耗时随长度增长 MAX_SQL_CHARS = 100_000 +#: 每段 SELECT 在语法树上记段号的 meta 键。条件、表达式要复制一份再改写,复制出来的子查询靠段号找回原来那一段 +#: 算好的来源(别名、外层):复制品没有父节点,自己找不到外层 +_SEGMENT = "sql_semantics_segment" # ========================================================================== @@ -249,9 +255,25 @@ class Aggregate: 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 不为空时其余都是空的。""" + """一条查询的各个方面。每一项都按在 SQL 里出现的先后排、按 key 去重;parse_error 不为空时其余都是空的。 + + limit 是最外层 LIMIT 的行数(写法);limits 是各段取哪几行(最外层在前,见 LimitFacet),只有写了 LIMIT 或 + OFFSET 的段才有。 + """ dialect: str | None = None tables: tuple[Part, ...] = () @@ -262,6 +284,7 @@ class Facets: 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]: @@ -417,6 +440,11 @@ class _Statement: 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] = {} for cte in root.find_all(exp.CTE): if cte.alias: @@ -552,36 +580,82 @@ class _Statement: # ---- 写成 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、loose 比较用,名字按 name_key,key 写真实表名、loose 不写;text、full 给人看, - text 只在好几张表都有这一列时写真实表名,full 一律写(关联条件说的就是两张表各出哪一列,用 full)。""" + text 只在好几张表都有这一列时写真实表名,full 一律写(关联条件说的就是两张表各出哪一列,用 full)。 + + node 里套着的子查询逐段改写:每一列按它所在那一段的来源认,子查询里的表去掉别名(列已经写成了真实表名)。 + 子查询里拿两张表的列互相比的(相关子查询),给人看时和关联条件一样写表名。 + """ from sqlglot import exp copy = node.copy() - for col in _own_columns(copy): - source, name = self.resolve(col, sources) - quoted = isinstance(col.this, exp.Identifier) and bool(col.this.args.get("quoted")) - if mode in ("key", "loose"): - col.set("this", exp.to_identifier(name_key(name))) - col.set("table", exp.to_identifier(source.key) if source is not None and mode == "key" else None) - col.set("db", None) - col.set("catalog", None) - continue - 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) + 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"): + 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))) return copy.sql(dialect=self.read, comments=False) + 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"): + col.set("this", exp.to_identifier(name_key(name))) + col.set("table", exp.to_identifier(source.key) if source is not None and mode == "key" 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: - """比较的两边排好顺序(列在左、字面量在右;两边都有列的按 key 排),IN 的取值排好序。返回新的树。""" + """比较的两边排好顺序(列在左、字面量在右;两边都有列的按 key 排),IN 的取值排好序。返回新的树。 + 套着的子查询里的比较按它那一段的来源排:拿外层的来源去认子查询的别名,排出来的先后就取决于别名怎么起。""" from sqlglot import exp def fix(n: Any) -> Any: @@ -589,8 +663,8 @@ class _Statement: 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)) - if (rc and not lc) or (lc and rc and self.render(left, sources, "key") - > self.render(right, sources, "key")): + here = self.scope(n, sources) if lc and rc else sources + if (rc and not lc) or (lc and rc and self.render(left, here, "key") > self.render(right, here, "key")): 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): @@ -656,9 +730,7 @@ class _Statement: joins += [self._join(d, sources) for d in drafts] group += self._group(select, sources) aggregates += self._aggregates(select, sources) - clause = select.args.get("having") - having += [self.part(self.canonical(c, sources), sources) - for c in (_conjuncts(clause.this) if clause is not None else [])] + having += self._having(select, sources) return Facets( dialect=self.read, tables=tuple(Part.plain(s.key, s.name) for s in self.tables), @@ -669,6 +741,7 @@ class _Statement: 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]]: @@ -757,25 +830,59 @@ class _Statement: return TimePredicate(column=self.part(cols[0], sources), op=op, value=value, condition=self.part(cond, sources, excludes=True)) - def _group(self, select: Any, sources: dict[str, _Source]) -> list[Part]: - """分组键。GROUP BY 1 换成 SELECT 里对应的表达式,GROUP BY 别名换成别名指的表达式(MySQL、SQLite 的读法)。""" + @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 = select.selects - by_alias = {name_key(s.alias): s.this for s in selects if isinstance(s, exp.Alias) and s.alias} + 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 item in group.expressions: - 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): - target = 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 by_alias: - target = by_alias[name_key(target.name)] - out.append(self.part(self.canonical(target, sources), sources)) + 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]: @@ -811,6 +918,76 @@ class _Statement: 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 "子查询" + # ========================================================================== # 比较 @@ -820,7 +997,8 @@ class _Statement: @dataclass(frozen=True) class FacetDiff: """一处不同。aspect:time(时间列)/ time_range(同一时间列、范围不同)/ filter / aggregate / join / tables / - group_by / having / limit / parse。a、b 是两边各自的写法(给人看,没有就是 None),text 是一句中文说明。 + group_by / having / limit(取哪几行:行数、跳过的行数、取前几行时的排序)/ parse。a、b 是两边各自的写法(给人看, + 没有就是 None),text 是一句中文说明。 detail 只在 time 上有:两边的时间列 [{table, column}],调用方据此对照数据目录(业务日期)。""" aspect: str @@ -910,6 +1088,12 @@ def _rows(limit: str | None) -> str: 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: """一次比较:两边的叫法、列的中文名,以及按哪种写法比。 @@ -1099,14 +1283,55 @@ class _Compare: for p in self.b.having if self.k(p) not in ka] def limit(self) -> list[FacetDiff]: - a, b = self.a.limit, self.b.limit - return [FacetDiff("limit", a, b, _say(self.la, _rows(a), ",", self.lb, _rows(b)))] if a != b else [] + """取哪几行:最外层两边比;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]: + for x, y in zip_longest(ia, ib): + out += self.rows(x, y) + return out + + def lk(self, x: LimitFacet | None) -> tuple[Any, ...] | None: + if x is None: + return None + return x.count, x.offset, tuple((self.k(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): + 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]: + 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、 - LIMIT。完全一致返回空列表。 + 取哪几行。完全一致返回空列表。 labels 是两边的叫法(缺省 A / B,双盲复算可以传「原流程」「盲算」);names 是列的中文名({「表.列」或「列」: 中文名},可以用 catalog_labels 从数据目录取),给了的话时间列写成「支付时间 paid_at」。任一边解析不了就只返回 @@ -1124,5 +1349,5 @@ def diff_facets(a: Facets, b: Facets, labels: tuple[str, str] = ("A", "B"), *, 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", +__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/tests/test_sql_semantics.py b/backend/tests/test_sql_semantics.py index 690fd351..d12f8c46 100644 --- a/backend/tests/test_sql_semantics.py +++ b/backend/tests/test_sql_semantics.py @@ -3,7 +3,8 @@ - normalize_sql:字面量换成占位符后的规范文本 + 字面量。只有字面量不同的两条 SQL 规范文本相同(定时运行每期的日期 不同,不能算模板变了);关键字大小写、空白、注释也不算不同;解析不了的退回逐词规范化,照样认得出只改了字面量 - facets:表、连接(类型 / 关联进来的表 / ON 条件)、WHERE 合取式、时间条件(列 / 运算符 / 取值)、分组、聚合 - (函数 / 参数 / 去重)、HAVING、LIMIT。别名换成真实表名、比较的两边排好顺序;解析不了带 parse_error,不抛 + (函数 / 参数 / 去重)、HAVING、取哪几行(各段的 LIMIT / OFFSET / 排序)。别名换成真实表名(条件里套着的子查询、 + HAVING 里的别名也是)、比较的两边排好顺序;解析不了带 parse_error,不抛 - diff_facets:每处不同一句中文:「A 按支付时间 paid_at 筛时间,B 按下单时间 created_at」「A 排除了 is_deleted = 1 的行,B 没有」「A 对 order_id 去重计数,B 没去重」……写法不同、口径相同的不报 @@ -186,6 +187,27 @@ def test_having_and_limit(): 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"), @@ -211,8 +233,18 @@ 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] - assert facets("SELECT orders.id FROM orders WHERE NOT EXISTS " - "(SELECT 1 FROM refunds WHERE orders.id = refunds.order_id)").where[1].key == f.where[1].key + 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(): @@ -380,11 +412,47 @@ def test_an_unparsable_side_is_reported_once_without_item_by_item_comparison(): ("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("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] From 132d126c6f8e3054e5d5acde872bfbc710cc3342 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 01:46:55 +0800 Subject: [PATCH 8/9] =?UTF-8?q?=E6=A0=B8=E6=9F=A5=20agent=20=E5=A4=8D?= =?UTF-8?q?=E6=A0=B8=E4=BF=AE=E6=AD=A3=EF=BC=9A=E6=95=B0=E5=80=BC=E5=AD=97?= =?UTF-8?q?=E6=AE=B5=E7=9A=84=E5=90=84=E7=A7=8D=E5=86=99=E6=B3=95=E9=83=BD?= =?UTF-8?q?=E6=8D=A2=E6=88=90=E6=A0=BC=E5=AD=90=E5=BC=95=E7=94=A8=E3=80=81?= =?UTF-8?q?=E8=AE=A4=E4=B8=8D=E5=87=BA=E7=9A=84=E7=BB=93=E6=9E=84=E4=B8=8B?= =?UTF-8?q?=E7=9A=84=E6=95=B0=E4=B8=8D=E9=87=87=E4=BF=A1=EF=BC=9B=E9=80=80?= =?UTF-8?q?=E5=9B=9E=E6=99=AE=E9=80=9A=E8=B0=83=E7=94=A8=E6=97=B6=E6=8A=8A?= =?UTF-8?q?=E7=BB=93=E6=9E=84=E5=86=99=E8=BF=9B=E6=B6=88=E6=81=AF=EF=BC=8C?= =?UTF-8?q?=E6=B2=A1=E4=BA=A4=E7=9A=84=E5=AD=97=E6=AE=B5=E4=B8=8D=E5=86=8D?= =?UTF-8?q?=E7=AE=97=E3=80=8C=E6=9F=A5=E4=B8=8D=E5=88=B0=E3=80=8D=EF=BC=9B?= =?UTF-8?q?=E6=94=B6=E5=B0=BE=E7=9A=84=E9=92=B1=E6=8C=89=E6=9C=80=E8=B4=B5?= =?UTF-8?q?=E7=9A=84=E6=A0=B7=E5=AD=90=E7=95=99=E3=80=81=E7=95=99=E8=BF=87?= =?UTF-8?q?=E7=9A=84=E4=B8=8D=E5=86=8D=E8=A2=AB=E4=B8=8A=E9=99=90=E6=8B=A6?= =?UTF-8?q?=EF=BC=9B=E6=B0=B8=E8=BF=9C=E4=B8=8D=E6=8A=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 数只认格子:原来只认 type: number / integer,可为 null 的数(anyOf,Pydantic 的 Optional[float] 就这么写)、$ref 引用进来的、allOf 包着的数值字段都没换成格子引用,模型写的数原样收下。现在先展开本地 $ref、并掉 allOf,anyOf / oneOf 里有一支能放数就整个算数值字段;取值时按值的形状挑分支。认不出的结构(没写类型、文本字段、没写 items 的数组、没写属性的对象)下模型写的数一律丢掉、照记下来,码值里的数照收。嵌套的引用、数组写成一段 JSON 文本的照样认,还不是数组的照实记下格式不对。数值字段引用到日期、名称这类不是数的格子,记成 unresolved,不拿去和原流程的数比。 退回普通调用:原来消息里只有「按给定结构填写」,结构却不在请求里,模型猜不到字段名,交回来对不上的字段全成了 missing,双盲复算会说盲算什么也没查到。现在把结构(JSON Schema)写进消息;没交的字段记 unresolved(模型没有交这个字段),交了 null 才是 missing;一个要求的字段都没交就是没按结构交结果,照错误交回。模型建结构化输出就报错的,直接走普通调用,不白算一次调用。 上限与金额:一条查询自己的时限到了(ToolTimeout 也是 TimeoutError)原来被当成整个核查的时长到了,现在先认 ToolTimeout,原因交回模型、核查照常往下走。收尾要把查到的结果交给模型,原来留钱按眼下的消息估,最后一条查询的大段结果一进来就不够,查完了付不起收尾、结果白查;现在按收尾最贵的样子留(查询结果按 _digest 的上限、最后的说明截到 2000 字、结构也算),留过钱的收尾不再被上限拦下,实际用量超出估算或并行的核查先用掉当天余额时,最多超出一次收尾调用的钱。每日上限缺省按设置里的「核查每日金额上限」,明确传 None 才是不限,质询接口不必各自记得去读。 永远不抛:每日账写不进去、建工具出错这类意外原来会直接抛出,节点失败、质询接口回 500;现在整段包住,按已经攒下的用量、查询照实交回。 --- backend/app/engine/verify_agent.py | 543 +++++++++++++++++++++++------ backend/tests/test_verify_agent.py | 320 ++++++++++++++++- 2 files changed, 748 insertions(+), 115 deletions(-) diff --git a/backend/app/engine/verify_agent.py b/backend/app/engine/verify_agent.py index 4294f043..ad0a39d2 100644 --- a/backend/app/engine/verify_agent.py +++ b/backend/app/engine/verify_agent.py @@ -8,8 +8,10 @@ **数从快照取。** 每次成功的查询是一份查询快照(工具自己落的 put_json),按这个 agent 自己的顺序编 Q1、Q2…。 最后交结果时,调用方结构里的数值字段一律改成格子引用 {alias, row, column},由代码按引用从快照里取值:模型自己 -写的数字不采信,取不到就是 None,不兜底成 0。理由文字里的数字要写成 [[v:Qn.rX.列]],用本 agent 自己的台账 -(evidence.build_catalog)组装成文档(evidence.compose_doc),裸数字会被标出来。 +写的数字不采信,取不到就是 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: @@ -17,11 +19,14 @@ query_artifact 和 role: verify——随运行封存、证据面板能下钻。 台账条目由它按自己的种类记(dual_check / challenge / inquiry)。 **上限。** 查询次数、时长、金额(这一次、全局每日)任何一项到顶就停止查询,按已经查到的结果收尾:收尾那次 -调用总是要做的(时长到顶时再给 SETTLE_GRACE_S 的宽限),所以金额上限每一轮都给它留出钱。模型不在价格目录里 -时估不出金额,金额上限对它无效,照裁判的说法记一条说明。每次调用都记进核查的每日账(engine/spend.py)。 +调用总是要做的(时长到顶时再给 SETTLE_GRACE_S 的宽限),所以每一轮都按收尾最贵的样子给它留出钱——收尾时要把 +查到的结果交给模型,结果有多长事先不知道,按 _digest 的上限算。留过钱的收尾不再被上限拦下:某一轮的实际用量 +超出了估算、或者并行的核查先用掉了当天的余额时,宁可超出上限一次收尾调用的钱,也不让已经查到的结果白查。 +每日上限不传就按设置里的「核查每日金额上限」,明确传 None 才是不限。模型不在价格目录里时估不出金额,金额上限 +对它无效,照裁判的说法记一条说明。每次调用都记进核查的每日账(engine/spend.py)。 -永远不抛:模型没配好、调用失败、格式不对、触顶,都落在返回值里(error / limit)。模型调用都经 get_chat_model, -测试换成剧本。 +永远不抛:模型没配好、调用失败、格式不对、触顶、库写不进去,都落在返回值里(error / limit)。模型调用都经 +get_chat_model,测试换成剧本。 """ from __future__ import annotations @@ -32,6 +37,7 @@ 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 @@ -42,9 +48,9 @@ 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 +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, _split_structured, _typed, explain_model_error +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 @@ -61,6 +67,14 @@ _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" #: 触顶原因 @@ -96,14 +110,23 @@ _RATIONALE_FIELD: dict[str, Any] = { 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 / missing(模型说查不到)/ unresolved(reason 写原因,模型直接写了 - 数字的记 model_value)。queries:成功的查询 {alias, artifact, sql, row_count, truncated, source, tool, call_id}。 + 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:模型是不是自己停下来的(没被上限掐断)。 @@ -139,30 +162,132 @@ def _is_object(schema: Any) -> bool: return isinstance(schema, dict) and ("object" in _types(schema) or "properties" in schema) -def _is_number(schema: Any) -> bool: - return bool(_types(schema) & {"number", "integer"}) +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]: - """交给模型的结构:调用方结构里的数值字段(number / integer,嵌套的对象、数组里的也算)一律换成可为 null 的 - 格子引用,另加一个理由字段(已有同名的字符串字段就沿用)。其余字段原样。""" + """交给模型的结构:调用方结构里的数值字段(见 _numeric_types;嵌套的对象、数组、anyOf 的分支里的也算,引用先 + 展开)一律换成可为 null 的格子引用,另加一个理由字段(已有同名的字符串字段就沿用)。其余字段原样。""" def rewrite(node: Any) -> Any: if not isinstance(node, dict): return node - if _is_number(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 - if _is_object(node): - props = node.get("properties") or {} - return {**node, "type": "object", "properties": {k: rewrite(v) for k, v in props.items()}} - if "array" in _types(node) and isinstance(node.get("items"), dict): - return {**node, "items": rewrite(node["items"])} - return node + 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 - root = rewrite(schema if _is_object(schema) else {"type": "object", "properties": {}}) + 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])) @@ -170,6 +295,42 @@ def result_schema(schema: dict[str, Any]) -> dict[str, Any]: 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: """按格子引用从查询快照里取值:快照经工件库读(复验哈希),同一份只读一次。""" @@ -187,15 +348,22 @@ class _Cells: self.memo[artifact] = (None, "查询快照与哈希不一致,可能已被修改") return self.memo[artifact] - def take(self, value: Any, path: str, schema: dict[str, Any]) -> Any: - """一个数值字段:按引用取值,记下出处,返回值(取不到为 None)。""" + 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": "模型直接写了数值,没有给出格子引用,不采信"}) + "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") @@ -220,25 +388,99 @@ class _Cells: self.cells.append({"path": path, "status": "unresolved", "alias": alias, "artifact": query["artifact"], "reason": f"{alias} {e}"}) return None - truth = _typed(raw, schema, column_kind(content, name)) - self.cells.append({"path": path, "status": "resolved", "ref": f"{alias}.r{row}.{name}", "alias": alias, - "row": row, "column": name, "artifact": query["artifact"], - "eid": cell_eid(query["artifact"], row, name), "value": truth}) + 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: - """按调用方的结构走一遍模型交的结果:数值字段换成从快照取到的值,其余原样。""" - if _is_number(schema): - return self.take(value, path, schema) - if _is_object(schema): - given = value if isinstance(value, dict) else {} - return {k: self.walk(sub, given.get(k), f"{path}.{k}" if path else k) - for k, sub in (schema.get("properties") or {}).items()} - if "array" in _types(schema) and isinstance(schema.get("items"), dict): - if not isinstance(value, list): + """按调用方的结构(_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 - return [self.walk(schema["items"], v, f"{path}[{i}]") for i, v in enumerate(value)] - return value + 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 # -------------------------------------------------------------------------- @@ -256,16 +498,25 @@ class _Budget: self.daily_max_usd = daily_max_usd self.spent = 0.0 - def estimate(self, text: str, out: int) -> float: - return pricing.estimate_cost(self.model_id, _tokens(text), out) if self.priced else 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) -> tuple[str | None, str | None]: - """(记在哪天, None) 或 (None, 拦下的上限)。keep 是要额外留出的钱(收尾那次调用):只用来判断,不预留。""" - 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" - cap = self.daily_max_usd if self.priced else None - if cap is not None and keep and (await spend.daily_spend())["usd"] + estimate + keep > cap + 1e-12: - return None, "daily_max_usd" + 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") @@ -326,16 +577,28 @@ def _text_of(messages: list[BaseMessage]) -> str: def _final_messages(system: str, task: str, queries: list[dict[str, Any]], said: str, limit_words: str | None ) -> list[BaseMessage]: """收尾那次调用的消息:任务、查到的结果(每次查询列出行号)、模型最后说的话,以及交结果的规则。 - limit_words:触顶时是哪个上限(「最多查询 6 次」),告诉模型为什么不能再查。""" + 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() or '(没有)'}") + 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): @@ -343,6 +606,30 @@ def _parsed(value: Any) -> dict[str, Any] | None: 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, @@ -353,7 +640,7 @@ async def run_verify_agent( max_queries: int | None = 6, timeout_s: float | None = 120, max_cost_usd: float | None = None, - daily_max_usd: float | None = None, + daily_max_usd: float | None | Default = Default.SETTINGS, max_steps: int | None = None, emit: Emit | None = None, run_id: str = "", @@ -364,37 +651,41 @@ async def run_verify_agent( - spec:用哪个模型(一般是 api/settings.py 的 verify_model_spec)。system / task:交给模型的角色说明和任务。 - sources:可以查的数据源名;不存在、停用了的照实记进 notes。 - - output_schema:调用方要的结果结构(JSON Schema 对象)。数值字段会换成格子引用交给模型、由代码取值。 + - output_schema:调用方要的结果结构(JSON Schema 对象,可以带 $defs / $ref)。数值字段会换成格子引用交给模型、 + 由代码取值。 - 上限:max_queries(db_query 的执行次数,被拒的语句也算)、timeout_s(整个循环的时长,收尾另给宽限)、 - max_cost_usd(这一次)、daily_max_usd(全局每日,和裁判分开记账);None 表示不限。max_steps 是模型最多想 - 几轮,缺省按查询次数推。 + 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 时上传源一律报错,不拿当前版本顶上)。 """ - started = time.monotonic() - deadline = started + timeout_s if timeout_s is not None else None - usage: dict[str, Any] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cost_usd": 0.0, "calls": 0, - "unpriced_calls": 0} - notes: list[str] = [] - queries: list[dict[str, Any]] = [] - ledger: list[dict[str, Any]] = [] + 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() - def outcome(**fields: Any) -> VerifyOutcome: - 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 queries] - base: dict[str, Any] = {"result": None, "cells": [], "queries": shown, "rationale": None, "usage": usage, - "notes": notes, "duration_ms": int((time.monotonic() - started) * 1000)} - return VerifyOutcome(**{**base, **fields}) - # ---- 模型 try: async with SessionLocal() as session: @@ -402,6 +693,12 @@ async def run_verify_agent( 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 数估算金额:金额上限(这次核查、每日)对其无效," @@ -423,6 +720,14 @@ async def run_verify_agent( 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]: @@ -457,6 +762,7 @@ async def run_verify_agent( # ---- 查库的循环 stopped: str | None = None answered = False + kept = False # 有一轮按收尾最贵的样子留过钱:收尾不再按金额上限拦 said = "" attempts = 0 done: dict[str, str] = {} # 成功过的查询(工具名 + 参数)→ 编号:一模一样的不再执行 @@ -466,17 +772,17 @@ async def run_verify_agent( stopped = "timeout_s" break estimate = budget.estimate(_text_of(messages) + tool_text, _OUT_STEP) - day, refused = await budget.reserve(estimate, keep=budget.estimate(_text_of(messages), _OUT_FINAL)) + 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(model=model_id, error=explain_model_error(failure, model_id) if failure - else "核查模型没有给出回复") + 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: @@ -494,10 +800,10 @@ async def run_verify_agent( messages.append(ToolMessage(content=f"错误:没有名为 {name} 的工具,可用的工具:" + ("、".join(tool_map) or "(无)"), tool_call_id=call_id)) continue - schema = args_model_of(tool) - if schema is not None: + args_schema = args_model_of(tool) + if args_schema is not None: try: - args, _ = prepare_args(schema, args) + 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)) @@ -535,10 +841,12 @@ async def run_verify_agent( else: stopped = "max_steps" - words = None - if stopped: - words = _limit_words(stopped, max_queries=max_queries, timeout_s=timeout_s, max_cost_usd=max_cost_usd, - daily_max_usd=daily_max_usd, steps=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}),已停止查询,按已有结果收尾") @@ -547,16 +855,12 @@ async def run_verify_agent( 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(model=model_id, limit=stopped, answered=answered, error=f"核查模型不可用({describe_exception(e)})") - schema = result_schema(output_schema) + return outcome(limit=stopped, answered=answered, error=f"核查模型不可用({describe_exception(e)})") final = _final_messages(head, task, queries, said, words) - estimate = budget.estimate(_text_of(final), _OUT_FINAL) - day, refused = await budget.reserve(estimate) - if refused: - limit_words = _limit_words(refused, max_queries=max_queries, timeout_s=timeout_s, max_cost_usd=max_cost_usd, - daily_max_usd=daily_max_usd, steps=steps) - return outcome(model=model_id, limit=stopped or refused, answered=answered, - error=f"已达{limit_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 @@ -564,38 +868,55 @@ async def run_verify_agent( def until_finish() -> float | None: return None if finish is None else finish - time.monotonic() - got, raw, failure = await ask(final, settler.with_structured_output(schema, include_raw=True).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 + 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 抠出来 + # 不支持结构化输出的模型:普通调用一次,结构写进消息,从回复里把 JSON 抠出来。结构化那次已经做过的, + # 这一次是补救,照常按上限拦;没做过的,这一次就是收尾,循环里留过钱 + plain = _plain_messages(final, schema) remaining = until_finish() if remaining is None or remaining > 0: - estimate = budget.estimate(_text_of(final), _OUT_FINAL) - day, refused = await budget.reserve(estimate) + 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, raw, failure = await ask(final, settler.ainvoke, estimate, day or "", "verify_result", remaining) + 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(model=model_id, limit=stopped, answered=answered, error=why) + 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(output_schema, parsed) if _is_object(output_schema) else {} + 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(model=model_id, result=result, cells=cells.cells, rationale=doc, limit=stopped, answered=answered) + 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, @@ -612,15 +933,17 @@ async def _run_tool(tool: Any, name: str, args: dict[str, Any], call_id: str, re 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 ToolTimeout as e: - 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 Exception as e: # noqa: BLE001 - 工具失败交回模型,让它自己改 message = f"工具执行失败:{describe_exception(e)}" await _announce(emit, EventType.TOOL_ERROR, tool=name, call_id=call_id, role="verify", @@ -656,4 +979,4 @@ async def _run_tool(tool: Any, name: str, args: dict[str, Any], call_id: str, re return content -__all__ = ["LIMITS", "RATIONALE", "SETTLE_GRACE_S", "VerifyOutcome", "result_schema", "run_verify_agent"] +__all__ = ["LIMITS", "RATIONALE", "SETTLE_GRACE_S", "Default", "VerifyOutcome", "result_schema", "run_verify_agent"] diff --git a/backend/tests/test_verify_agent.py b/backend/tests/test_verify_agent.py index 7e72083f..0a8c04d0 100644 --- a/backend/tests/test_verify_agent.py +++ b/backend/tests/test_verify_agent.py @@ -4,9 +4,11 @@ - 每次成功的查询是一份查询快照,按这个 agent 自己的顺序编 Q1、Q2…;节点里调用时发 tool.end(带 query_artifact、 role: verify),但不往证据台账里写 query:核查的查询进了报告的写作目录,写作者就会把它们当成可以引用的证据 - 最后按调用方的结构交结果:数值字段一律是格子引用 {alias, row, column},由代码从快照取值,模型自己写的数字 - 不采信;理由里的数字要写成 [[v:Qn.rX.列]],用本 agent 自己的台账组装成文档,裸数字会被标出 -- 查询次数、时长、金额(单次和每日)任何一项到顶就停止查询、按已有结果收尾;模型不在价格目录里时金额上限 - 无效,照裁判的说法记一条说明 + 不采信(type 列表、anyOf、$ref、allOf 里的数也是;认不出的结构下写的数一律丢掉);没交的字段和交了 null 分开记; + 理由里的数字要写成 [[v:Qn.rX.列]],用本 agent 自己的台账组装成文档,裸数字会被标出 +- 查询次数、时长、金额(单次和每日)任何一项到顶就停止查询、按已有结果收尾;收尾的钱在查询之前就按最贵的样子 + 留好;每日上限缺省按设置;模型不在价格目录里时金额上限无效,照裁判的说法记一条说明 +- 永远不抛:库写不进去、模型不支持结构化输出,都落在返回值里 - 上传表格按调用方给的数据版本查(运行固定的那一版),不拿当前版本顶上 模型一律是 MockChatModel 剧本,数据源是现造的 SQLite 文件,表名、数字都是虚构的。 @@ -31,7 +33,8 @@ 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.verify_agent import run_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 @@ -48,9 +51,12 @@ 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 == spend.SPEND_KEY)) + 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 @@ -272,6 +278,32 @@ async def test_a_model_call_that_runs_past_the_deadline_is_cut(monkeypatch, shop 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 美元结算:金额上限的算术好核对。""" @@ -279,6 +311,88 @@ def priced(monkeypatch): 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)] @@ -394,6 +508,131 @@ async def test_nested_numbers_and_arrays_are_references_too(monkeypatch, shop): 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 + + # -------------------------------------------------------------------------- # 数据版本 # -------------------------------------------------------------------------- @@ -482,3 +721,74 @@ async def test_a_final_answer_in_plain_text_json_is_still_read(monkeypatch, shop 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 From d5e9ef39aebff35b57110d57f817c72822fb296e Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 5 Oct 2026 02:37:33 +0800 Subject: [PATCH 9/9] =?UTF-8?q?SQL=20=E8=AF=AD=E4=B9=89=E5=A4=8D=E6=A0=B8?= =?UTF-8?q?=E4=BF=AE=E6=AD=A3=EF=BC=88=E7=AC=AC=202=20=E8=BD=AE=EF=BC=89?= =?UTF-8?q?=EF=BC=9ACTE=E3=80=81=E5=AD=90=E6=9F=A5=E8=AF=A2=E8=B5=B7?= =?UTF-8?q?=E7=9A=84=E5=90=8D=E5=AD=97=E4=B8=8D=E5=86=8D=E8=BF=9B=E6=AF=94?= =?UTF-8?q?=E8=BE=83=EF=BC=8C=E6=9D=A1=E4=BB=B6=E9=87=8C=E7=9A=84=E5=AD=90?= =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E9=80=89=E5=87=BA=E6=9D=A5=E7=9A=84=E5=88=97?= =?UTF-8?q?=E5=90=8D=E4=B9=9F=E4=B8=8D=E7=AE=97=EF=BC=9B=E6=B4=BE=E7=94=9F?= =?UTF-8?q?=E6=9D=A5=E6=BA=90=E6=8C=89=E4=BD=8D=E7=BD=AE=E7=BC=96=E5=8F=B7?= =?UTF-8?q?=EF=BC=8C=E7=BC=96=E5=8F=B7=E9=94=99=E5=BC=80=E6=97=B6=E5=86=8D?= =?UTF-8?q?=E6=8C=89=E4=B8=8D=E5=86=99=E7=BC=96=E5=8F=B7=E7=9A=84=E5=86=99?= =?UTF-8?q?=E6=B3=95=E9=85=8D=E4=B8=80=E9=81=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 上一轮只去掉了条件里子查询的表别名,子查询整段写进外层条件的 key 时还带着各自起的名字:标量子查询、IN 子查询给选出来的列起的名字(AVG(amount) AS avg_amount),子查询里 FROM 的子查询别名(s 换成 agg),以及外层引用 CTE、FROM 里的子查询算出来的列时的来源名(派生来源的 key 就是别名)。CTE 改个名,外层的 total > 100 就被报成两边各有一个、互相矛盾的条件;双盲复算里盲算的名字本来就是自己起的,会被说成换了口径。 现在派生来源(引用 CTE 的表、FROM / JOIN 里的子查询和 LATERAL 子查询)按位置编号作比较用的 key:从外往里一层一层地编,CTE 在第一次用到的地方编号,写成 CTE 还是 FROM 里的子查询、CTE 定义的先后都不影响;给人看的写法照写名字。不拿里面那段查询的写法作 key:里面差一个条件,外层按它的列写的条件就全跟着变,同一处不同会报好几遍。条件里的子查询写进外层 key 时,引用 CTE 的表、子查询(含 LATERAL)、子查询自带的 CTE 定义都换成编号,只看取值的子查询(IN、EXISTS、标量子查询)去掉选出来的列名;FROM 里的子查询、CTE 定义选出来的列名留着,外层按名字引用它们。LATERAL 的别名原来没登记成来源,按「别名写错了」照写,一并改成派生来源。 编号只看位置:一边多出一个排在前面的 CTE(比如盲算先关联一个筛 VIP 的 CTE),后面的编号就错开了,同名的 CTE 反倒比原来按名字比时多报出两条自相矛盾的话。Part 加 anon(派生来源一律写「@」、不写编号),diff_facets 每个方面先按编号比,比不上的再按 anon 一对一地配,配上的不算不同;比较两边的先后先按 anon 排、分不出再按编号;连接先配完全一样的再按关联的表配,免得错开的两个 CTE 被配成一对、说成关联条件不同。编号对得上时照样分得清是哪个 CTE 的同名列(p.total > r.total 写反成 r.total > p.total 照报)。筛选条件配对改按列的比较写法(Part.column_keys,CTE、子查询的列写「@.列」),CTE 名不同、阈值不同时说成「按 total > 100 筛选,盲算按 total > 200 筛选」,不再拆成两条各说各的。 仍按名字比的只有 CTE、子查询选出来、外层按名字引用的列(一边叫 total、一边叫 sum_amt),写进了模块说明。 --- backend/app/data/sql_semantics.py | 357 +++++++++++++++++++++------- backend/tests/test_sql_semantics.py | 231 +++++++++++++++++- 2 files changed, 503 insertions(+), 85 deletions(-) diff --git a/backend/app/data/sql_semantics.py b/backend/app/data/sql_semantics.py index 0d7a2ad8..613f5e41 100644 --- a/backend/app/data/sql_semantics.py +++ b/backend/app/data/sql_semantics.py @@ -17,6 +17,13 @@ 写在 WHERE 里的列等式算作关联条件,不算筛选。条件里套着的子查询(EXISTS、IN (SELECT …))既单算一段,也按它自己 那一段的来源整个写进外层那个条件:别名在子查询里定义,拿外层的来源去认,同一个条件只因别名不同就会被报成不同。 +**各自起的名字不进比较。** 表的别名换成真实表名;CTE、FROM 里的子查询没有真实表名可换,比较时按它在语句里的位置 +编号(@1、@2…,从外往里、按第一次用到的先后),给人看时照写名字;条件里的子查询(IN、EXISTS、标量子查询)选出来 +的列名不影响取值,比较时去掉。这些名字都是各自起的,双盲复算的盲算更是:按名字比,同一个口径会被说成换了口径。 +编号只看位置,一边多出一个排在前面的 CTE、子查询,后面的编号就错开了:diff_facets 先按编号比,比不上的再按不写 +编号的写法一对一地配,配上的不算不同。仍按名字比的只有 CTE、子查询选出来、外层按名字引用的列(一边叫 total、 +一边叫 sum_amt,算两列)。 + **和 sqlcheck 一样宁可漏报。** 解析不了(语法错、方言特性、多条语句、不是查询)就带 parse_error 返回,从不抛 异常。没加表前缀、又有好几张表可能有这一列的,照写列名(给了表结构才去对)。只读语法树,不执行 SQL。 """ @@ -24,6 +31,7 @@ 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 @@ -42,6 +50,11 @@ 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 = "@" # ========================================================================== @@ -201,11 +214,14 @@ ColumnRef = tuple[str | None, str] class Part: """一个条件、一个分组键、一个聚合参数…… - key 用来比较:别名换成真实表名,表名、列名按 name_key,等式和比较的两边排好顺序;loose 同 key 但不写表名,两边 - 查的表换了一张(orders 换成 orders_archive)时拿它比其余各方面,免得同一个条件只因表名不同被报成不同。 - text 给人看:好几张表都有的列写「表.列」(真实表名,不是别名),其余只写列名;full 一律写表名,两边的 text - 一样、key 却不同时改用它说明。excludes 只在 WHERE 的「排除」写法上有:<>、NOT、NOT IN、IS NOT NULL 去掉 - 否定之后排除的是什么(IS NOT NULL 写成「列 为空」)。 + 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 @@ -214,11 +230,13 @@ class Part: excludes: str | None = None loose: str = "" full: str = "" + column_keys: tuple[str, ...] = () + anon: str = "" @staticmethod def plain(key: str, text: str) -> "Part": - """没有列的一项(表名、COUNT(*) 的「*」):四种写法相同。""" - return Part(key=key, text=text, loose=key, full=text) + """没有列的一项(表名、COUNT(*) 的「*」):比较用的写法都是 key,给人看的都是 text。""" + return Part(key=key, text=text, loose=key, full=text, anon=key) @dataclass(frozen=True) @@ -358,6 +376,53 @@ def _own_columns(node: Any) -> list[Any]: 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 @@ -423,11 +488,12 @@ class _Sources(dict): @dataclass class _Source: - """一段 SELECT 的一个来源。kind:table(实表)/ derived(CTE、子查询)/ other(表函数、认不出的别名)。""" + """一段 SELECT 的一个来源。kind:table(实表)/ derived(CTE、子查询、LATERAL 子查询)/ other(表函数、认不出的 + 别名)。""" kind: str - name: str # 给人看:表名(有库名就带上)、CTE 名、子查询别名 - key: str # 比较用:name_key + name: str # 给人看:表名(有库名就带上)、CTE 名、子查询(含 LATERAL)的别名 + key: str # 比较用:实表是 name_key,派生来源是编号(见 _Statement._number_derived) idents: tuple[Any, ...] = () # 表名各段的 Identifier:给人看的写法照原样写回去 query: Any = None # derived:里面那段查询 @@ -446,9 +512,13 @@ class _Statement: 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): @@ -470,6 +540,61 @@ class _Statement: 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 @@ -485,21 +610,19 @@ class _Statement: if id(select) in self._memo: return self._memo[id(select)] - items = ([select.args["from_"].this] if select.args.get("from_") is not None else []) + \ - [j.this for j in select.args.get("joins") or []] out = _Sources() - for i, node in enumerate(items): - alias = node.alias_or_name if isinstance(node, (exp.Table, exp.Subquery)) else "" + 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, name_key(node.name), (node.this.copy(),), + 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 isinstance(node, exp.Subquery): - out[slot] = _Source("derived", alias or "子查询", slot, (exp.to_identifier(alias),) if alias else (), - node.this) + 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 @@ -598,11 +721,13 @@ class _Statement: 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、loose 比较用,名字按 name_key,key 写真实表名、loose 不写;text、full 给人看, - text 只在好几张表都有这一列时写真实表名,full 一律写(关联条件说的就是两张表各出哪一列,用 full)。 + """五种写法(见 Part):key、anon、loose 比较用,名字按 name_key,key、anon 写真实表名、loose 不写(CTE、子查询 + 只有 key 写编号);text、full 给人看,text 只在好几张表都有这一列时写真实表名,full 一律写(关联条件说的就是 + 两张表各出哪一列,用 full)。 node 里套着的子查询逐段改写:每一列按它所在那一段的来源认,子查询里的表去掉别名(列已经写成了真实表名)。 - 子查询里拿两张表的列互相比的(相关子查询),给人看时和关联条件一样写表名。 + 子查询里拿两张表的列互相比的(相关子查询),给人看时和关联条件一样写表名。比较用的写法里各自起的名字也去掉 + (见 _drop_names)。 """ from sqlglot import exp @@ -623,21 +748,50 @@ class _Statement: 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"): + 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"): + 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(source.key) if source is not None and mode == "key" else None) + col.set("table", exp.to_identifier(owner) if owner else None) col.set("db", None) col.set("catalog", None) return @@ -654,7 +808,7 @@ class _Statement: col.set("catalog", heads[-2].copy() if len(heads) > 1 else None) def canonical(self, node: Any, sources: dict[str, _Source]) -> Any: - """比较的两边排好顺序(列在左、字面量在右;两边都有列的按 key 排),IN 的取值排好序。返回新的树。 + """比较的两边排好顺序(列在左、字面量在右;两边都有列的按 _rank 排),IN 的取值排好序。返回新的树。 套着的子查询里的比较按它那一段的来源排:拿外层的来源去认子查询的别名,排出来的先后就取决于别名怎么起。""" from sqlglot import exp @@ -664,7 +818,7 @@ class _Statement: 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.render(left, here, "key") > self.render(right, here, "key")): + 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): @@ -673,17 +827,24 @@ class _Statement: 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], None] = {} + refs: dict[tuple[str | None, str], str] = {} for col in _own_columns(node): source, name = self.resolve(col, sources) - refs.setdefault((source.name if source is not None else None, name)) + 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) + 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 为空」。 @@ -748,12 +909,10 @@ class _Statement: from sqlglot import exp first = select.args.get("from_") - first_alias = first.this.alias_or_name if first is not None and \ - isinstance(first.this, (exp.Table, exp.Subquery)) else "" + 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 []): - target = join.this - alias = target.alias_or_name if isinstance(target, (exp.Table, exp.Subquery)) else "" + 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: @@ -790,12 +949,14 @@ class _Statement: def _join(self, draft: dict[str, Any], sources: dict[str, _Source]) -> JoinFacet: source = draft["source"] - table = Part.plain(source.key, source.name) if source is not None else Part.plain("?", "?") - triples = sorted((self.render(c, sources, "key"), self.render(c, sources, "loose"), - self.render(c, sources, "full")) for c in draft["conds"]) - on = Part(key=" AND ".join(t[0] for t in triples), text=" AND ".join(t[2] for t in triples), - loose=" AND ".join(sorted(t[1] for t in triples)), full=" AND ".join(t[2] for t in triples)) \ - if triples else None + 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"] @@ -1071,11 +1232,11 @@ def _labeler(names: Mapping[str, str] | None) -> Any: def _single_column(part: Part, *, loose: bool) -> str | None: - """只涉及一列的条件:那一列(表.列;loose 时只看列名),用来把两边对同一列的不同写法配成一对。""" - if len(part.columns) != 1: + """只涉及一列的条件:那一列(Part.column_keys:实表.列,CTE、子查询的列写「@.列」,两边起的名字不同、编号错开 + 也配得上;loose 时只看列名),用来把两边对同一列的不同写法配成一对。""" + if len(part.columns) != 1 or len(part.column_keys) != 1: return None - table, column = part.columns[0] - return name_key(column) if loose else f"{name_key(table or '')}.{name_key(column)}" + return name_key(part.columns[0][1]) if loose else part.column_keys[0] def _verb(g: Aggregate) -> str: @@ -1099,6 +1260,8 @@ 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: @@ -1115,6 +1278,32 @@ class _Compare: 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]: """两边给人看的写法;读起来一样(各自只写列名)却不是同一个东西时,改成一律写表名的写法。""" @@ -1143,7 +1332,7 @@ class _Compare: 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 {self.k(p) for p in ca} != {self.k(p) for p in cb}: + 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) @@ -1157,7 +1346,7 @@ class _Compare: 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 {self.k(t.condition) for t in a.time} != {self.k(t.condition) for t in b.time}: + 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))] @@ -1167,8 +1356,7 @@ class _Compare: 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 = [p for p in fa if self.k(p) not in {self.k(q) for q in fb}] - only_b = [p for p in fb if self.k(p) not in {self.k(q) for q in fa}] + 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) @@ -1190,6 +1378,9 @@ class _Compare: 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): @@ -1200,8 +1391,7 @@ class _Compare: ys.remove(y) return pairs - ga = [g for g in self.a.aggregates if key(g) not in {key(h) for h in self.b.aggregates}] - gb = [g for g in self.b.aggregates if key(g) not in {key(h) for h in self.a.aggregates}] + 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) @@ -1229,30 +1419,31 @@ class _Compare: return out def joins(self) -> list[FacetDiff]: + """连接:先把完全一样的(类型、关联的表、条件)两边配掉,剩下的按关联的是哪张表配成一对、说类型或条件哪里 + 不同,再剩下的是只有一边关联了。直接按表配的话,编号错开时会把两个不相干的 CTE 配成一对,说成「关联条件 + 不同」;先配完全一样的,同一张表关联了两次、先后不同也不算不同。""" la, lb = self.la, self.lb - - def index(f: Facets) -> dict[str, list[JoinFacet]]: - out: dict[str, list[JoinFacet]] = {} - for j in f.joins: - out.setdefault(j.table.key, []).append(j) - return out - - ja, jb = index(self.a), index(self.b) + 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 key in dict.fromkeys([*ja, *jb]): - xs, ys = ja.get(key, []), jb.get(key, []) - for x, y in zip(xs, ys): - 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): - 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 xs[len(ys):]] - out += [FacetDiff("join", None, y.table.text, _say(lb, "关联了", y.table.text, "表,", la, "没有")) - for y in ys[len(xs):]] + 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] @@ -1264,7 +1455,7 @@ class _Compare: def group(self) -> list[FacetDiff]: a, b = self.a.group_by, self.b.group_by - if {self.k(p) for p in a} == {self.k(p) for p in b}: + 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: @@ -1276,11 +1467,9 @@ class _Compare: return [FacetDiff("group_by", wa or None, wb or None, text)] def having(self) -> list[FacetDiff]: - ka, kb = {self.k(p) for p in self.a.having}, {self.k(p) for p in self.b.having} - return [FacetDiff("having", p.text, None, self.one_sided(self.la, self.lb, p, group=True)) - for p in self.a.having if self.k(p) not in kb] + \ - [FacetDiff("having", None, p.text, self.one_sided(self.lb, self.la, p, group=True)) - for p in self.b.having if self.k(p) not in ka] + 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、子查询里的按出现的先后一一配对(两边给它起的名字可以不同),整组一样就不报。""" @@ -1290,15 +1479,18 @@ class _Compare: 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]: + 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) -> tuple[Any, ...] | None: + def lk(self, x: LimitFacet | None, *, anon: bool = False) -> tuple[Any, ...] | None: + """一段取哪几行的比较用写法:行数、跳过的行数、排序键。anon 时排序键按 ak(见 unmatched)。""" if x is None: return None - return x.count, x.offset, tuple((self.k(p), desc) for p, desc in x.order) + 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: @@ -1308,7 +1500,7 @@ class _Compare: def rows(self, x: LimitFacet | None, y: LimitFacet | None) -> list[FacetDiff]: """一对段(一边可以没有)取的行有什么不同:行数、跳过的行数,两边都限了行时再比排序。""" - if self.lk(x) == self.lk(y): + 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 @@ -1320,7 +1512,8 @@ class _Compare: 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]: + 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) diff --git a/backend/tests/test_sql_semantics.py b/backend/tests/test_sql_semantics.py index d12f8c46..1e62c3ab 100644 --- a/backend/tests/test_sql_semantics.py +++ b/backend/tests/test_sql_semantics.py @@ -4,11 +4,13 @@ 不同,不能算模板变了);关键字大小写、空白、注释也不算不同;解析不了的退回逐词规范化,照样认得出只改了字面量 - facets:表、连接(类型 / 关联进来的表 / ON 条件)、WHERE 合取式、时间条件(列 / 运算符 / 取值)、分组、聚合 (函数 / 参数 / 去重)、HAVING、取哪几行(各段的 LIMIT / OFFSET / 排序)。别名换成真实表名(条件里套着的子查询、 - HAVING 里的别名也是)、比较的两边排好顺序;解析不了带 parse_error,不抛 + HAVING 里的别名也是),CTE、子查询按位置编号、不按起的名字比,条件里的子查询选出来的列名不算;比较的两边排好 + 顺序;解析不了带 parse_error,不抛 - diff_facets:每处不同一句中文:「A 按支付时间 paid_at 筛时间,B 按下单时间 created_at」「A 排除了 - is_deleted = 1 的行,B 没有」「A 对 order_id 去重计数,B 没去重」……写法不同、口径相同的不报 + is_deleted = 1 的行,B 没有」「A 对 order_id 去重计数,B 没去重」……写法不同、口径相同的不报,一边多出一个排在 + 前面的 CTE、编号错开了也不报 -夹具全是虚构的表:订单 orders、退款 refunds、门店 stores,以及一张中文表「订单」。 +夹具全是虚构的表:订单 orders、退款 refunds、门店 stores、用户 users、付款 payments,以及一张中文表「订单」。 """ from __future__ import annotations @@ -435,6 +437,229 @@ def test_a_real_difference_inside_a_subquery_is_still_reported(): "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",