"""数据目录模型起草的引擎（data/catalog_draft.py）：分段、按失败原因处理、捡回被截断的部分、起草要求、用量。 用假模型复刻实测查明的几种失败（不调用真实模型）： - 思考加输出用完额度：finish_reason=length 且正文为空（只有思考）、LengthFinishReasonError → 拆成两半再问； - 只给空壳 → 拆开再问；拆不开的换纯文本问一次，写出来了这次起草后面都走纯文本； - 被截断的 JSON：已经完整的表、列对象捡回来，只补剩下的； - 只回一部分表：只补缺的，补一次； - 服务方嫌额度太大 → 退回 8192，后面都用它；结构化输出不支持 → 后面改走纯文本； - 网络错误、5xx → 退避重试，最多 2 次，报错写明重试过几次；超时当作这一段太大，拆小再问，不原样重试。 验收后修正的几条（用假模型复现）：宽表的列分段均分、没有很小的尾段；拆不开的空壳纯文本也再问一次；额度用完或超时 累计 2 次后新开始的段列数减半；一段最多调用 8 次；最近 6 次调用 4 次没写出来时提醒、少于 2 次时清掉；记超时和没写 出来的调用；模型留空的列计数；窄表不被宽表的列分段插队；原因写清是哪几段、哪几列。 表名、列名、说明文字都是虚构的。 """ from __future__ import annotations import asyncio import inspect import json from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Callable import pytest from langchain_core.messages import AIMessage from app.data import catalog from app.data import catalog_draft as D from tests.fixtures.catalog.fake_model import answer, requested # -------------------------------------------------------------------------- # 夹具 # -------------------------------------------------------------------------- def _col(name: str, type_: str = "INTEGER") -> dict: return {"name": name, "type": type_, "nullable": True, "comment": None} def _table(n: int, prefix: str = "c", comment: str | None = None) -> dict: return {"qualified": "", "schema": None, "columns": [_col(f"{prefix}{i:03d}") for i in range(1, n + 1)], "primary_key": [f"{prefix}001"], "comment": comment, "is_view": False, "foreign_keys": []} def _source(tables: dict, **kw) -> SimpleNamespace: for name, meta in tables.items(): meta["qualified"] = name return SimpleNamespace(id="src", name="shop", kind="sqlite", origin="manual", description=kw.pop("description", ""), schema_cache={"tables": tables}) def _reply(table: str, meta: dict) -> dict: """一张表的完整草稿：表级三项 + 每列中文名和度量类型。""" return {"label": f"{table}的中文名", "grain": "每条记录一行", "kind": "fact", "columns": [{"name": c["name"], "label": f"{c['name']}的中文名", "measure": "attribute"} for c in meta["columns"]]} @dataclass class Reply: """一次回复。obj 是正常的 JSON 对象；text 给了就原样作正文（截断的 JSON 之类）。""" obj: Any = None text: str | None = None finish: str = "stop" thinking: str = "" usage: tuple[int, int] = (0, 0) def render(self, mode: str, include_raw: bool) -> Any: if self.text is not None: content = self.text else: content = json.dumps(self.obj, ensure_ascii=False) if self.obj is not None else "" msg = AIMessage( content=content, additional_kwargs={"reasoning_content": self.thinking} if self.thinking else {}, response_metadata={"finish_reason": self.finish}, usage_metadata={"input_tokens": self.usage[0], "output_tokens": self.usage[1], "total_tokens": sum(self.usage)}, ) if mode == "text": return msg parsed = self.obj if self.text is None else None if not include_raw: return parsed return {"raw": msg, "parsed": parsed, "parsing_error": None if parsed is not None else ValueError("坏的")} @dataclass class Call: mode: str text: str asked: dict n: int class Scripted: """每次调用交给 handler(call) 决定怎么回：返回 Reply，或者抛异常；handler 可以是协程。""" def __init__(self, handler: Callable[[Call], Any], name: str = "scripted-model") -> None: self.handler, self.model_name = handler, name self.calls: list[Call] = [] def with_structured_output(self, schema: Any, include_raw: bool = False, **_: Any) -> Any: owner = self class _S: async def ainvoke(self, messages): return await owner._run(messages, "structured", include_raw) return _S() async def ainvoke(self, messages: list[Any]) -> Any: return await self._run(messages, "text", False) async def _run(self, messages: list[Any], mode: str, include_raw: bool) -> Any: text = "\n".join(str(getattr(m, "content", m)) for m in messages) call = Call(mode=mode, text=text, asked=requested(text), n=len(self.calls)) self.calls.append(call) out = self.handler(call) if inspect.isawaitable(out): out = await out return out.render(mode, include_raw) def answering(tables: dict, **kw) -> Callable[[Call], Reply]: """照请求写：每张表按请它写的那部分给出 _reply 的内容。""" def handle(call: Call) -> Reply: return Reply(obj={"tables": [answer(_reply(t, tables[t]), t, *call.asked[t]) for t in call.asked]}, **kw) return handle class _Status(Exception): def __init__(self, status: int, text: str) -> None: super().__init__(text) self.status_code = status class LengthFinishReasonError(Exception): """和 openai 的同名异常一样：解析结构化输出时发现 finish_reason=length。""" def __init__(self) -> None: super().__init__("Could not parse response content as the length limit was reached") self.completion = None class ConnectError(Exception): pass @pytest.fixture(autouse=True) def _no_wait(monkeypatch): """退避不真等：记下要等多久。""" waited: list[float] = [] async def fake_sleep(seconds: float) -> None: waited.append(seconds) monkeypatch.setattr(D, "_sleep", fake_sleep) return waited def _status(outcomes: dict) -> dict: return {t: o.status for t, o in outcomes.items()} def _hollow(call: Call) -> Reply: return Reply(obj={"tables": [{"table": t} for t in call.asked]}) def _exhausted(seg: str = "这一段", budget: int = 32768) -> str: return D.REASON_EXHAUSTED.format(budget=budget, seg=seg) EXHAUSTED = Reply(text="", finish="length", thinking="……") class Reports: """进度回调 call_done 收到的每次调用的结果。""" def __init__(self) -> None: self.items: list[D.CallReport] = [] async def __call__(self, report: D.CallReport) -> None: self.items.append(report) def hooks(self) -> D.Hooks: return D.Hooks(call_done=self) @property def failures(self) -> list[str | None]: return [r.failure for r in self.items] # -------------------------------------------------------------------------- # 思考用完额度、空壳：拆开再问 # -------------------------------------------------------------------------- async def test_thinking_used_up_the_budget_splits_the_segment(): tables = {"orders_x": _table(24)} ok = answering(tables) def handle(call: Call) -> Reply: if call.n == 0: # 只想不写：正文为空、只有思考，额度用完 return Reply(text="", finish="length", thinking="先逐列想一遍……") return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["orders_x"]) o = got["orders_x"] assert o.status == "ok" and o.error is None and o.items == 3 + 24 * 2 assert len(model.calls) == 3 cols = [f"c{i:03d}" for i in range(1, 25)] assert model.calls[1].asked == {"orders_x": ("table+columns", cols[:12])} assert model.calls[2].asked == {"orders_x": ("columns", cols[12:])} # 后半段在前半段之后问，带着刚起草的表级内容做背景 assert catalog.CATALOG_DRAFT_PROMPT_LABELS["drafted"] in model.calls[2].text assert "orders_x的中文名" in model.calls[2].text async def test_length_finish_reason_error_also_splits(): tables = {"orders_x": _table(24)} ok = answering(tables) def handle(call: Call) -> Reply: if call.n == 0: raise LengthFinishReasonError() return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["orders_x"]) assert got["orders_x"].status == "ok" and len(model.calls) == 3 async def test_hollow_shell_splits_the_pack_by_table(): tables = {"a_tab": _table(6), "b_tab": _table(6)} ok = answering(tables) def handle(call: Call) -> Reply: if call.n == 0: return Reply(obj={"tables": [{"table": "a_tab"}, {"table": "b_tab"}]}) return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["a_tab", "b_tab"]) assert _status(got) == {"a_tab": "ok", "b_tab": "ok"} assert [set(c.asked) for c in model.calls] == [{"a_tab", "b_tab"}, {"a_tab"}, {"b_tab"}] assert all(c.asked[t][0] == "full" for c in model.calls[1:] for t in c.asked) async def test_splitting_stops_at_three_levels_and_five_columns(monkeypatch): monkeypatch.setattr(catalog, "MODEL_PART_MAX_CALLS", 99) monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) # 不让按比例止损先停下来 tables = {"wide40": _table(40)} model = Scripted(lambda call: EXHAUSTED) got = await D.run_model_draft(model, _source(tables), ["wide40"]) o = got["wide40"] assert len(model.calls) == 1 + 2 + 4 + 8, "40 → 20 → 10 → 5 列，最多拆三层" assert min(len(c.asked["wide40"][1] or []) for c in model.calls[1:]) == 5 assert o.status == "failed" and o.items == 0 assert o.error == "模型的思考和输出用完了额度（32768 个 token），这一段没写出来；已拆小重试仍没写出" small = Scripted(lambda call: Reply(text="", finish="length")) got = await D.run_model_draft(small, _source({"t9": _table(9)}), ["t9"]) assert len(small.calls) == 1, "不到 10 列拆不出两段 5 列" assert got["t9"].error == "模型的思考和输出用完了额度（32768 个 token），这一段没写出来" async def test_one_part_makes_at_most_eight_calls(monkeypatch): """一段（连同它拆出来的、补问的、重试的）最多调用 8 次：40 列一直用完额度，拆到第 8 次就停，剩下没问的记没写出来、 原因写清。以前一段 40 列拆到底要 15 次。""" monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) assert catalog.MODEL_PART_MAX_CALLS == 8 model = Scripted(lambda call: EXHAUSTED) got = await D.run_model_draft(model, _source({"t40": _table(40)}), ["t40"]) assert len(model.calls) == 8 o = got["t40"] assert o.status == "failed" assert o.error == (_exhausted() + "；每段最多调用 8 次模型，这一段到了次数仍没写完，已停止重试" + D.REASON_SPLIT) # 网络错误的重试也算在这 8 次里 def flaky(call: Call) -> Reply: raise ConnectError("All connection attempts failed") monkeypatch.setattr(catalog, "MODEL_PART_MAX_CALLS", 2) model = Scripted(flaky) got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"]) assert len(model.calls) == 2 assert got["t4"].error == "每段最多调用 2 次模型，这一段到了次数仍没写完，已停止重试" async def test_an_unsplittable_hollow_shell_tries_text_once_then_stays_on_text(): """有的网关的结构化输出只给空壳：拆不开时换纯文本问一次，写出来了，这次起草后面直接走纯文本。 用宽表看先后：表级那段（拆不开）写完之后才问列分段。""" tables = {"wide45": _table(45)} ok = answering(tables) def handle(call: Call) -> Reply: if call.mode == "structured": return Reply(obj={"tables": [{"table": t} for t in call.asked]}) return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["wide45"]) assert got["wide45"].status == "ok" and len(got["wide45"].notes["columns"]) == 45 assert [c.mode for c in model.calls] == ["structured", "text", "text", "text"] assert catalog.CATALOG_DRAFT_TEXT_PROMPT in model.calls[1].text, "纯文本那条路要说清 JSON 的格式" def _text_only(handle_text: Callable[[Call], Reply]) -> Callable[[Call], Reply]: """结构化输出不支持（这次起草一开始就改走纯文本），纯文本交给 handle_text。""" def handle(call: Call) -> Reply: if call.mode == "structured": raise _Status(400, "This response_format type is unavailable now") return handle_text(call) return handle async def test_an_unsplittable_hollow_shell_is_asked_again_in_text_mode_too(): """已经走纯文本时，拆不开的段（不到 10 列）只给空壳：原样再问一次，写出来就用；仍是空壳才记没写出来。 以前纯文本模式下不补问，整张表直接记没写出来。""" tables = {"t8": _table(8)} ok = answering(tables) model = Scripted(_text_only(lambda call: _hollow(call) if call.n == 1 else ok(call))) got = await D.run_model_draft(model, _source(tables), ["t8"]) assert [c.mode for c in model.calls] == ["structured", "text", "text"] assert got["t8"].status == "ok" and len(got["t8"].notes["columns"]) == 8 model = Scripted(_text_only(_hollow)) got = await D.run_model_draft(model, _source(tables), ["t8"]) assert [c.mode for c in model.calls] == ["structured", "text", "text"], "再问一次，不多问" assert got["t8"].status == "failed" and got["t8"].error == catalog.MODEL_EMPTY_REASON # -------------------------------------------------------------------------- # 截断、缺表、缺列 # -------------------------------------------------------------------------- async def test_truncated_json_keeps_complete_columns_and_asks_only_the_rest(): tables = {"t12": _table(12)} full = json.dumps({"tables": [{"table": "t12", **_reply("t12", tables["t12"])}]}, ensure_ascii=False) cut = full[:full.index('"c006"') + 8] # 第 6 列写到一半 ok = answering(tables) def handle(call: Call) -> Reply: if call.n == 0: return Reply(text=cut, finish="length") return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t12"]) o = got["t12"] assert o.status == "ok" and set(o.notes["columns"]) == {f"c{i:03d}" for i in range(1, 13)} assert o.notes["label"]["value"] == "t12的中文名" assert len(model.calls) == 2 assert model.calls[1].asked == {"t12": ("columns", [f"c{i:03d}" for i in range(6, 13)])} async def test_truncated_tool_call_arguments_are_salvaged_too(): """结构化输出走工具调用时，截断的参数在 invalid_tool_calls 里（一段字符串）。""" tables = {"t12": _table(12)} full = json.dumps({"tables": [{"table": "t12", **_reply("t12", tables["t12"])}]}, ensure_ascii=False) cut = full[:full.index('"c004"') + 3] # 第 4 列的列名写到一半 ok = answering(tables) class ToolModel(Scripted): async def _run(self, messages, mode, include_raw): if not self.calls: self.calls.append(Call(mode, "", {}, 0)) raw = AIMessage(content="", invalid_tool_calls=[{"name": "catalog_draft", "args": cut, "id": "x", "error": None, "type": "invalid_tool_call"}], response_metadata={"finish_reason": "length"}) return {"raw": raw, "parsed": None, "parsing_error": ValueError("截断")} return await super()._run(messages, mode, include_raw) model = ToolModel(ok) got = await D.run_model_draft(model, _source(tables), ["t12"]) assert got["t12"].status == "ok" assert model.calls[1].asked == {"t12": ("columns", [f"c{i:03d}" for i in range(4, 13)])} async def test_missing_tables_are_asked_again_once(): """缺了的表只补缺的、补一次；补问时整张表还是一项都没有，按只给空壳处理（拆不开就换纯文本问一次）。""" tables = {"a_tab": _table(4), "b_tab": _table(4), "c_tab": _table(4)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) keep = {0: {"a_tab"}, 1: {"b_tab"}}.get(call.n, set()) reply.obj["tables"] = [e for e in reply.obj["tables"] if e["table"] in keep] return reply model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["a_tab", "b_tab", "c_tab"]) assert len(model.calls) == 3 assert model.calls[1].asked == {"b_tab": ("full", None), "c_tab": ("full", None)} assert model.calls[2].asked == {"c_tab": ("full", None)} and model.calls[2].mode == "text" assert _status(got) == {"a_tab": "ok", "b_tab": "ok", "c_tab": "failed"} assert got["c_tab"].error == "模型的回复里没有这张表" async def test_a_table_still_missing_after_the_followup_is_split_like_a_hollow_shell(): """补问时整张表还是一项都没有：照第一次的规则拆小（这里按列拆成两半），不直接记失败。""" tables = {"a_tab": _table(4), "b_tab": _table(12)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) if call.n < 2: reply.obj["tables"] = [e for e in reply.obj["tables"] if e["table"] == "a_tab"] return reply model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["a_tab", "b_tab"]) assert _status(got) == {"a_tab": "ok", "b_tab": "ok"} cols = [f"c{i:03d}" for i in range(1, 13)] assert [c.asked for c in model.calls[2:]] == [{"b_tab": ("table+columns", cols[:6])}, {"b_tab": ("columns", cols[6:])}] async def test_columns_still_missing_after_the_followup_are_not_written(): """补问过仍然没列出来的列算没写出来：表是 partial，说清哪部分没写出来（不算拿不准，旧的模型起草项不删）。""" tables = {"t8": _table(8)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) for entry in reply.obj["tables"]: entry["columns"] = [c for c in entry.get("columns", []) if c["name"] not in ("c007", "c008")] return reply model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t8"]) assert len(model.calls) == 2 and model.calls[1].asked == {"t8": ("columns", ["c007", "c008"])} o = got["t8"] assert o.status == "partial" assert o.error == ("这张表 8 列，有 2 列（c007、c008）没写出来：模型的回复里漏了这些列，补问一次仍没有给出。" "可以只重试这张表") assert set(o.notes["columns"]) == {f"c{i:03d}" for i in range(1, 7)} assert o.blank == [], "没写出来的列不算留空" compat = (await catalog.draft_with_model(Scripted(handle), _source(tables), ["t8"]))["t8"] assert compat.covered == frozenset(), "没写全的不算完整起草过" async def test_columns_listed_blank_in_the_followup_are_left_blank_not_failed(): """补问时模型把漏的列列出来了、字段却留空：是拿不准（记留空），不算「漏了这些列」。""" tables = {"t8": _table(8)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) entry = reply.obj["tables"][0] if call.n == 0: entry["columns"] = [c for c in entry["columns"] if c["name"] not in ("c003", "c007")] else: entry["columns"] = [{"name": c["name"]} for c in entry["columns"]] return reply model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), ["t8"], hooks=reports.hooks()) assert model.calls[1].asked == {"t8": ("columns", ["c003", "c007"])} o = got["t8"] assert o.status == "ok" and o.error is None and o.blank == ["c003", "c007"] # 补问的那次把请写的列都列出来了（字段留空是拿不准）：用量和提醒里也不算一次「只给空壳」 assert reports.failures == [None, None] async def test_a_followup_that_comes_back_as_an_empty_shell_leaves_the_table_partial(): """40 列的表第一次只回了 20 列，补问只给一个空壳：没写出来的 20 列照实记下，表是 partial。""" tables = {"t40": _table(40)} ok = answering(tables) def handle(call: Call) -> Reply: if call.n == 0: reply = ok(call) reply.obj["tables"][0]["columns"] = reply.obj["tables"][0]["columns"][:20] return reply return Reply(obj={"tables": [{"table": "t40"}]}) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t40"]) o = got["t40"] assert len(model.calls) == 2 assert model.calls[1].asked == {"t40": ("columns", [f"c{i:03d}" for i in range(21, 41)])} assert o.status == "partial" and len(o.notes["columns"]) == 20 assert o.error == ("这张表 40 列，有 20 列（c021 到 c040）没写出来：模型的回复里漏了这些列，补问一次仍没有给出。" "可以只重试这张表") async def test_columns_listed_with_empty_fields_are_left_to_the_model(): """列出了、字段留空的列是模型拿不准：不补问，也不算没写出来；记成留空（表的状态不因此变）。""" tables = {"t8": _table(8)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) for entry in reply.obj["tables"]: entry["columns"] = [c if c["name"] not in ("c007", "c008") else {"name": c["name"]} for c in entry.get("columns", [])] return reply model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t8"]) assert len(model.calls) == 1 assert got["t8"].status == "ok" and got["t8"].error is None assert set(got["t8"].notes["columns"]) == {f"c{i:03d}" for i in range(1, 7)} assert got["t8"].blank == ["c007", "c008"] async def test_blank_columns_are_the_ones_without_a_label_or_meaning(): """留空的列：请模型写了、中文名和含义却都没给（只给了度量类型、单位，或者给的都没通过校验）。已确认的、库注释已经说清 的不算：给模型的提示词里它们以此为准，模型不写也不缺。""" meta = _table(9) meta["columns"][5]["comment"] = "会员等级" # c006：库注释就是中文名 meta["columns"][8]["comment"] = "C009" # c009：注释就是列名，说明不了什么 tables = {"t9": meta} existing = {"t9": {"columns": {"c005": {"label": catalog.make_item("已确认的名字", "human")}}}} def handle(call: Call) -> Reply: return Reply(obj={"tables": [{"table": "t9", "label": "某表", "columns": [ {"name": "c001", "label": "第一列"}, {"name": "c002", "meaning": "第二列的含义"}, {"name": "c003", "measure": "attribute", "unit": "元"}, # 中文名、含义都没给 {"name": "c004", "label": "=c001+c002"}, # 像公式：校验丢掉 {"name": "c005"}, {"name": "c006"}, {"name": "c007"}, {"name": "c008"}, {"name": "c009"}]}]}) got = await D.run_model_draft(Scripted(handle), _source(tables), ["t9"], existing=existing) o = got["t9"] assert o.status == "ok" and o.error is None assert o.blank == ["c003", "c004", "c007", "c008", "c009"] # -------------------------------------------------------------------------- # 额度被拒、结构化输出不支持、网络错误、超时 # -------------------------------------------------------------------------- async def test_budget_rejection_falls_back_to_8192_for_the_rest(): tables = {"wide45": _table(45)} def refuse(call: Call) -> Reply: raise _Status(400, "max_tokens: 32768 > 8192, which is the maximum allowed number of output tokens") big, small = Scripted(refuse), Scripted(answering(tables)) model = catalog.DraftModel(big, small, "m") got = await D.run_model_draft(model, _source(tables), ["wide45"]) assert got["wide45"].status == "ok" assert len(big.calls) == 1 and len(small.calls) == 3, "退回之后，表级那段重来、两个列分段都用小额度" assert model.budget == 8192 async def test_exhausted_message_names_the_budget_actually_used(): def refuse(call: Call) -> Reply: raise _Status(400, "max_tokens is too large") model = catalog.DraftModel(Scripted(refuse), Scripted(lambda c: Reply(text="", finish="length")), "m") got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"]) assert got["t4"].error == "模型的思考和输出用完了额度（8192 个 token），这一段没写出来" async def test_structured_output_unsupported_switches_to_text_for_the_rest(): tables = {"wide45": _table(45)} ok = answering(tables) def handle(call: Call) -> Reply: if call.mode == "structured": raise _Status(400, "This response_format type is unavailable now") return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["wide45"]) assert got["wide45"].status == "ok" assert [c.mode for c in model.calls] == ["structured", "text", "text", "text"] @pytest.mark.parametrize("thinking", ["", "先想想每一列是什么……"], ids=["plain", "with-thinking"]) async def test_structured_output_that_comes_back_as_prose_falls_back_to_text(thinking): """网关接受了结构化输出的参数却没照做（模型写成了一段说明文字）：换纯文本问，不当成额度用完去拆，也不直接判失败。""" tables = {"t40": _table(40)} ok = answering(tables) def handle(call: Call) -> Reply: if call.mode == "structured": return Reply(text="好的。这张表记录的是一笔笔业务，各列的含义如下：第一列是编号……", thinking=thinking) return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t40"]) assert got["t40"].status == "ok" and len(got["t40"].notes["columns"]) == 40 assert [c.mode for c in model.calls] == ["structured", "text"] assert catalog.CATALOG_DRAFT_TEXT_PROMPT in model.calls[1].text async def test_prose_on_both_paths_is_a_format_failure_not_a_budget_one(): def handle(call: Call) -> Reply: return Reply(text="这张表看起来是业务明细，具体含义需要结合业务确认。", thinking="……") model = Scripted(handle) got = await D.run_model_draft(model, _source({"t40": _table(40)}), ["t40"]) assert [c.mode for c in model.calls] == ["structured", "text"], "格式不对不拆小重试" assert got["t40"].status == "failed" and got["t40"].error == "模型的回复不是约定的格式" async def test_a_model_without_structured_output_goes_to_text(): class TextOnly(Scripted): def with_structured_output(self, schema, **_): raise NotImplementedError tables = {"t4": _table(4)} model = TextOnly(answering(tables)) got = await D.run_model_draft(model, _source(tables), ["t4"]) assert got["t4"].status == "ok" and [c.mode for c in model.calls] == ["text"] async def test_network_errors_are_retried_with_backoff(_no_wait): tables = {"t4": _table(4)} ok = answering(tables) def handle(call: Call) -> Reply: if call.n < 2: raise ConnectError("All connection attempts failed") return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t4"]) assert got["t4"].status == "ok" and len(model.calls) == 3 assert _no_wait == list(catalog.MODEL_RETRY_DELAYS_S) async def test_server_errors_give_up_after_two_retries(): def handle(call: Call) -> Reply: raise _Status(503, "Service Unavailable") model = Scripted(handle) got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"]) assert len(model.calls) == 3 assert got["t4"].status == "failed" assert got["t4"].error == "调用模型出错：Service Unavailable（已重试 2 次）" async def test_other_errors_are_not_retried(): def handle(call: Call) -> Reply: raise RuntimeError("网关拒绝了这次请求") model = Scripted(handle) got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"]) assert len(model.calls) == 1 and got["t4"].error == "调用模型出错：网关拒绝了这次请求" async def test_a_call_past_the_time_limit_is_split_not_retried(monkeypatch, _no_wait): """超时当作这一段太大：不原样重试同样大小的段，拆成两半再问（同额度用完）。超时的调用拿不到用量，记 0。""" monkeypatch.setattr(catalog, "MODEL_TIMEOUT_S", 0.05) tables = {"t24": _table(24)} ok = answering(tables, usage=(100, 50)) async def handle(call: Call) -> Reply: if call.n == 0: await asyncio.sleep(1) return ok(call) model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), ["t24"], hooks=reports.hooks()) cols = [f"c{i:03d}" for i in range(1, 25)] assert got["t24"].status == "ok" assert [c.asked for c in model.calls] == [{"t24": ("full", None)}, {"t24": ("table+columns", cols[:12])}, {"t24": ("columns", cols[12:])}] assert reports.failures == ["timeout", None, None] assert [(r.input_tokens, r.output_tokens) for r in reports.items] == [(0, 0), (100, 50), (100, 50)] assert _no_wait == [], "不退避重试" async def test_a_small_part_past_the_time_limit_is_not_retried(monkeypatch): monkeypatch.setattr(catalog, "MODEL_TIMEOUT_S", 0.05) async def slow(call: Call) -> Reply: await asyncio.sleep(1) return Reply(obj={"tables": []}) model = Scripted(slow) got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"]) assert len(model.calls) == 1, "拆不开、也不原样重试" assert got["t4"].status == "failed" and got["t4"].error == "模型 0.05 秒内没有写完这一段" # -------------------------------------------------------------------------- # 分段：宽表拆成表级一段 + 列分段，窄表拼段 # -------------------------------------------------------------------------- async def test_a_wide_table_is_split_into_a_table_part_and_column_parts(): tables = {"wide95": _table(95, comment="虚构的宽表")} model = Scripted(answering(tables)) done: list[str] = [] async def part_done(table: str) -> None: done.append(table) got = await D.run_model_draft(model, _source(tables), ["wide95"], hooks=D.Hooks(part_done=part_done)) o = got["wide95"] assert o.status == "ok" and o.parts_total == 4 and o.parts_done == 4 and done == ["wide95"] * 4 assert len(o.notes["columns"]) == 95 and o.notes["label"]["value"] == "wide95的中文名" first, *rest = model.calls assert first.asked == {"wide95": ("table", None)} assert catalog.CATALOG_DRAFT_PROMPT_LABELS["brief_columns"].format(n=95) in first.text assert sorted(len(c.asked["wide95"][1]) for c in rest) == [31, 32, 32] for call in rest: assert call.asked["wide95"][0] == "columns" assert "wide95的中文名" in call.text and "虚构的宽表" in call.text, "列分段带着表注释和表级起草结果" assert {c for call in rest for c in call.asked["wide95"][1]} == {f"c{i:03d}" for i in range(1, 96)} @pytest.mark.parametrize("n, sizes", [(164, [33, 33, 33, 33, 32]), (95, [32, 32, 31]), (41, [21, 20]), (80, [40, 40]), (81, [27, 27, 27])]) async def test_column_parts_are_even_without_a_small_tail(monkeypatch, n, sizes): """列分段数 = 列数 / 40 向上取整，各段列数平均（相差不超过 1）：不再留下几列的尾段（实测会思考的模型常对这种小段只给 空壳）。表级那段照旧在最前面。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) tables = {"wide": _table(n)} model = Scripted(answering(tables)) got = await D.run_model_draft(model, _source(tables), ["wide"]) o = got["wide"] assert o.status == "ok" and o.parts_total == 1 + len(sizes) == D.parts_for(n) assert model.calls[0].asked == {"wide": ("table", None)} parts = [c.asked["wide"][1] for c in model.calls[1:]] assert [len(p) for p in parts] == sizes assert [c for p in parts for c in p] == [f"c{i:03d}" for i in range(1, n + 1)], "按表结构的顺序连续切开" async def test_a_wide_table_with_one_failed_part_is_partial_and_says_which(): tables = {"wide95": _table(95)} ok = answering(tables) def handle(call: Call) -> Reply: part, cols = call.asked["wide95"] if part == "columns" and "c041" in cols: raise RuntimeError("网关拒绝") return ok(call) model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["wide95"]) o = got["wide95"] assert o.status == "partial" and len(o.notes["columns"]) == 63 assert o.error == ("这张表 95 列，分 4 段起草，第 3 段（c033 到 c064，共 32 列）没写出来：调用模型出错：网关拒绝。" "可以只重试这张表") compat = (await catalog.draft_with_model(Scripted(handle), _source(tables), ["wide95"]))["wide95"] assert compat.covered == frozenset() and compat.error == o.error, "没写全的不算完整起草过，旧的模型项不删" async def test_reasons_name_each_part_and_say_these_parts(monkeypatch): """几段因为同一个原因没写出来：原因里写「这几段」；每段写上列的范围（列名照表结构原样）；拆过的才补「已拆小重试」。""" monkeypatch.setattr(catalog, "MODEL_SHRINK_AFTER", 99) monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) tables = {"wide95": _table(95)} ok = answering(tables) def handle(call: Call) -> Reply: part, cols = call.asked["wide95"] if part == "columns" and cols[0] >= "c033": # 第 3、4 段（c033 起）一直用完额度 return EXHAUSTED return ok(call) got = await D.run_model_draft(Scripted(handle), _source(tables), ["wide95"]) o = got["wide95"] assert o.status == "partial" and len(o.notes["columns"]) == 32 assert o.error == ("这张表 95 列，分 4 段起草，第 3 段（c033 到 c064，共 32 列）、第 4 段（c065 到 c095，共 31 列）没写出来：" + _exhausted("这几段") + D.REASON_SPLIT + "。可以只重试这张表") async def test_a_part_that_was_never_halved_does_not_say_it_was(): """被截断后只补剩下的（不是拆小）：补的那次也用完了额度、又拆不开，原因里不写「已拆小重试」。""" tables = {"t9": _table(9)} full = json.dumps({"tables": [{"table": "t9", **_reply("t9", tables["t9"])}]}, ensure_ascii=False) cut = full[:full.index('"c005"') + 8] def handle(call: Call) -> Reply: return Reply(text=cut, finish="length") if call.n == 0 else EXHAUSTED model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["t9"]) o = got["t9"] assert len(model.calls) == 2 and model.calls[1].asked == {"t9": ("columns", [f"c{i:03d}" for i in range(5, 10)])} assert o.status == "partial" assert o.error == "这张表 9 列，有 5 列（c005 到 c009）没写出来：" + _exhausted() + "。可以只重试这张表" @pytest.mark.parametrize("missing, said", [ (["c004"], "有 1 列（c004）"), (["c003", "c007"], "有 2 列（c003、c007）"), (["c002", "c004", "c006", "c008"], "有 4 列（c002、c004、c006 等）"), (["c003", "c004", "c005", "c006"], "有 4 列（c003 到 c006）"), ]) async def test_columns_not_written_are_named(missing, said): tables = {"t8": _table(8)} ok = answering(tables) def handle(call: Call) -> Reply: reply = ok(call) for entry in reply.obj["tables"]: entry["columns"] = [c for c in entry.get("columns", []) if c["name"] not in missing] return reply got = await D.run_model_draft(Scripted(handle), _source(tables), ["t8"]) assert got["t8"].error == f"这张表 8 列，{said}没写出来：{D.REASON_UNANSWERED}。可以只重试这张表" async def test_tables_are_packed_up_to_40_columns_and_4_tables(): widths = {"ta": 15, "tb": 15, "tc": 15, "td": 5, "te": 5, "tf": 5, "tg": 5, "th": 41} tables = {name: _table(n) for name, n in widths.items()} model = Scripted(answering(tables)) got = await D.run_model_draft(model, _source(tables), list(widths)) assert set(_status(got).values()) == {"ok"} asked = sorted((tuple(sorted(c.asked)), c.asked[next(iter(c.asked))][0]) for c in model.calls) assert asked == sorted([ (("ta", "tb"), "full"), (("tc", "td", "te", "tf"), "full"), (("tg",), "full"), (("th",), "table"), (("th",), "columns"), (("th",), "columns"), ]) # -------------------------------------------------------------------------- # 报错去重、止损、几次起草轮流用模型 # -------------------------------------------------------------------------- async def test_the_same_reason_is_said_once_even_when_some_parts_were_split(monkeypatch): """宽表的表级那段拆不开、列分段拆过：原因只说一次，「已拆小重试」只跟在真拆过的那几段后面说。""" monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) monkeypatch.setattr(catalog, "MODEL_PART_MAX_CALLS", 99) model = Scripted(_hollow) got = await D.run_model_draft(model, _source({"wide45": _table(45)}), ["wide45"]) o = got["wide45"] assert o.status == "failed" assert o.error == ("这张表 45 列，分 3 段起草，第 1 段（表级的项）、第 2 段（c001 到 c023，共 23 列）、" "第 3 段（c024 到 c045，共 22 列）都没写出来：模型只给了空壳，没有写出可用的内容；" "其中第 2、3 段已拆小重试仍没写出") def test_reasons_are_joined_by_their_base(): """(原因, 第几段)：同样的原因只说一次；同一个原因落在几段上写「这几段」，同一段拆出来的几块仍是「这一段」。 「已拆小重试仍没写出」只跟在真拆过的那几段后面说：说到的段都拆过时不点名，只有一部分拆过时写「其中第几段」。""" W = D._Why ex = D.REASON_EXHAUSTED.format(budget=32768, seg=D.SEG) assert D._join([(W(D.REASON_HOLLOW), 1), (W(D.REASON_HOLLOW, split=True), 2)]) == \ D.REASON_HOLLOW + "；其中第 2 段已拆小重试仍没写出" assert D._join([(W(ex, split=True), 2), (W(D.REASON_HOLLOW), 1)]) == \ f"{_exhausted()}；{D.REASON_HOLLOW}；其中第 2 段已拆小重试仍没写出" assert D._join([(W(ex, split=True), 2), (W(ex, split=True), 3)]) == _exhausted("这几段") + D.REASON_SPLIT assert D._join([(W(ex, split=True), 3), (W(ex), 4), (W(ex, split=True), 2)]) == \ _exhausted("这几段") + "；其中第 2、3 段已拆小重试仍没写出" assert D._join([(W(ex), 2), (W(ex), 3)]) == _exhausted("这几段") assert D._join([(W(ex), 2), (W(ex, split=True), 2)]) == _exhausted() + D.REASON_SPLIT, "同一段里有一块拆过就算拆过" assert D._join([(W("甲"), 1), (W("乙"), 1), (W("丙", split=True), 1)]) == "甲；乙", "没说出来的原因不补后缀" halted = D.REASON_HALTED.format(window=catalog.MODEL_HALT_WINDOW, on=catalog.MODEL_HALT_ON) assert halted == "；最近 10 次调用有 7 次以上没写出来，已停止拆小重试，建议更换起草模型" assert D._join([(W(D.REASON_HOLLOW, halted=halted), 1), (W(D.REASON_HOLLOW), 2)]) == D.REASON_HOLLOW + halted assert D._join([(W(D.REASON_HOLLOW, halted=halted), 1), (W(D.REASON_HOLLOW, halted=halted), 3)]) == \ D.REASON_HOLLOW + halted, "止损只说一次" assert D._join([(W("调用模型出错：{坏的}"), 1)]) == "调用模型出错：{坏的}", "报错原文里的花括号原样留着" def _by_table(model: Scripted) -> dict[str, int]: out: dict[str, int] = {} for call in model.calls: for name in call.asked: out[name] = out.get(name, 0) + 1 return out HALTED = "；最近 10 次调用有 7 次以上没写出来，已停止拆小重试，建议更换起草模型" async def test_when_most_recent_calls_fail_the_rest_are_asked_once_until_the_ratio_drops(monkeypatch): """按比例止损：这次起草最近 10 次调用有 7 次以上没写出来，之后的段不再拆小、每块只问一次，报错说明已停止拆小重试； 比例降下去（最近 10 次里少于 7 次）就恢复拆小。偶尔一次写出来不会清零：t4 到 t6 写出来了，最近 10 次里仍有 7 次没写 出来，照旧不拆；t7 写出来以后才降下去，t8 照常拆小。一次只起草一段（并发 1），先后是确定的。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) assert (catalog.MODEL_HALT_WINDOW, catalog.MODEL_HALT_ON) == (10, 7) tables = {f"t{i}": _table(10) for i in range(1, 9)} ok = answering(tables) bad = {"t1", "t2", "t3", "t8"} model, reports = Scripted(lambda call: EXHAUSTED if set(call.asked) & bad else ok(call)), Reports() got = await D.run_model_draft(model, _source(tables), list(tables), hooks=reports.hooks()) assert _by_table(model) == {"t1": 3, "t2": 3, "t3": 1, "t4": 1, "t5": 1, "t6": 1, "t7": 1, "t8": 3} assert got["t1"].error == _exhausted() + D.REASON_SPLIT assert got["t3"].error == _exhausted() + HALTED assert {got[f"t{i}"].status for i in range(4, 8)} == {"ok"} assert got["t8"].error == _exhausted() + D.REASON_SPLIT, "比例降下去以后照常拆小" assert len(reports.items) == len(model.calls) async def test_while_halted_a_chunk_is_not_asked_again_for_missing_columns(monkeypatch): """止损期间每块只问一次：回复里漏了几列不补问，漏的列记没写出来、说明已停止拆小重试。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {**{f"x{i}": _table(4) for i in range(1, 8)}, "t8": _table(8)} ok = answering(tables) def handle(call: Call) -> Reply: if "t8" not in call.asked: raise RuntimeError("网关拒绝了这次请求") reply = ok(call) reply.obj["tables"][0]["columns"] = reply.obj["tables"][0]["columns"][:6] return reply model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), list(tables)) assert _by_table(model)["t8"] == 1, "漏了 c007、c008 也不补问" o = got["t8"] assert o.status == "partial" assert o.error == ("这张表 8 列，有 2 列（c007、c008）没写出来：模型的回复里漏了这些列" + HALTED + "。可以只重试这张表") #: 200 列的宽表分 6 段起草（表级一段、列分段 5 段），逐段写上列的范围 WIDE200_PARTS = ("这张表 200 列，分 6 段起草，第 1 段（表级的项）、第 2 段（c001 到 c040，共 40 列）、" "第 3 段（c041 到 c080，共 40 列）、第 4 段（c081 到 c120，共 40 列）、第 5 段（c121 到 c160，共 40 列）、" "第 6 段（c161 到 c200，共 40 列）") async def test_a_wide_table_skips_its_remaining_parts_after_three_dead_parts_in_a_row(monkeypatch): """同一张表连续 3 个原始段都一项没写出来：这张表剩下的段跳过不问，原因写清跳过了几段、可以换一个起草模型后只重试 这张表。200 列分 6 段（表级一段、列分段 5 段）。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) assert catalog.MODEL_TABLE_GIVE_UP_AFTER == 3 tables = {"wide200": _table(200)} def refuse(call: Call) -> Reply: raise RuntimeError("网关拒绝") model = Scripted(refuse) got = await D.run_model_draft(model, _source(tables), ["wide200"]) o = got["wide200"] assert len(model.calls) == 3, "第 4 到 6 段跳过不问" assert (o.status, o.parts_total, o.parts_done) == ("failed", 6, 6) assert o.error == (WIDE200_PARTS + "都没写出来：调用模型出错：网关拒绝；" "这张表接连 3 段都没写出来，剩下的 3 段已跳过，可以换一个起草模型后只重试这张表") # 表级那段写出来了、后面接连 3 段没写出来：只写出一部分，剩下的 2 段跳过；不再另说「可以只重试这张表」 ok = answering(tables) model = Scripted(lambda call: ok(call) if call.n == 0 else refuse(call)) got = await D.run_model_draft(model, _source(tables), ["wide200"]) o = got["wide200"] assert len(model.calls) == 4 and o.status == "partial" assert o.error == ("这张表 200 列，分 6 段起草，第 2 段（c001 到 c040，共 40 列）、第 3 段（c041 到 c080，共 40 列）、" "第 4 段（c081 到 c120，共 40 列）、第 5 段（c121 到 c160，共 40 列）、第 6 段（c161 到 c200，共 40 列）" "没写出来：调用模型出错：网关拒绝；这张表接连 3 段都没写出来，剩下的 2 段已跳过，" "可以换一个起草模型后只重试这张表") async def test_a_part_that_writes_resets_the_dead_parts_of_its_table(monkeypatch): """中间有一段写出来了，就不算「接连」：200 列的第 2、3、5、6 段没写出来、第 4 段写出来了，每段都照常问。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) tables = {"wide200": _table(200)} ok = answering(tables) def handle(call: Call) -> Reply: part, cols = call.asked["wide200"] if part == "table" or "c081" in (cols or []): return ok(call) raise RuntimeError("网关拒绝") model = Scripted(handle) got = await D.run_model_draft(model, _source(tables), ["wide200"]) assert len(model.calls) == 6 and got["wide200"].status == "partial" assert "已跳过" not in got["wide200"].error async def test_the_column_parts_of_a_wide_table_stop_splitting_too(monkeypatch): """按比例止损在宽表里同样起作用，接连没写出来的宽表剩下的段跳过：95 列一直只给空壳（已经走纯文本），第 2 段拆小到 止损；之后的第 3 段只问一次；接连 3 段没写出来，第 4 段跳过。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) model = Scripted(_text_only(_hollow)) got = await D.run_model_draft(model, _source({"wide95": _table(95)}), ["wide95"]) o = got["wide95"] asked = [c for c in model.calls if c.mode == "text"] assert len(asked) <= 2 + 8 + 1 assert not any("c065" in (c.asked["wide95"][1] or []) for c in asked), "第 4 段没问" assert o.status == "failed" and o.error.startswith(WIDE95_PARTS + "都没写出来：") assert HALTED in o.error assert o.error.endswith("这张表接连 3 段都没写出来，剩下的 1 段已跳过，可以换一个起草模型后只重试这张表") async def test_a_wide_table_that_never_comes_back_is_bounded_by_the_first_wave(): """默认并发（3 段同时起草）：294 列（8 个列分段）一直只给空壳：表级那段 2 次，头一批同时起草的列分段拆小到止损或 封顶；接连 3 段没写出来以后，没开始的段跳过不问。以前所有列分段一下子都排上去、一层一层地拆，要 168 次。""" model = Scripted(_hollow) got = await D.run_model_draft(model, _source({"wide294": _table(294)}), ["wide294"]) o = got["wide294"] assert o.status == "failed" and o.parts_total == 9 and o.parts_done == 9 assert len(model.calls) <= 2 + 4 * 8 assert "已跳过" in o.error asked = {c for call in model.calls for c in call.asked["wide294"][1] or []} assert "c294" not in asked, "最后一段没问" async def test_a_big_draft_does_not_hold_the_model_ahead_of_a_later_one(monkeypatch): """全进程的调用名额先来先放：一次起草不把它的所有段一下子都排上去，后开始的那次起草（另一个数据源）能轮上。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) big = {f"a{i}": _table(4) for i in range(1, 5)} small = {"b1": _table(4)} order: list[str] = [] release = asyncio.Event() def handler(tables: dict) -> Callable[[Call], Any]: ok = answering(tables) async def handle(call: Call) -> Reply: order.append(next(iter(call.asked))) if len(order) == 1: await release.wait() return ok(call) return handle first = asyncio.create_task(D.run_model_draft(Scripted(handler(big)), _source(big), list(big))) while not order: await asyncio.sleep(0) second = asyncio.create_task(D.run_model_draft(Scripted(handler(small)), _source(small), ["b1"])) for _ in range(20): await asyncio.sleep(0) release.set() got_big, got_small = await asyncio.gather(first, second) assert set(_status(got_big).values()) == {"ok"} and _status(got_small) == {"b1": "ok"} assert order == ["a1", "b1", "a2", "a3", "a4"] #: 95 列的宽表分 4 段起草，逐段写上列的范围 WIDE95_PARTS = ("这张表 95 列，分 4 段起草，第 1 段（表级的项）、第 2 段（c001 到 c032，共 32 列）、" "第 3 段（c033 到 c064，共 32 列）、第 4 段（c065 到 c095，共 31 列）") def _order(model: Scripted) -> list[tuple[str, str]]: return [(name, part) for call in model.calls for name, (part, _) in call.asked.items()] async def test_narrow_tables_are_not_cut_in_by_the_column_parts_of_a_wide_one(monkeypatch): """一个任务里按表的先后排队：宽表的表级那段写完，它的列分段排到已经排着的窄表段后面（先来先服务），不插队。 以前插到队首，窄表要等宽表的列分段全部派完。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {"wide": _table(45), "na": _table(4), "nb": _table(4)} model = Scripted(answering(tables)) got = await D.run_model_draft(model, _source(tables), ["wide", "na", "nb"]) assert set(_status(got).values()) == {"ok"} assert _order(model) == [("wide", "table"), ("na", "full"), ("nb", "full"), ("wide", "columns"), ("wide", "columns")] async def test_after_two_exhausted_or_timed_out_calls_new_parts_are_halved(monkeypatch): """额度用完或超时累计 2 次，之后新开始的段列数上限减半（40 → 20 → 10，最低 10）；已经开始的段照旧（按原因拆小）。 每次调用的结果里带着当前的上限（part_columns）。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {f"t{i}": _table(40) for i in range(1, 4)} ok = answering(tables) def handle(call: Call) -> Reply: (_, (part, cols)), = call.asked.items() width = 40 if part == "full" else len(cols) return EXHAUSTED if width > 10 else ok(call) # 这个模型一次只写得出 10 列 model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), list(tables), hooks=reports.hooks()) assert set(_status(got).values()) == {"ok"} asked = {name: [(part, len(cols or [])) for c in model.calls for n, (part, cols) in c.asked.items() if n == name] for name in tables} # t1：整张 40 列开始，用完额度拆成 20、20，又各拆成 10、10 assert asked["t1"][0] == ("full", 0) and len(asked["t1"]) == 7 # 第 2 次用完额度后上限是 20：t2 一开始就按 20 列切成两段（各自再拆一次） assert asked["t2"][0] == ("table+columns", 20) and ("columns", 20) in asked["t2"] # 第 4 次之后上限是 10：t3 一开始就切成 4 段 10 列，一次都没用完额度 assert asked["t3"] == [("table+columns", 10), ("columns", 10), ("columns", 10), ("columns", 10)] caps = [r.part_columns for r in reports.items] assert list(dict.fromkeys(caps)) == [40, 20, 10] assert caps[:2] == [40, 20], "第 2 次用完额度的那次调用之后就是 20" assert min(caps) == 10 == catalog.MODEL_MIN_PART_COLUMNS async def test_the_notice_comes_when_four_of_six_calls_fail_and_goes_when_fewer_than_two(monkeypatch): """最近 6 次调用有 4 次以上没写出来：提醒（写明次数和原因，建议换起草模型、减少表或停止）；之后最近 6 次里没写出来 的少于 2 次就清掉。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {f"t{i}": _table(4) for i in range(1, 10)} ok = answering(tables) def handle(call: Call) -> Reply: name = next(iter(call.asked)) if name in ("t1", "t2"): raise RuntimeError("网关拒绝了这次请求") if name == "t3": return EXHAUSTED if name == "t4": return _hollow(call) return ok(call) reports = Reports() got = await D.run_model_draft(Scripted(handle), _source(tables), list(tables), hooks=reports.hooks()) assert reports.failures == ["error", "error", "exhausted", "hollow", "hollow"] + [None] * 5 notices = [r.notice for r in reports.items] assert notices[:3] == [None] * 3 assert notices[3] == ("最近 4 次调用有 4 次没写出来（额度用完 1 次、只给空壳 1 次、出错 2 次）。建议在「设置 → 偏好设置」" "的「起草模型」换一个模型、缩小起草范围，或者停止这次起草") assert notices[4].startswith("最近 5 次调用有 5 次没写出来") assert notices[7].startswith("最近 6 次调用有 3 次没写出来"), "还有 2 次以上没写出来：留着，次数照实写" assert notices[8].startswith("最近 6 次调用有 2 次没写出来") assert notices[9] is None, "最近 6 次只有 1 次没写出来：清掉" assert got["t9"].status == "ok" # -------------------------------------------------------------------------- # 复核后的修正：每个原始段一共 8 次、一直超时的模型几轮就停、读回复出错也记用量、目录里已有的不算留空、补问列出了不算 # 没写出来、一项都没写出来的宽表逐段写范围、止损后纯文本不再原样再问 # -------------------------------------------------------------------------- def _width(call: Call, tables: dict) -> int: """这次调用请模型写几列（只写表级的项是 0）。""" return sum(len(tables[name]["columns"]) if part == "full" else len(cols or []) for name, (part, cols) in call.asked.items()) async def test_the_chunks_of_a_part_cut_by_the_smaller_cap_share_its_eight_calls(monkeypatch): """每个原始段（连同段大小减半后按上限分开的各块、拆出来的、补问的、重试的）一共最多调用 8 次。以前分开的每一块各有 8 次：上限缩到 10 列以后一直只给空壳，4 张窄表拼成的 40 列段调用了 20 次，单张 40 列的表也是 20 次。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) # 不让止损先停下来 tables = {"w40": _table(40), **{f"n{i}": _table(10) for i in range(1, 5)}, "x40": _table(40)} def handle(call: Call) -> Reply: return EXHAUSTED if "w40" in call.asked else _hollow(call) model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), list(tables), hooks=reports.hooks()) first = sum(1 for c in model.calls if "w40" in c.asked) assert first == 8 and reports.items[first - 1].part_columns == 10, "w40 一直用完额度，上限缩到 10 列" packed = [c for c in model.calls if set(c.asked) & {"n1", "n2", "n3", "n4"}] assert len(packed) == 8, "拼成的 40 列段按 10 列分开成 4 块，4 块一共 8 次" assert all(len(c.asked) == 1 and _width(c, tables) <= 10 for c in packed) assert {got[f"n{i}"].status for i in range(1, 5)} == {"failed"} assert "每段最多调用 8 次模型，这一段到了次数仍没写完，已停止重试" in {got[f"n{i}"].error for i in range(1, 5)}, \ "次数用完时还没问的块记「每段最多调用 8 次」" wide = [c for c in model.calls if "x40" in c.asked] assert len(wide) == 8 and _width(wide[0], tables) == 10, "单张 40 列的表按 10 列分开，各块一共 8 次" assert len(reports.items) == len(model.calls) class AlwaysSlow: """每次调用都超时（立刻抛 TimeoutError，和到了单次调用的时限走同一条路），fits(列数) 为真的照常写。 记下每次调用在第几轮：每次调用一样长（一个时限），一次调用开始时，已经结束的调用里最晚的那一轮的下一轮—— 轮数乘以时限就是用掉的时间。""" def __init__(self, tables: dict, fits: Callable[[int], bool] = lambda n: False) -> None: self.tables, self.fits, self.ok = tables, fits, answering(tables) self.rounds: list[int] = [] self._ended: list[int] = [] async def __call__(self, call: Call) -> Reply: r = 1 + max(self._ended, default=0) self.rounds.append(r) for _ in range(5): await asyncio.sleep(0) self._ended.append(r) if self.fits(_width(call, self.tables)): return self.ok(call) raise TimeoutError() async def test_a_model_that_always_times_out_is_bounded_by_the_part_cap(): """每次调用都超时：超时拆出来的两半同时问，最多拆到这一段的调用次数用完（8 次）或止损。以前 40→20→10→5 一层一层串着问， 8 次调用约 7 轮；现在两半同时问，至多 4 轮。""" tables = {"t40": _table(40)} slow = AlwaysSlow(tables) model, reports = Scripted(slow), Reports() got = await D.run_model_draft(model, _source(tables), ["t40"], hooks=reports.hooks()) assert len(model.calls) <= catalog.MODEL_PART_MAX_CALLS and max(slow.rounds) <= 4 assert [_width(c, tables) for c in model.calls[:3]] == [40, 20, 20], "两半同时问" assert set(reports.failures) == {"timeout"} and len(reports.items) == len(model.calls) o = got["t40"] assert o.status == "failed" and "模型 420 秒内没有写完这一段" + D.REASON_SPLIT in o.error assert o.chunks_done == o.chunks_total > 1, "拆出来的块都结束了" async def test_three_parts_that_always_time_out_stop_together(): """首轮同时起草的 3 段都一直超时：最近 10 次调用 7 次以上没写出来以后，拆出来还没问的块不再问，不再各自拆满 8 次。 以前（接连超时的止损）12 次调用；更早要 24 次。""" tables = {f"t{i}": _table(40) for i in range(1, 4)} slow = AlwaysSlow(tables) model = Scripted(slow) got = await D.run_model_draft(model, _source(tables), list(tables)) assert len(model.calls) <= 9 and max(slow.rounds) <= 4 assert {o.status for o in got.values()} == {"failed"} assert any(HALTED in o.error for o in got.values()) async def test_timeouts_that_go_away_at_ten_columns_are_still_split_down_and_written(): """超时和段大小有关（40、20 列超时，10 列写得完）：照旧拆到 10 列写出来，不因为前面超时多就停下。""" tables = {"t40": _table(40)} slow = AlwaysSlow(tables, fits=lambda n: n <= 10) model = Scripted(slow) got = await D.run_model_draft(model, _source(tables), ["t40"]) assert got["t40"].status == "ok" and len(got["t40"].notes["columns"]) == 40 assert len(model.calls) == 7 and max(slow.rounds) == 4 async def test_odd_values_in_a_reply_are_dropped_not_a_crash(): """模型把列名、表类型、业务主键、业务日期的列、度量类型写成了数组或对象：这几项丢掉，别的照收；不当成「起草这张表 时出错」。以前列名是数组时校验里抛 TypeError，整张表记出错，这次调用也没有记进用量。""" tables = {"t4": _table(4)} def handle(call: Call) -> Reply: return Reply(obj={"tables": [{ "table": "t4", "label": "某表", "kind": ["fact"], "keys": [["c001"]], "business_date": {"column": ["c002"], "rule": "按下单时间"}, "columns": [{"name": ["c001"], "label": "数组列名"}, {"name": {"x": 1}, "label": "对象列名"}, {"name": "c002", "label": "第二列", "measure": ["flow"]}, {"name": "c003", "label": "第三列", "measure": {"x": "flow"}}, {"name": "c004", "label": "第四列"}]}]}) model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), ["t4"], hooks=reports.hooks()) o = got["t4"] assert len(reports.items) == len(model.calls) == 2, "c001 没有按列名列出来：补问一次" assert o.status == "partial" and "c001" in o.error and "出错" not in o.error assert set(o.notes) == {"label", "columns"} and set(o.notes["columns"]) == {"c002", "c003", "c004"} assert all(set(items) == {"label"} for items in o.notes["columns"].values()) async def test_a_call_whose_reply_cannot_be_read_still_counts(monkeypatch): """读回复时出了意料之外的错：这张表记出错，这次调用照样记进用量（算一次出错），调用次数和用量对得上。""" tables = {"t4": _table(4)} def broken(entry, meta, *, system, **_): raise RuntimeError("校验时出了意外") monkeypatch.setattr(catalog, "_from_model", broken) model, reports = Scripted(answering(tables, usage=(100, 50))), Reports() got = await D.run_model_draft(model, _source(tables), ["t4"], hooks=reports.hooks()) assert len(model.calls) == 1 and reports.failures == ["error"] assert [(r.input_tokens, r.output_tokens) for r in reports.items] == [(100, 50)] assert got["t4"].status == "failed" and got["t4"].error == "起草这张表时出错：校验时出了意外" async def test_columns_the_catalog_already_names_are_not_counted_blank(): """模型这次留空、目录里已经有中文名或含义的列（不论来源：以前模型起草的、库注释、人工；驳回的不算）不算留空：照界面 的建议补了起草要求再起草，以前起草好的留着（重新起草不丢内容），这一列不该再说成「模型拿不准、留空了」。""" tables = {"t6": _table(6)} item = catalog.make_item existing = {"t6": {"columns": { "c002": {"label": item("以前起草的名字", "llm")}, "c003": {"meaning": item("库注释说清的含义", "comment")}, "c004": {"label": item("驳回的名字", "llm", "rejected")}, "c005": {"unit": item("元", "llm")}, }}} def handle(call: Call) -> Reply: cols = [{"name": "c001", "label": "第一列"}] + [{"name": f"c00{i}"} for i in range(2, 7)] return Reply(obj={"tables": [{"table": "t6", "label": "某表", "columns": cols}]}) got = await D.run_model_draft(Scripted(handle), _source(tables), ["t6"], existing=existing) assert got["t6"].status == "ok" and got["t6"].blank == ["c004", "c005", "c006"] async def test_a_wide_table_that_wrote_nothing_names_every_part(): """一项都没写出来的宽表同样逐段写上列的范围（以前只写原因和「这一段」「这几段」，看不出是哪几段）。""" def handle(call: Call) -> Reply: raise RuntimeError("网关拒绝") got = await D.run_model_draft(Scripted(handle), _source({"wide95": _table(95)}), ["wide95"]) o = got["wide95"] assert o.status == "failed" and o.error == (WIDE95_PARTS + "都没写出来：调用模型出错：网关拒绝；" "这张表接连 3 段都没写出来，剩下的 1 段已跳过，可以换一个起草模型后只重试这张表") narrow = await D.run_model_draft(Scripted(handle), _source({"t8": _table(8)}), ["t8"]) assert narrow["t8"].error == "调用模型出错：网关拒绝", "只有一段的表：这一段就是整张表，照旧只写原因" async def test_a_halted_part_that_could_be_split_is_not_asked_again_in_text_mode(monkeypatch): """止损以后每块只问一次：已经走纯文本的空壳不再原样再问、也不拆小。t1：结构化不支持 1 次（只管通信的重来，不进 止损的窗口）、纯文本 1 次、拆成两半各问一次再原样再问一次（5 次没写出来）；t2 的第一次是第 6 次，照常拆小，前一半 是第 7 次没写出来，之后止损：后一半不再问；t3 只问一次、不拆小。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {f"t{i}": _table(10) for i in range(1, 5)} model = Scripted(_text_only(_hollow)) got = await D.run_model_draft(model, _source(tables), list(tables)) assert _by_table(model) == {"t1": 6, "t2": 2, "t3": 1, "t4": 1} assert got["t2"].error == D.REASON_HOLLOW + D.REASON_SPLIT + HALTED assert got["t3"].error == catalog.MODEL_EMPTY_REASON + HALTED async def test_while_halted_structured_output_still_gets_one_probe_in_text(monkeypatch): """止损期间结构化输出只给空壳：这次起草还没换纯文本问过的，照旧换纯文本问一次（看看是不是结构化输出的问题）；纯文本 写出来了，后面都走纯文本。问过一次（纯文本也只给空壳）就不再问。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {**{f"x{i}": _table(4) for i in range(1, 8)}, "t8": _table(8), "t9": _table(8)} ok = answering(tables) def handle(text_works: bool) -> Callable[[Call], Reply]: def go(call: Call) -> Reply: if next(iter(call.asked)).startswith("x"): raise RuntimeError("网关拒绝了这次请求") if call.mode == "structured": return _hollow(call) return ok(call) if text_works else _hollow(call) return go model = Scripted(handle(True)) got = await D.run_model_draft(model, _source(tables), list(tables)) assert [c.mode for c in model.calls if "x1" not in c.asked and not next(iter(c.asked)).startswith("x")] == \ ["structured", "text", "text"] assert got["t8"].status == got["t9"].status == "ok" model = Scripted(handle(False)) got = await D.run_model_draft(model, _source(tables), list(tables)) assert _by_table(model)["t8"] == 2 and _by_table(model)["t9"] == 1, "纯文本问过一次，不再问" assert got["t9"].error == catalog.MODEL_EMPTY_REASON + HALTED # -------------------------------------------------------------------------- # 起草要求、用量 # -------------------------------------------------------------------------- async def test_guidance_goes_into_the_prompt_after_the_rules(): tables = {"t4": _table(4)} model = Scripted(answering(tables)) src = _source(tables, description="虚构的景区票务库") await D.run_model_draft(model, src, ["t4"], background=D.Background(guidance="金额一律按元写单位", note="这次重点看会员相关的列")) text = model.calls[0].text block = ("使用者补充的背景和要求：\n数据源说明：虚构的景区票务库\n起草要求：金额一律按元写单位\n" "本次说明：这次重点看会员相关的列") assert block in text rule = "它和这里的规矩冲突时，以这里的规矩为准" assert rule in catalog.CATALOG_DRAFT_SYSTEM assert text.index(rule) < text.index(block) < text.index("请为下面") < text.index("表 t4") plain = Scripted(answering(tables)) await D.run_model_draft(plain, _source(tables), ["t4"]) assert "使用者补充的背景和要求：" not in plain.calls[0].text, "都没有就不出现这一段（系统提示词里那条规矩照旧在）" only_note = Scripted(answering(tables)) await D.run_model_draft(only_note, _source(tables), ["t4"], background=D.Background(note="只看金额")) assert "使用者补充的背景和要求：\n本次说明：只看金额\n\n请为下面" in only_note.calls[0].text async def test_the_prompt_asks_not_to_draft_every_column_while_thinking(): assert "思考时不必把要输出的内容逐列草拟一遍" in catalog.CATALOG_DRAFT_SYSTEM async def test_usage_is_added_up_through_the_hooks(): tables = {"orders_x": _table(24)} ok = answering(tables, usage=(120, 60)) def handle(call: Call) -> Reply: if call.n == 0: return Reply(text="", finish="length", thinking="……", usage=(100, 32768)) return ok(call) reports = Reports() await D.run_model_draft(Scripted(handle), _source(tables), ["orders_x"], hooks=reports.hooks()) assert [(r.input_tokens, r.output_tokens, r.failure) for r in reports.items] == \ [(100, 32768, "exhausted"), (120, 60, None), (120, 60, None)] assert {r.notice for r in reports.items} == {None} and {r.part_columns for r in reports.items} == {40} def test_budget_words_match_the_other_drafts(): from app.data import recipe_ai assert D.BUDGET_WORDS == recipe_ai._BUDGET_WORDS # -------------------------------------------------------------------------- # 从正文里取 JSON # -------------------------------------------------------------------------- @pytest.mark.parametrize("text, tables, truncated", [ # 前面有一句话、带括号的说明、围栏 ('按 {table} 的格式写好了：\n```json\n{"tables": [{"table": "a", "label": "甲"}]}\n```', ["a"], False), ('[说明] 草稿：{"tables": [{"table": "a"}]} 以上', ["a"], False), # 外面多包了一层 ('{"result": {"tables": [{"table": "a"}, {"table": "b"}]}}', ["a", "b"], False), # 截断在字符串、数字、字面量、键名中间：只收写完了的 ('{"tables": [{"table": "a", "label": "甲"}, {"table": "b", "label": "乙', ["a", "b"], True), ('{"tables": [{"table": "a", "columns": [{"name": "x", "label": "某"}, {"name": "y", "len": 12', ["a"], True), ('{"tables": [{"table": "a", "keys": ["x", "y"', ["a"], True), ('{"tables": [{"table": "a"}, {"tab', ["a"], True), ('{"tables": [{"table": "a", "flag": tr', ["a"], True), ]) def test_json_is_taken_from_the_text(text, tables, truncated): reply = D._parse_text(text, cut=False) assert reply is not None and [e["table"] for e in reply.entries] == tables assert reply.truncated is truncated def test_cut_values_are_dropped_and_complete_ones_kept(): text = '{"tables": [{"table": "a", "label": "甲", "keys": ["x", "y"], "columns": [{"name": "x", "label": "某"}, ' \ '{"name": "y", "label": "半' entry = D._parse_text(text, cut=True).entries[0] assert entry == {"table": "a", "label": "甲", "keys": ["x", "y"], "columns": [{"name": "x", "label": "某"}]} cut_keys = D._parse_text('{"tables": [{"table": "a", "label": "甲", "keys": ["x", "y"', cut=True).entries[0] assert cut_keys == {"table": "a", "label": "甲"}, "没写完的业务主键整项不要，不能只留前半" @pytest.mark.parametrize("text", ["好的，我来写", "{坏的 json}", ""]) def test_not_json_is_not_a_reply(text): assert D._parse_text(text, cut=False) is None # -------------------------------------------------------------------------- # 闭环第一步的起草遗留：逐段交出去落盘、块进度、「已拆小重试」只说拆过的段 # -------------------------------------------------------------------------- class Written: """进度回调 part_written 收到的：(表名, 已写出内容的段数, 这时这张表拿到的列)；chunks_changed 收到的块进度。""" def __init__(self) -> None: self.parts: list[tuple[str, int, set[str], bool]] = [] self.chunks: dict[str, list[tuple[int, int]]] = {} async def part_written(self, table: str, written: D.PartWrite) -> None: cols = set((written.notes.get("columns") or {})) self.parts.append((table, written.parts_written, cols, "label" in written.notes)) async def chunks_changed(self, table: str, total: int, done: int) -> None: self.chunks.setdefault(table, []).append((total, done)) def hooks(self) -> D.Hooks: return D.Hooks(part_written=self.part_written, chunks_changed=self.chunks_changed) async def test_each_part_of_a_wide_table_is_handed_over_as_soon_as_it_is_written(monkeypatch): """宽表每写出一段就交给调用方并入目录（最后一段由 table_done 收尾）：交出去的是这张表到这时拿到的全部内容。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) tables = {"wide95": _table(95)} seen = Written() got = await D.run_model_draft(Scripted(answering(tables)), _source(tables), ["wide95"], hooks=seen.hooks()) o = got["wide95"] assert o.status == "ok" and o.parts_written == 4 assert [(t, n, len(cols), label) for t, n, cols, label in seen.parts] == [ ("wide95", 1, 0, True), ("wide95", 2, 32, True), ("wide95", 3, 64, True)] def handle(call: Call) -> Reply: part, cols = call.asked["wide95"] if part == "columns" and "c033" in cols: raise RuntimeError("网关拒绝") return answering(tables)(call) seen = Written() got = await D.run_model_draft(Scripted(handle), _source(tables), ["wide95"], hooks=seen.hooks()) assert got["wide95"].status == "partial" and got["wide95"].parts_written == 3 assert [n for _, n, _, _ in seen.parts] == [1, 2], "没写出来的那段不交" async def test_a_narrow_table_is_handed_over_only_once_at_the_end(): tables = {"t8": _table(8)} seen = Written() got = await D.run_model_draft(Scripted(answering(tables)), _source(tables), ["t8"], hooks=seen.hooks()) assert seen.parts == [] and got["t8"].parts_written == 1 failed = await D.run_model_draft(Scripted(lambda call: EXHAUSTED), _source({"t4": _table(4)}), ["t4"]) assert failed["t4"].parts_written == 0 async def test_chunks_count_what_the_table_is_actually_asked_in(monkeypatch): """块进度：一开始是这张表分的段数；一段被拆成两半多一块、被缩小的上限分开成几块就多几块；一块结束算一块。 多表拼成的段按表拆开时，每张表还是一块。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) monkeypatch.setattr(catalog, "MODEL_SHRINK_AFTER", 1) tables = {"a12": _table(12), "b40": _table(40)} ok = answering(tables) seen = Written() got = await D.run_model_draft(Scripted(lambda call: EXHAUSTED if call.n == 0 else ok(call)), _source(tables), ["a12", "b40"], hooks=seen.hooks()) a, b = got["a12"], got["b40"] assert a.status == b.status == "ok" assert (a.chunks_total, a.chunks_done) == (2, 2), "额度用完拆成两半" assert (b.chunks_total, b.chunks_done) == (2, 2), "上限缩到 20 列，40 列的段分开成两块" assert seen.chunks["a12"][-1] == (2, 2) and seen.chunks["b40"][-1] == (2, 2) assert all(done <= total for pairs in seen.chunks.values() for total, done in pairs) wide = {"wide95": _table(95)} got = await D.run_model_draft(Scripted(answering(wide)), _source(wide), ["wide95"]) assert (got["wide95"].chunks_total, got["wide95"].chunks_done) == (4, 4) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 4) packed = {"p6": _table(6), "q6": _table(6)} okp = answering(packed) model = Scripted(lambda call: _hollow(call) if call.n == 0 else okp(call)) got = await D.run_model_draft(model, _source(packed), ["p6", "q6"]) assert [set(c.asked) for c in model.calls] == [{"p6", "q6"}, {"p6"}, {"q6"}], "按表拆开" assert {(o.chunks_total, o.chunks_done) for o in got.values()} == {(1, 1)} async def test_only_the_parts_that_were_split_say_so(monkeypatch): """一张表几段没写出来、只有一部分真拆小重试过：「已拆小重试仍没写出」只跟在拆过的那几段后面说。95 列：第 3 段一直用完 额度（拆小重试过），第 4 段调用出错（没拆过）。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_SHRINK_AFTER", 99) monkeypatch.setattr(catalog, "MODEL_HALT_ON", 99) tables = {"wide95": _table(95)} ok = answering(tables) def handle(call: Call) -> Reply: part, cols = call.asked["wide95"] if part == "columns" and cols[0] >= "c065": raise RuntimeError("网关拒绝") if part == "columns" and cols[0] >= "c033": return EXHAUSTED return ok(call) got = await D.run_model_draft(Scripted(handle), _source(tables), ["wide95"]) o = got["wide95"] assert o.status == "partial" assert o.error == ("这张表 95 列，分 4 段起草，第 3 段（c033 到 c064，共 32 列）、第 4 段（c065 到 c095，共 31 列）没写出来：" + _exhausted() + "；调用模型出错：网关拒绝；其中第 3 段已拆小重试仍没写出。可以只重试这张表") # -------------------------------------------------------------------------- # 复核后的修正：止损的窗口只记每次调用的最终结局。通信层的重来（服务方嫌额度太大退回 8192、结构化输出不支持或回了 # 说明文字改走纯文本、429 / 5xx 退避重试）不花额度，也说明不了模型写不出来：照旧报给进度回调（用量、提醒照算）， # 不进止损的窗口。以前都记成没写出来，OpenAI 兼容网关（额度上限 8192、不支持 json_schema）上并发的头几次调用全是 # 重来，第一次额度用完就凑满 10 次里 7 次，那一段直接记没写出来、不再拆小。 # -------------------------------------------------------------------------- def _six_tables() -> dict: """6 张 40 列的表：每张自成一段，默认并发 3 段同时起草。""" return {f"t{i}": _table(40, prefix=f"t{i}_") for i in range(6)} def _halved(model: Scripted, table: str) -> bool: """这张表被拆成两半问过（表级的项加前 20 列那一半）。""" return any(part == "table+columns" and len(cols or []) == 20 for call in model.calls for name, (part, cols) in call.asked.items() if name == table) async def test_budget_and_structured_output_fallbacks_do_not_halt_a_concurrent_draft(): tables = _six_tables() ok = answering(tables) spent: set[str] = set() async def refuse(call: Call) -> Reply: await asyncio.sleep(0.01) raise _Status(400, "max_tokens: 32768 > 8192, which is the maximum allowed number of output tokens") async def small(call: Call) -> Reply: await asyncio.sleep(0.01) if call.mode == "structured": raise _Status(400, "This response_format type is unavailable now") if "t0" in call.asked and "t0" not in spent: spent.add("t0") return EXHAUSTED return ok(call) big, sm, reports = Scripted(refuse), Scripted(small), Reports() got = await D.run_model_draft(catalog.DraftModel(big, sm, "m"), _source(tables), list(tables), hooks=reports.hooks()) assert _status(got) == dict.fromkeys(tables, "ok"), {t: o.error for t, o in got.items()} assert _halved(sm, "t0"), "额度用完的那一段照常拆小重试" assert len(big.calls) == 3 and reports.failures.count("error") == 6, "退回小额度、改走纯文本的重来照旧报给进度回调" assert reports.failures.count("exhausted") == 1 async def test_rate_limit_retries_do_not_halt_the_draft(): tables = _six_tables() ok = answering(tables) limited: dict[str, int] = {} spent: set[str] = set() async def handle(call: Call) -> Reply: await asyncio.sleep(0.01) name = next(iter(call.asked)) if name in ("t0", "t1", "t2") and limited.get(name, 0) < 2: limited[name] = limited.get(name, 0) + 1 raise _Status(429, "Too Many Requests") if name == "t0" and not spent: spent.add(name) return EXHAUSTED return ok(call) model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), list(tables), hooks=reports.hooks()) assert _status(got) == dict.fromkeys(tables, "ok"), {t: o.error for t, o in got.items()} assert _halved(model, "t0"), "额度用完的那一段照常拆小重试" assert reports.failures.count("error") == 6, "429 之后重试成功的那几次照旧报给进度回调" async def test_structured_output_that_comes_back_as_prose_does_not_count_toward_halting(monkeypatch): """结构化输出回了说明文字、这一块改走纯文本再问：也是通信方式的问题，那一次不进止损的窗口（止损门槛调低到 3 次， 并发的 3 块都回说明文字也不止损）；纯文本写出来的、没写出来的照常算。""" monkeypatch.setattr(catalog, "MODEL_HALT_ON", 3) tables = _six_tables() ok = answering(tables) spent: set[str] = set() async def handle(call: Call) -> Reply: await asyncio.sleep(0.01) if call.mode == "structured": return Reply(text="好的。这张表记录的是一笔笔业务，各列的含义如下：第一列是编号……") if "t0" in call.asked and "t0" not in spent: spent.add("t0") return EXHAUSTED return ok(call) model, reports = Scripted(handle), Reports() got = await D.run_model_draft(model, _source(tables), list(tables), hooks=reports.hooks()) assert _status(got) == dict.fromkeys(tables, "ok"), {t: o.error for t, o in got.items()} assert _halved(model, "t0") assert reports.failures.count("error") == 3, "回了说明文字的那几次照旧报给进度回调" async def test_a_call_that_still_fails_after_its_retries_counts_toward_halting(monkeypatch): """重试用完仍然出错的那一次是这次调用的最终结局，照常进止损的窗口：一直 503 的网关，止损照样起作用。""" monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1) monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1) tables = {f"t{i}": _table(10) for i in range(1, 9)} ok = answering(tables) def handle(call: Call) -> Reply: if set(call.asked) & {"t1", "t2", "t3", "t4", "t5", "t6", "t7"}: raise _Status(503, "Service Unavailable") return EXHAUSTED if "t8" in call.asked else ok(call) got = await D.run_model_draft(Scripted(handle), _source(tables), list(tables)) assert got["t7"].error == "调用模型出错：Service Unavailable（已重试 2 次）" assert got["t8"].error == _exhausted() + HALTED, "最近 10 次的最终结局有 7 次以上没写出来：t8 不再拆小"