mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 02:11:58 +00:00
数据目录模型起草换成分段、按失败原因处理的引擎(data/catalog_draft.py):每次最多 4 张表、40 列,窄表按顺序拼段,宽表拆成「表级一段 + 每 40 列一段」,列分段带着表注释和表级起草结果做背景,最后合并照旧逐项校验;思考和输出用完额度、只给空壳时拆成两半再问(最少 5 列、最多 3 层),拆不开的空壳换纯文本问一次,写出来了后面都走纯文本;被截断的 JSON 先捡回已完整的表和列、只补剩下的;缺了的表和列只补缺的、补一次;服务方嫌额度大时退回 8192、结构化输出不支持时改走纯文本,都管到这次起草结束;网络错误、429、5xx、超时退避重试最多 2 次;单次调用 420 秒;全进程最多 3 个调用同时进行;报错写原因和办法(第几段、共几列没写出来、可以只重试这张表);提示词加上起草要求一段(数据源说明、起草要求、本次说明,规矩在前、冲突以规矩为准)和「思考时不必逐列草拟」;同步起草改用新引擎,写入抽成每张表一结束就能调的 write_draft
This commit is contained in:
+158
-199
@@ -1130,10 +1130,14 @@ async def save_patch(session: AsyncSession, source_id: str, table: str, changes:
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TableDraft:
|
||||
"""一张表的起草结果。covered:这一轮完整算过的来源(merge_notes 据此删掉这些来源里已经没有的旧项)。"""
|
||||
"""一张表的起草结果。covered:这一轮完整算过的来源(merge_notes 据此删掉这些来源里已经没有的旧项)。
|
||||
|
||||
error:模型起草只写出了一部分(有的段没写出来)时,没写出来的是哪部分、为什么。这时 covered 不含 llm:没写全
|
||||
就不算完整起草过,已有的模型起草项不会因此被删。"""
|
||||
|
||||
notes: dict[str, Any]
|
||||
covered: frozenset[str] = frozenset()
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def system_notes_source(source: Any) -> bool:
|
||||
@@ -1446,10 +1450,25 @@ status 状态;attribute 属性)。
|
||||
- 只写数据事实,不写计算公式,不写指标怎么算。
|
||||
- 拿不准的字段留空,不要猜。列名、表名照原样写,不要翻译或改写。
|
||||
- 标着「系统说明」的内容是导入时按核对结果生成的,以它为准:不要重复,也不要与之矛盾。
|
||||
- 标着「已确认」的目录项由人工确认过,以它为准。"""
|
||||
- 标着「已确认」的目录项由人工确认过,以它为准。
|
||||
- 表名下一行写着「这次只写……」的表,只写它说的那部分:只写表级的项时不写列,只写列出的几列时不写表级的项、\
|
||||
也不写别的列。
|
||||
- 使用者补充的背景和要求帮你理解业务、统一用词;它和这里的规矩冲突时,以这里的规矩为准。
|
||||
- 思考时不必把要输出的内容逐列草拟一遍,想清楚后直接按格式输出:思考和输出共用同一份长度额度,额度用完时还没输出的\
|
||||
内容就全丢了。"""
|
||||
|
||||
CATALOG_DRAFT_PROMPT = "请为下面 {n} 张表编写数据目录,每张表输出一项,table 写表名。\n\n{tables}"
|
||||
|
||||
#: 使用者补充的背景和要求(数据源说明、起草要求、本次说明),放在表结构前面单独一段;都没有就不出现这一段
|
||||
CATALOG_DRAFT_BACKGROUND_PROMPT = "使用者补充的背景和要求:\n{lines}"
|
||||
|
||||
#: 纯文本那条路(结构化输出用不了时):没有 JSON Schema 约束,要把格式说清
|
||||
CATALOG_DRAFT_TEXT_PROMPT = """只输出一个 JSON 对象,不要输出别的文字,格式如下(拿不准的字段不写):
|
||||
{"tables": [{"table": "表名", "label": "中文名", "grain": "一行代表什么", "kind": "fact", "keys": ["列名"], \
|
||||
"business_date": {"column": "列名", "rule": "一句规则", "timezone": "时区"}, \
|
||||
"columns": [{"name": "列名", "label": "中文名", "meaning": "含义", "unit": "单位", "measure": "flow"}]}]}
|
||||
kind 取 fact、dimension、snapshot、log、config 之一;measure 取 flow、stock、ratio、identifier、status、attribute 之一。"""
|
||||
|
||||
#: 拼表结构描述用的固定片段(写给模型)
|
||||
CATALOG_DRAFT_PROMPT_LABELS = {
|
||||
"table": "表 {name}",
|
||||
@@ -1465,6 +1484,16 @@ CATALOG_DRAFT_PROMPT_LABELS = {
|
||||
"col_system": "系统说明:{text}",
|
||||
"unique": "唯一约束:{cols}",
|
||||
"confirmed": "已确认的目录项:",
|
||||
# 分段起草:表名下一行说这次只写哪部分(宽表的表级一段、列分段,拆小重问时的半段)
|
||||
"only_table": "这次只写表级的项(中文名、粒度、表类型、业务主键、业务日期),不写列",
|
||||
"only_columns": "这次只写下面列出的 {n} 列(这张表共 {total} 列),不写表级的项",
|
||||
"table_and_columns": "这次只写表级的项和这 {n} 列:{cols}",
|
||||
"brief_columns": "列(共 {n} 列,这里只列名字和类型):",
|
||||
"drafted": "已起草的表级内容:",
|
||||
# 使用者补充的背景和要求
|
||||
"description": "数据源说明:{text}",
|
||||
"guidance": "起草要求:{text}",
|
||||
"note": "本次说明:{text}",
|
||||
}
|
||||
|
||||
CATALOG_DRAFT_PROMPT_SCHEMA: dict[str, Any] = {
|
||||
@@ -1509,9 +1538,6 @@ CATALOG_DRAFT_PROMPT_SCHEMA: dict[str, Any] = {
|
||||
"required": ["tables"],
|
||||
}
|
||||
|
||||
_P = CATALOG_DRAFT_PROMPT_LABELS
|
||||
|
||||
|
||||
class DraftModel:
|
||||
"""起草用的模型:先按 MODEL_MAX_TOKENS 的额度调;服务方嫌额度太大(budget_rejected)时换成 MODEL_FALLBACK_MAX_TOKENS
|
||||
的那一个,之后的调用都用它(shrink)。一次起草(一个起草任务、一次同步起草)共用一个,退回之后不再先失败一次。
|
||||
@@ -1572,58 +1598,6 @@ def _confirmed_only(notes: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
return _assemble({k: v for k, v in _slots(notes).items() if v.get("status") in ("confirmed", "verified")})
|
||||
|
||||
|
||||
def _describe_for_model(name: str, meta: Mapping[str, Any], *, system: bool, existing: Mapping[str, Any] | None) -> str:
|
||||
"""一张表交给模型的结构描述:列、类型、主键、外键、唯一约束、注释;不含任何数据值。"""
|
||||
lines = [(_P["view"] if meta.get("is_view") else _P["table"]).format(name=name)]
|
||||
comment = (meta.get("comment") or "").strip()
|
||||
if comment:
|
||||
lines.append((_P["system"] if system else _P["comment"]).format(text=comment))
|
||||
pk = set(meta.get("primary_key") or [])
|
||||
fk_of: dict[str, str] = {}
|
||||
for fk in meta.get("foreign_keys") or []:
|
||||
if isinstance(fk, dict) and len(fk.get("columns") or []) == 1:
|
||||
to_cols = fk.get("to_columns") or []
|
||||
fk_of[fk["columns"][0]] = f"{fk.get('to_table')}.{to_cols[0]}" if to_cols else str(fk.get("to_table"))
|
||||
lines.append(_P["columns"])
|
||||
for col in meta.get("columns") or []:
|
||||
marks = []
|
||||
if col["name"] in pk:
|
||||
marks.append(_P["pk"])
|
||||
if not col.get("nullable", True):
|
||||
marks.append(_P["not_null"])
|
||||
if col["name"] in fk_of:
|
||||
marks.append(_P["fk"].format(target=fk_of[col["name"]]))
|
||||
if text := (col.get("comment") or "").strip():
|
||||
marks.append((_P["col_system"] if system else _P["col_comment"]).format(text=text))
|
||||
line = _P["column"].format(name=col["name"], type=col.get("type") or "")
|
||||
lines.append(line + ("," + ",".join(marks) if marks else ""))
|
||||
for group in meta.get("unique") or []:
|
||||
if isinstance(group, list) and group:
|
||||
lines.append(_P["unique"].format(cols="、".join(group)))
|
||||
confirmed = _notes_lines(_confirmed_only(existing), meta=meta, system_notes=False)
|
||||
if confirmed:
|
||||
lines.append(_P["confirmed"])
|
||||
lines += confirmed
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _batches(names: list[str], tables: Mapping[str, Any]) -> list[list[str]]:
|
||||
"""按表分批:每批不超过 MODEL_BATCH_TABLES 张、合计不超过 MODEL_BATCH_COLUMNS 列(一张表超了就自成一批)。"""
|
||||
out: list[list[str]] = []
|
||||
batch: list[str] = []
|
||||
width = 0
|
||||
for name in names:
|
||||
n = len((tables.get(name) or {}).get("columns") or [])
|
||||
if batch and (len(batch) >= MODEL_BATCH_TABLES or width + n > MODEL_BATCH_COLUMNS):
|
||||
out.append(batch)
|
||||
batch, width = [], 0
|
||||
batch.append(name)
|
||||
width += n
|
||||
if batch:
|
||||
out.append(batch)
|
||||
return out
|
||||
|
||||
|
||||
#: 模型给了这张表、但一项可用内容都没有(只有表名的空壳,或者给的值全不合规):这张表记模型失败,不算起草成功。
|
||||
#: 不然界面只看到「新增 0 项」,会说成「没有变化」
|
||||
MODEL_EMPTY_REASON = "模型没有给出这张表的可用内容"
|
||||
@@ -1653,39 +1627,6 @@ def _model_tables(raw: Any) -> list[dict[str, Any]] | None:
|
||||
return [t for t in raw if isinstance(t, dict)]
|
||||
|
||||
|
||||
async def _ask_model(model: Any, messages: list[Any],
|
||||
usable: Callable[[list[dict[str, Any]]], bool] | None = None) -> list[dict[str, Any]]:
|
||||
"""问一批。先走结构化输出;模型不支持、调用报错、给回来的不是约定的结构,或者 usable 说它不够用(实测见过
|
||||
结构化输出只给 {"tables":[{"table":"x"}]} 这样的空壳),再退回纯文本、从正文里取 JSON。
|
||||
|
||||
结构化输出报错也退回(同 judge._ask):有的网关不支持工具调用,一调就报错,纯文本却能用。代价是真的
|
||||
网络故障会多试一次,两次加起来仍受 MODEL_TIMEOUT_S 约束。
|
||||
|
||||
两条路都拿到了东西时,结构化的在前、纯文本的在后一起交回(同一张表各有一份,调用方挑能用的那份);纯文本那次
|
||||
失败了就只交回结构化的那份(调用方据此给空壳的表记失败)。
|
||||
"""
|
||||
try:
|
||||
parsed = _model_tables(await model.with_structured_output(CATALOG_DRAFT_PROMPT_SCHEMA).ainvoke(messages))
|
||||
except Exception: # noqa: BLE001 - 退回纯文本
|
||||
parsed = None
|
||||
if parsed is not None and (usable is None or usable(parsed)):
|
||||
return parsed
|
||||
from app.engine.state import message_text
|
||||
|
||||
try:
|
||||
reply = await model.ainvoke(messages)
|
||||
except Exception: # noqa: BLE001 - 结构化那份还在:交回它,空壳的表由调用方记失败
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
raise
|
||||
text = _model_tables(message_text(reply))
|
||||
if text is None:
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
raise ValueError("模型的回复不是约定的格式")
|
||||
return [*(parsed or []), *text]
|
||||
|
||||
|
||||
def _clean_text(value: Any, limit: int = _TEXT_MAX) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
@@ -1740,65 +1681,30 @@ def _from_model(entry: Mapping[str, Any], meta: Mapping[str, Any], *, system: bo
|
||||
|
||||
|
||||
async def draft_with_model(model: Any, source: Any, tables: list[str], *,
|
||||
existing: Mapping[str, Mapping[str, Any]] | None = None) -> dict[str, TableDraft | str]:
|
||||
existing: Mapping[str, Mapping[str, Any]] | None = None,
|
||||
guidance: str | None = None, note: str | None = None) -> dict[str, TableDraft | str]:
|
||||
"""用模型起草这些表的中文名、粒度、表类型、业务主键、业务日期和每列的中文名、含义、单位、度量类型。
|
||||
|
||||
model 要有 with_structured_output(schema) 或 ainvoke(messages)(测试注入假模型)。按表分批(_batches),
|
||||
最多 MODEL_CONCURRENCY 批同时发。返回 {表名: TableDraft 或失败原因}:一批失败只影响这一批的表,
|
||||
其余照常;失败的表不算 covered,已有的模型起草项不会因此被删。
|
||||
实现在 data/catalog_draft.py(分段、按失败原因处理、捡回被截断的部分);这里是给老调用方的形状:返回
|
||||
{表名: TableDraft 或失败原因}。全部写出来的 covered 是 {"llm"};只写出一部分的仍是 TableDraft,covered 为空、
|
||||
error 写明哪部分没写出来;一项都没写出来的是失败原因。表结构里没有的表不在结果里。
|
||||
|
||||
交给模型的只有表结构(列、类型、主键、外键、唯一约束、注释),不发任何数据值;existing 里已确认、
|
||||
有确证的项作为「以此为准」的背景一起交过去。
|
||||
model 要有 with_structured_output(schema) 或 ainvoke(messages)(测试注入假模型)。交给模型的只有表结构和使用者
|
||||
补充的背景(数据源说明、guidance 起草要求、note 本次说明),不发任何数据值;existing 里已确认、有确证的项作为
|
||||
「以此为准」的背景一起交过去。
|
||||
"""
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from app.data import catalog_draft
|
||||
|
||||
all_tables = (getattr(source, "schema_cache", None) or {}).get("tables") or {}
|
||||
system = system_notes_source(source)
|
||||
names = [t for t in tables if isinstance(all_tables.get(t), dict)]
|
||||
results: dict[str, TableDraft | str] = {}
|
||||
gate = asyncio.Semaphore(MODEL_CONCURRENCY)
|
||||
|
||||
async def run(batch: list[str]) -> None:
|
||||
described = "\n\n".join(_describe_for_model(n, all_tables[n], system=system,
|
||||
existing=(existing or {}).get(n)) for n in batch)
|
||||
messages = [SystemMessage(content=CATALOG_DRAFT_SYSTEM),
|
||||
HumanMessage(content=CATALOG_DRAFT_PROMPT.format(n=len(batch), tables=described))]
|
||||
def drafted(entries: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
||||
"""这批里每张表能用的那份:同一张表给了几份(结构化、纯文本各一份)时取第一份有可用内容的。"""
|
||||
out: dict[str, dict[str, Any]] = {}
|
||||
for entry in entries:
|
||||
name = str(entry.get("table"))
|
||||
if name in batch and name not in out:
|
||||
notes = _from_model(entry, all_tables[name], system=system)
|
||||
if notes:
|
||||
out[name] = notes
|
||||
return out
|
||||
|
||||
try:
|
||||
async with gate:
|
||||
got = await asyncio.wait_for(
|
||||
_ask_model(model, messages, usable=lambda entries: len(drafted(entries)) == len(batch)),
|
||||
timeout=MODEL_TIMEOUT_S)
|
||||
except asyncio.TimeoutError:
|
||||
for n in batch:
|
||||
results[n] = f"模型起草超时:{MODEL_TIMEOUT_S} 秒内没有收到回复"
|
||||
return
|
||||
except Exception as e: # noqa: BLE001 - 一批失败不影响其他批和其他来源
|
||||
from app.core.errors import describe_exception, first_line
|
||||
|
||||
for n in batch:
|
||||
results[n] = f"模型起草失败:{first_line(e) or describe_exception(e)}"
|
||||
return
|
||||
named = {str(t.get("table")) for t in got}
|
||||
usable = drafted(got)
|
||||
for n in batch:
|
||||
if n in usable:
|
||||
results[n] = TableDraft(notes=usable[n], covered=frozenset({"llm"}))
|
||||
else:
|
||||
results[n] = MODEL_EMPTY_REASON if n in named else "模型的回复里没有这张表"
|
||||
|
||||
await asyncio.gather(*(run(b) for b in _batches(names, all_tables)))
|
||||
return {n: results[n] for n in names if n in results}
|
||||
outcomes = await catalog_draft.run_model_draft(
|
||||
model, source, tables, existing=existing, background=catalog_draft.Background(guidance=guidance, note=note))
|
||||
out: dict[str, TableDraft | str] = {}
|
||||
for name, outcome in outcomes.items():
|
||||
if outcome.status == "failed":
|
||||
out[name] = outcome.error or MODEL_EMPTY_REASON
|
||||
else:
|
||||
out[name] = TableDraft(notes=outcome.notes, covered=frozenset({"llm"}) if outcome.status == "ok"
|
||||
else frozenset(), error=outcome.error)
|
||||
return out
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
@@ -1808,7 +1714,8 @@ async def draft_with_model(model: Any, source: Any, tables: list[str], *,
|
||||
|
||||
@dataclass
|
||||
class TableDraftResult:
|
||||
"""一张表这次起草的结果。error:这张表没起草(表结构里没有、写入一直冲突);model_error:只是模型那部分失败。"""
|
||||
"""一张表这次起草的结果。error:这张表没起草(表结构里没有、写入一直冲突);model_error:只是模型那部分失败,
|
||||
或者只写出了一部分(哪部分没写出来、为什么)。"""
|
||||
|
||||
table: str
|
||||
added: int = 0
|
||||
@@ -1817,9 +1724,14 @@ class TableDraftResult:
|
||||
version: int = 0
|
||||
error: str | None = None
|
||||
model_error: str | None = None
|
||||
#: 模型给了这张表几项可用内容(并入目录之前,按项数)。没用模型、模型这部分失败时为 None。
|
||||
#: 模型给了这张表几项可用内容(并入目录之前,按项数)。没用模型、模型这部分一项都没写出来时为 None。
|
||||
#: 大于 0 而新增、更新、删除都是 0,才是「模型给了内容、只是和现有目录一致」
|
||||
model_items: int | None = None
|
||||
#: ok 写入了(用模型时各段都写出来了)、partial 模型只写出一部分、failed 没写入或模型一段都没写出来、
|
||||
#: skipped 表结构里没有这张表
|
||||
status: str = "ok"
|
||||
#: 模型分几段起草(不用模型时为 None)
|
||||
parts_total: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -1834,35 +1746,102 @@ class DraftReport:
|
||||
model_error: str | None = None
|
||||
|
||||
|
||||
async def draft_catalog(session: AsyncSession, source: Any, *, tables: list[str] | None = None,
|
||||
use_model: bool = False, model: Any = None, actor: str | None = None) -> DraftReport:
|
||||
"""起草这些表的目录并写入,同步完成。
|
||||
#: 表结构里没有要起草的表(被删了、改名了,或者名字写错了)
|
||||
def missing_table_error(name: str) -> str:
|
||||
return f"表结构里没有 {name},请先重新探查结构"
|
||||
|
||||
|
||||
#: 写入一直撞上别人刚改过(乐观锁重读重并三次仍冲突)
|
||||
WRITE_CONFLICT_ERROR = "这张表的数据目录正被其他人修改,请稍后再起草"
|
||||
|
||||
|
||||
async def write_draft(session: AsyncSession, source: Any, table: str, modeled: Any = None, *, at: str,
|
||||
actor: str | None) -> TableDraftResult:
|
||||
"""一张表的起草并入目录并写入:数据库注释、外键约束、命名推断照做,modeled(catalog_draft.ModelOutcome,
|
||||
没用模型时为 None)有可用内容的再并上。同步起草和后台起草任务(每张表一结束就调)共用这一处。
|
||||
|
||||
- source:手工源的 DataSource,或上传源绑定当前快照的视图(table_versions.resolve_source);按 source.id 存。
|
||||
- tables 为空时按使用次数(table_usage)取前 DRAFT_DEFAULT_TABLES 张,次数相同按表结构里的顺序。
|
||||
- 每张表:数据库注释、外键约束、命名推断照做;use_model 时再加模型起草。model 可注入(测试用假模型),
|
||||
不给就按设置解析(resolve_draft_model);模型用不了、某一批失败,都不影响其余来源照常写入。
|
||||
- 并入已有目录按 merge_notes 的规则:人工确认、驳回过的项不动。内容没变不写、不升版本。
|
||||
- 模型只写出一部分的,模型那份不算完整起草过(covered 不含 llm),已有的模型起草项不会因此被删。
|
||||
- 写入撞上别人刚改过(乐观锁)就重读重并,最多三次。
|
||||
"""
|
||||
cache = getattr(source, "schema_cache", None) or {}
|
||||
all_tables = cache.get("tables") or {}
|
||||
report = DraftReport()
|
||||
if tables:
|
||||
wanted = list(dict.fromkeys(str(t) for t in tables))
|
||||
result = TableDraftResult(table=table)
|
||||
structure = draft_structure(source, table)
|
||||
model_draft: TableDraft | None = None
|
||||
if modeled is not None:
|
||||
result.parts_total = modeled.parts_total
|
||||
result.model_error = modeled.error
|
||||
if modeled.items:
|
||||
result.model_items = modeled.items
|
||||
model_draft = TableDraft(notes=modeled.notes, covered=frozenset({"llm"}) if modeled.status == "ok"
|
||||
else frozenset())
|
||||
for _attempt in range(3):
|
||||
entry = await read_entry(session, source.id, table)
|
||||
version = entry.version if entry else 0
|
||||
merged, stats = merge_notes(entry.notes if entry else {}, structure.notes, covered=structure.covered, at=at)
|
||||
if model_draft is not None:
|
||||
merged, more = merge_notes(merged, model_draft.notes, covered=model_draft.covered, at=at)
|
||||
stats += more
|
||||
result.added, result.updated, result.removed = stats.added, stats.updated, stats.removed
|
||||
result.version = version
|
||||
if entry is None and not merged:
|
||||
break # 什么也没起草出来,不建空目录
|
||||
try:
|
||||
written = await write_entry(session, source.id, table, merged, if_version=version, actor=actor)
|
||||
except CatalogConflict:
|
||||
continue
|
||||
result.version = written.version
|
||||
break
|
||||
else:
|
||||
usage = await table_usage(session, source)
|
||||
order = {name: i for i, name in enumerate(all_tables)}
|
||||
wanted = sorted(all_tables, key=lambda n: (-usage.get(n, 0), order[n]))[:DRAFT_DEFAULT_TABLES]
|
||||
result.error = WRITE_CONFLICT_ERROR
|
||||
result.added = result.updated = result.removed = 0
|
||||
if result.error or (modeled is not None and modeled.status == "failed"):
|
||||
result.status = "failed"
|
||||
elif modeled is not None and modeled.status == "partial":
|
||||
result.status = "partial"
|
||||
return result
|
||||
|
||||
|
||||
def pick_tables(cache: Mapping[str, Any], wanted: Iterable[str]) -> tuple[list[str], list[str]]:
|
||||
"""要起草的表 → (表结构里的键,去重、按给的顺序;表结构里没有的原样名字)。"""
|
||||
picked: list[str] = []
|
||||
for name in wanted:
|
||||
unknown: list[str] = []
|
||||
for name in dict.fromkeys(str(t) for t in wanted):
|
||||
key = resolve_table_name(cache, name)
|
||||
if key is None:
|
||||
report.tables.append(TableDraftResult(table=name, error=f"表结构里没有 {name},请先重新探查结构"))
|
||||
unknown.append(name)
|
||||
elif key not in picked:
|
||||
picked.append(key)
|
||||
return picked, unknown
|
||||
|
||||
llm: dict[str, TableDraft | str] = {}
|
||||
|
||||
async def default_tables(session: AsyncSession, source: Any) -> list[str]:
|
||||
"""不指定表时起草哪些:按使用次数(table_usage)取前 DRAFT_DEFAULT_TABLES 张,次数相同按表结构里的顺序。"""
|
||||
all_tables = (getattr(source, "schema_cache", None) or {}).get("tables") or {}
|
||||
usage = await table_usage(session, source)
|
||||
order = {name: i for i, name in enumerate(all_tables)}
|
||||
return sorted(all_tables, key=lambda n: (-usage.get(n, 0), order[n]))[:DRAFT_DEFAULT_TABLES]
|
||||
|
||||
|
||||
async def draft_catalog(session: AsyncSession, source: Any, *, tables: list[str] | None = None,
|
||||
use_model: bool = False, model: Any = None, actor: str | None = None,
|
||||
guidance: str | None = None, note: str | None = None) -> DraftReport:
|
||||
"""起草这些表的目录并写入,同步完成(POST /catalog/draft;界面改用后台起草任务,见 data/catalog_jobs.py)。
|
||||
|
||||
- source:手工源的 DataSource,或上传源绑定当前快照的视图(table_versions.resolve_source);按 source.id 存。
|
||||
- tables 为空时按使用次数取前 DRAFT_DEFAULT_TABLES 张(default_tables)。
|
||||
- 每张表:数据库注释、外键约束、命名推断照做;use_model 时再加模型起草(data/catalog_draft.py:分段、按失败
|
||||
原因处理)。model 可注入(测试用假模型),不给就按设置解析(resolve_draft_model);模型用不了、某一段失败,
|
||||
都不影响其余来源照常写入。guidance 是数据源级的起草要求,note 是本次说明,和数据源说明一起交给模型。
|
||||
- 并入、写入见 write_draft。
|
||||
"""
|
||||
cache = getattr(source, "schema_cache", None) or {}
|
||||
report = DraftReport()
|
||||
wanted = list(tables) if tables else await default_tables(session, source)
|
||||
picked, unknown = pick_tables(cache, wanted)
|
||||
for name in unknown:
|
||||
report.tables.append(TableDraftResult(table=name, error=missing_table_error(name), status="skipped"))
|
||||
|
||||
outcomes: dict[str, Any] = {}
|
||||
if use_model and picked:
|
||||
try:
|
||||
if model is None:
|
||||
@@ -1872,41 +1851,17 @@ async def draft_catalog(session: AsyncSession, source: Any, *, tables: list[str]
|
||||
except CatalogModelUnavailable as e:
|
||||
report.model_error = str(e)
|
||||
else:
|
||||
from app.data import catalog_draft
|
||||
|
||||
report.model_used = True
|
||||
existing = await read_catalog(session, source.id)
|
||||
llm = await draft_with_model(model, source, picked, existing={t: e.notes for t, e in existing.items()})
|
||||
outcomes = await catalog_draft.run_model_draft(
|
||||
model, source, picked, existing={t: e.notes for t, e in existing.items()},
|
||||
background=catalog_draft.Background(guidance=guidance, note=note))
|
||||
|
||||
at = now_iso()
|
||||
for table in picked:
|
||||
result = TableDraftResult(table=table)
|
||||
structure = draft_structure(source, table)
|
||||
modeled = llm.get(table)
|
||||
if isinstance(modeled, str):
|
||||
result.model_error = modeled
|
||||
elif isinstance(modeled, TableDraft):
|
||||
result.model_items = len(_slots(modeled.notes))
|
||||
for _attempt in range(3):
|
||||
entry = await read_entry(session, source.id, table)
|
||||
version = entry.version if entry else 0
|
||||
merged, stats = merge_notes(entry.notes if entry else {}, structure.notes, covered=structure.covered,
|
||||
at=at)
|
||||
if isinstance(modeled, TableDraft):
|
||||
merged, more = merge_notes(merged, modeled.notes, covered=modeled.covered, at=at)
|
||||
stats += more
|
||||
result.added, result.updated, result.removed = stats.added, stats.updated, stats.removed
|
||||
result.version = version
|
||||
if entry is None and not merged:
|
||||
break # 什么也没起草出来,不建空目录
|
||||
try:
|
||||
written = await write_entry(session, source.id, table, merged, if_version=version, actor=actor)
|
||||
except CatalogConflict:
|
||||
continue
|
||||
result.version = written.version
|
||||
break
|
||||
else:
|
||||
result.error = "这张表的数据目录正被其他人修改,请稍后再起草"
|
||||
result.added = result.updated = result.removed = 0
|
||||
report.tables.append(result)
|
||||
report.tables.append(await write_draft(session, source, table, outcomes.get(table), at=at, actor=actor))
|
||||
return report
|
||||
|
||||
|
||||
@@ -2271,12 +2226,16 @@ def frozen_catalog(entries: Mapping[str, CatalogEntry], tables: Iterable[str]) -
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CARDINALITIES", "CARDINALITY_LABEL", "CATALOG_DRAFT_PROMPT", "CATALOG_DRAFT_PROMPT_SCHEMA",
|
||||
"CATALOG_DRAFT_SYSTEM", "COLUMN_FIELDS", "COLUMN_FIELD_LABEL", "COLUMN_MEASURES", "CatalogConflict",
|
||||
"CARDINALITIES", "CARDINALITY_LABEL", "CATALOG_DRAFT_BACKGROUND_PROMPT", "CATALOG_DRAFT_PROMPT",
|
||||
"CATALOG_DRAFT_PROMPT_LABELS", "CATALOG_DRAFT_PROMPT_SCHEMA", "CATALOG_DRAFT_SYSTEM", "CATALOG_DRAFT_TEXT_PROMPT",
|
||||
"COLUMN_FIELDS", "COLUMN_FIELD_LABEL", "COLUMN_MEASURES", "CatalogConflict",
|
||||
"CatalogEntry", "CatalogInvalid", "CatalogModelUnavailable", "CatalogPathError", "DRAFT_DEFAULT_TABLES",
|
||||
"DraftReport", "INFERRED_MARK", "ITEM_SOURCES", "ITEM_STATUSES", "JoinEdge", "KIND_LABEL", "MEASURE_LABEL",
|
||||
"MODEL_EMPTY_REASON", "MergeStats", "NEW_RELATION_PATH", "PATCH_MAX_CHANGES", "PROFILE_NOTE_PREFIX",
|
||||
"PROFILE_VERIFY_COVERAGE",
|
||||
"DraftModel", "DraftReport", "INFERRED_MARK", "ITEM_SOURCES", "ITEM_STATUSES", "JoinEdge", "KIND_LABEL",
|
||||
"MEASURE_LABEL", "MODEL_BATCH_COLUMNS", "MODEL_BATCH_TABLES", "MODEL_CONCURRENCY", "MODEL_EMPTY_REASON",
|
||||
"MODEL_FALLBACK_MAX_TOKENS", "MODEL_MAX_TOKENS", "MODEL_RETRIES", "MODEL_RETRY_DELAYS_S", "MODEL_SPLIT_DEPTH",
|
||||
"MODEL_SPLIT_MIN_COLUMNS", "MODEL_TIMEOUT_S", "MergeStats", "NEW_RELATION_PATH", "PATCH_MAX_CHANGES",
|
||||
"PROFILE_NOTE_PREFIX", "PROFILE_VERIFY_COVERAGE", "WRITE_CONFLICT_ERROR", "default_tables",
|
||||
"missing_table_error", "pick_tables", "write_draft",
|
||||
"PatchChange", "PatchPlan", "QUERY_SNAPSHOT_KIND", "REVIEW_ACTIONS", "TABLE_FIELDS", "TABLE_FIELD_LABEL",
|
||||
"TABLE_KINDS", "TableDraft", "TableDraftResult", "UI_KIND_LABEL", "UI_MEASURE_LABEL", "apply_human_edit",
|
||||
"apply_patch", "codes_complete", "describe_slot", "draft_catalog", "draft_structure", "draft_with_model",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+55
-7
@@ -3,10 +3,11 @@
|
||||
model = FakeModel({"visits": {"label": "入园记录", "kind": "fact"}}, fail={"orders"})
|
||||
await catalog.draft_with_model(model, source, ["visits"])
|
||||
|
||||
- 结构化输出(with_structured_output(...).ainvoke)从提示词里认出这一批有哪些表(「表 <表名>」),
|
||||
返回 {"tables": [{"table": 表名, **replies[表名]}, …]};
|
||||
- 结构化输出(with_structured_output(...).ainvoke)从提示词里认出这一段请它写哪些表、每张表写哪部分(requested:
|
||||
整张表、只写表级的项、只写列出的几列、表级的项加几列),返回 {"tables": [{"table": 表名, …}, …]},只给请它写的
|
||||
那部分(整张表给 replies[表名] 的全部;只写几列时只给 replies 里这几列);
|
||||
- 直接 ainvoke 是纯文本那条路:同样的内容写成一段正文(AIMessage),前面带一句话;
|
||||
- fail 里的表所在的那一批抛异常,用来验证「一批失败只影响这一批」;
|
||||
- fail 里的表所在的那一段抛异常,用来验证「一段失败只影响这一段」;
|
||||
- calls 记下每次收到的提示词全文,用来核对交给模型的内容(只有结构,没有数据);modes 记下每次走的是
|
||||
structured 还是 text。
|
||||
|
||||
@@ -22,6 +23,51 @@ from __future__ import annotations
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from app.data.catalog import CATALOG_DRAFT_PROMPT_LABELS as _P
|
||||
|
||||
#: 表名下一行写着「这次只写…」时,请模型写的是哪一部分(data/catalog_draft.py 按这几句写提示词)
|
||||
_ONLY_TABLE = _P["only_table"]
|
||||
_ONLY_COLUMNS = _P["only_columns"].split("{")[0]
|
||||
_TABLE_AND = _P["table_and_columns"].split("{")[0]
|
||||
|
||||
|
||||
def requested(text: str) -> dict[str, tuple[str, list[str] | None]]:
|
||||
"""提示词里请模型写哪些表、每张表写哪部分:{表名: (哪部分, 列)}。哪部分是 full(整张表)、table(只写表级的项)、
|
||||
columns(只写列出的列,列是表结构描述里列出的那几列)、table+columns(表级的项加「这次只写」那一行点名的列)。"""
|
||||
out: dict[str, tuple[str, list[str] | None]] = {}
|
||||
for chunk in text.split("\n\n"):
|
||||
lines = chunk.split("\n")
|
||||
head = lines[0]
|
||||
if not head.startswith("表 "):
|
||||
continue
|
||||
name = head[2:]
|
||||
if name.endswith("(视图)"):
|
||||
name = name[:-len("(视图)")]
|
||||
second = lines[1] if len(lines) > 1 else ""
|
||||
listed = [line[2:].split(" ")[0] for line in lines if line.startswith("- ")]
|
||||
if second.startswith(_ONLY_TABLE):
|
||||
out[name] = ("table", None)
|
||||
elif second.startswith(_ONLY_COLUMNS):
|
||||
out[name] = ("columns", listed)
|
||||
elif second.startswith(_TABLE_AND):
|
||||
out[name] = ("table+columns", second.split(":", 1)[1].split("、"))
|
||||
else:
|
||||
out[name] = ("full", None)
|
||||
return out
|
||||
|
||||
|
||||
def answer(reply: dict, name: str, part: str, columns: list[str] | None) -> dict:
|
||||
"""按请它写的那部分,从 replies[表名] 里取。"""
|
||||
if part == "full":
|
||||
return {"table": name, **reply}
|
||||
out: dict[str, Any] = {"table": name}
|
||||
if part in ("table", "table+columns"):
|
||||
out.update({k: v for k, v in reply.items() if k != "columns"})
|
||||
if part in ("columns", "table+columns"):
|
||||
wanted = set(columns or [])
|
||||
out["columns"] = [c for c in reply.get("columns") or [] if isinstance(c, dict) and c.get("name") in wanted]
|
||||
return out
|
||||
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self, replies: dict[str, dict], fail: set[str] = frozenset(), *,
|
||||
@@ -44,17 +90,19 @@ class FakeModel:
|
||||
text = "\n".join(str(getattr(m, "content", m)) for m in messages)
|
||||
self.calls.append(text)
|
||||
self.modes.append(mode)
|
||||
names = [t for t in self.replies if f"表 {t}\n" in text or f"表 {t}(" in text]
|
||||
if any(t in self.fail for t in names):
|
||||
asked = requested(text)
|
||||
names = [t for t in asked if t in self.replies]
|
||||
if any(t in self.fail for t in asked):
|
||||
raise RuntimeError("网关超时")
|
||||
shape = self.structured if mode == "structured" else self.text
|
||||
if shape == "error":
|
||||
raise NotImplementedError("这条路用不了")
|
||||
if shape == "hollow":
|
||||
return {"tables": [{"table": t} for t in names]}
|
||||
answers = [answer(self.replies[t], t, *asked[t]) for t in names]
|
||||
if shape == "single":
|
||||
return {"table": names[0], **self.replies[names[0]]} if names else {}
|
||||
return {"tables": [{"table": t, **self.replies[t]} for t in names]}
|
||||
return answers[0] if answers else {}
|
||||
return {"tables": answers}
|
||||
|
||||
|
||||
class _Structured:
|
||||
|
||||
@@ -274,10 +274,11 @@ async def test_model_single_table_object_without_wrapper_is_recognized():
|
||||
|
||||
|
||||
async def test_hollow_structured_output_retries_with_text():
|
||||
"""结构化输出只给了 {"tables":[{"table":"visits"}]} 这样的空壳:改走纯文本再问一次,用纯文本给的内容。"""
|
||||
"""结构化输出只给了 {"tables":[{"table":"visits"}]} 这样的空壳、这一段又拆不开(不到 10 列):改走纯文本再问一次,
|
||||
用纯文本给的内容。纯文本只给了表级的项、没给列:再补问一次没给的列(只问缺的,data/catalog_draft.py)。"""
|
||||
model = FakeModel({"visits": {"label": "入园记录", "kind": "fact"}}, structured="hollow", text="tables")
|
||||
got = await catalog.draft_with_model(model, _source(_scenic_tables()), ["visits"])
|
||||
assert model.modes == ["structured", "text"]
|
||||
assert model.modes == ["structured", "text", "text"]
|
||||
assert got["visits"].notes == {"label": make_item("入园记录", "llm"), "kind": make_item("fact", "llm")}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,555 @@
|
||||
"""数据目录模型起草的引擎(data/catalog_draft.py):分段、按失败原因处理、捡回被截断的部分、起草要求、用量。
|
||||
|
||||
用假模型复刻实测查明的几种失败(不调用真实模型):
|
||||
- 思考加输出用完额度:finish_reason=length 且正文为空(只有思考)、LengthFinishReasonError → 拆成两半再问;
|
||||
- 只给空壳 → 拆开再问;拆不开的换纯文本问一次,写出来了这次起草后面都走纯文本;
|
||||
- 被截断的 JSON:已经完整的表、列对象捡回来,只补剩下的;
|
||||
- 只回一部分表:只补缺的,补一次;
|
||||
- 服务方嫌额度太大 → 退回 8192,后面都用它;结构化输出不支持 → 后面改走纯文本;
|
||||
- 网络错误、5xx、超时 → 退避重试,最多 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()}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 思考用完额度、空壳:拆开再问
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
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():
|
||||
tables = {"wide40": _table(40)}
|
||||
model = Scripted(lambda call: Reply(text="", finish="length", thinking="……"))
|
||||
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_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 的格式"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 截断、缺表、缺列
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
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) == 2
|
||||
assert model.calls[1].asked == {"b_tab": ("full", None), "c_tab": ("full", None)}
|
||||
assert _status(got) == {"a_tab": "ok", "b_tab": "ok", "c_tab": "failed"}
|
||||
assert got["c_tab"].error == "模型的回复里没有这张表"
|
||||
|
||||
|
||||
async def test_columns_the_model_skips_are_asked_once_and_then_left_empty():
|
||||
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 == "ok" and o.error is None, "补问过仍然没给的列算拿不准、留空"
|
||||
assert set(o.notes["columns"]) == {f"c{i:03d}" for i in range(1, 7)}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 额度被拒、结构化输出不支持、网络错误、超时
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
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"]
|
||||
|
||||
|
||||
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_retried_then_reported(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) == 3
|
||||
assert got["t4"].error == "模型 0.05 秒内没有写完这一段(已重试 2 次)"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 分段:宽表拆成表级一段 + 列分段,窄表拼段
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
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) == [15, 40, 40]
|
||||
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)}
|
||||
|
||||
|
||||
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"]) == 55
|
||||
assert o.error == "这张表 95 列,分 4 段起草,第 3 段没写出来(共 40 列):调用模型出错:网关拒绝。可以只重试这张表"
|
||||
compat = (await catalog.draft_with_model(Scripted(handle), _source(tables), ["wide95"]))["wide95"]
|
||||
assert compat.covered == frozenset() and compat.error == o.error, "没写全的不算完整起草过,旧的模型项不删"
|
||||
|
||||
|
||||
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_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)
|
||||
|
||||
seen: list[tuple[int, int]] = []
|
||||
|
||||
async def usage(i: int, o: int) -> None:
|
||||
seen.append((i, o))
|
||||
|
||||
await D.run_model_draft(Scripted(handle), _source(tables), ["orders_x"], hooks=D.Hooks(usage=usage))
|
||||
assert seen == [(100, 32768), (120, 60), (120, 60)]
|
||||
|
||||
|
||||
def test_budget_words_match_the_other_drafts():
|
||||
from app.data import recipe_ai
|
||||
|
||||
assert D.BUDGET_WORDS == recipe_ai._BUDGET_WORDS
|
||||
Reference in New Issue
Block a user