diff --git a/backend/app/data/catalog.py b/backend/app/data/catalog.py index 281633a7..c48b99bc 100644 --- a/backend/app/data/catalog.py +++ b/backend/app/data/catalog.py @@ -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", diff --git a/backend/app/data/catalog_draft.py b/backend/app/data/catalog_draft.py new file mode 100644 index 00000000..2b1d71bf --- /dev/null +++ b/backend/app/data/catalog_draft.py @@ -0,0 +1,1048 @@ +"""数据目录的模型起草：分段、按失败原因处理、捡回被截断的部分、记用量。 + +对外的名字留在 data/catalog.py（draft_with_model、draft_catalog、write_draft、CATALOG_DRAFT_SYSTEM、各个 MODEL_* +常量……），这里是实现。常量一律在调用时按 catalog.X 读：测试改 catalog 上的常量就能生效。 + +**为什么要分段、按原因处理。** 实测（经 OpenAI 兼容网关的会思考的模型）：思考算在输出额度里，一段上百列时思考就把 +额度全用掉，正文 0 字或只给一个空壳；同一批表有时成功有时失败，看模型这次想多久。以前失败了退回纯文本、发同一个提示词， +同样失败、耗时翻倍；超时把第一次拿到的部分也丢掉；一批里有一张表是空的就整批重问。 + +**分段。** 每次调用最多 MODEL_BATCH_TABLES 张表、合计 MODEL_BATCH_COLUMNS 列：列数不超过它的表按顺序拼在一起 +（整表段）；超过它的宽表拆成「表级一段 + 列分段」——表级那段给全部列名和类型做背景、只要表级的项（中文名、粒度、 +表类型、业务主键、业务日期）；列分段每段 MODEL_BATCH_COLUMNS 列，带着表名、表注释和表级起草结果做背景，只要这几列。 +最后合并成这张表的一份，照旧经 catalog._from_model 逐项校验（起草要求不能让模型绕过校验）。 + +**按失败原因处理，不原样重发。** + +- 思考加输出用完额度（finish_reason=length、LengthFinishReasonError、正文为空而有思考）、只给空壳：把这一段拆成 + 两半再问（多表的段按表拆，单表的按列拆），最小 MODEL_SPLIT_MIN_COLUMNS 列，最多拆 MODEL_SPLIT_DEPTH 层。 + 拆不开的空壳（结构化输出给的）换纯文本问一次：有的网关的结构化输出只给空壳、纯文本却能写——写出来了，这次起草 + 后面都走纯文本。 +- 被截断的 JSON：先把已经完整的表、列对象捡回来，只补剩下的。 +- 一段里缺了的表、列只补缺的，补一次。补过仍然没给的列算模型拿不准、留空，不算没写出来；整张表一项都没有才算。 +- 结构化输出这个接口不支持（400 / 422 提到 response_format、json_schema、tool……，或者模型类不支持）：这次起草 + 后面改走纯文本、从正文取 JSON（纯文本那条路另外说清 JSON 的格式）。 +- 服务方嫌额度太大：退回 MODEL_FALLBACK_MAX_TOKENS 重来，这次起草后面都用它（catalog.DraftModel.shrink）。 +- 网络错误、429、5xx、超时：等一会儿重试（最多 MODEL_RETRIES 次，退避）。单次调用的时限是 MODEL_TIMEOUT_S。 +- 别的错误：这一段没写出来，原因照实写。 + +**并发。** 全进程最多 MODEL_CONCURRENCY 个调用同时进行（model_gate：按事件循环各一个信号量，所有起草任务和同步 +起草共用）。 + +**进度。** 调用方给 Hooks：开始调用、每次调用的用量、一段结束、一张表的模型部分全部结束。后台起草任务据此落库， +每张表一结束就并入目录（data/catalog_jobs.py）。 + +交给模型的只有表结构（列、类型、主键、外键、唯一约束、注释）和使用者补充的背景，不发任何数据值。写给模型的文字都在 +catalog.py 名字带 _SYSTEM / _PROMPT 的常量里；这里的 REASON_* 是给人看的（拼进这张表的 model_error）。 +""" +from __future__ import annotations + +import asyncio +import copy +import json +import logging +import re +import weakref +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from app.data import catalog as C +from app.data.names import name_key + +logger = logging.getLogger(__name__) + +#: 退避等待用的 sleep（测试里换掉，不真等） +_sleep = asyncio.sleep + +#: 服务方嫌额度太大时报错里会出现的字样（同 copilot._BUDGET_WORDS、recipe_ai._BUDGET_WORDS，有测试核对） +BUDGET_WORDS = ("max_tokens", "max_completion_tokens", "max_output_tokens", "maximum context length") +#: 结构化输出这个接口不支持时报错里会出现的字样：json_schema 的 response_format（有的服务只认 json_object）、 +#: 工具调用（有的网关不支持）、Anthropic 的 output_config +STRUCTURED_WORDS = ("response_format", "json_schema", "json_object", "tool", "function", "output_config", + "output_format", "structured output") +#: 一张表的表级的项在「没写出来」记录里的键（不会和列名撞上） +_TABLE = "\0table" + +# ---- 一段没写出来的原因（给人看，拼进这张表的 model_error） +REASON_EXHAUSTED = "模型的思考和输出用完了额度（{budget} 个 token），这一段没写出来" +REASON_SPLIT = "；已拆小重试仍没写出" +REASON_HOLLOW = "模型只给了空壳，没有写出可用的内容" +REASON_TRUNCATED = "模型的输出被截断，这一段没写完" +REASON_FORMAT = "模型的回复不是约定的格式" +REASON_MISSING = "模型的回复里没有这张表" +REASON_TIMEOUT = "模型 {seconds} 秒内没有写完这一段" +REASON_CALL = "调用模型出错：{line}" +REASON_RETRIED = "（已重试 {n} 次）" +REASON_CRASH = "起草这张表时出错：{line}" + + +# ========================================================================== +# 对外的形状 +# ========================================================================== + + +@dataclass(frozen=True) +class Background: + """使用者补充的背景和要求：guidance 是数据源级的起草要求，note 是这一次的补充说明。数据源说明（description） + 从源上读。都没有就不出现那一段。""" + + guidance: str | None = None + note: str | None = None + + +@dataclass +class ModelOutcome: + """一张表的模型起草结果。 + + status：ok 各段都写出来了；partial 有的段写出来、有的没有（error 写清哪部分、为什么）；failed 一项都没写出来 + （error 写原因）。notes 是写出来的那些（已逐项校验，llm / proposed）；items 是项数。parts_total 分几段起草 + （整表段是 1，宽表是 1 + 列分段数），parts_done 已经结束了几段。""" + + status: str + notes: dict[str, Any] + items: int + error: str | None + parts_total: int + parts_done: int + + +@dataclass +class Hooks: + """进度回调，都是 async，出错只记日志、不影响起草。 + + - call_started(表名列表)：一次调用拿到并发名额、马上要发出去了； + - usage(输入 token, 输出 token)：每次调用一回（拿不到用量的是 0），条数就是调用次数； + - part_done(表名)：这张表的一段结束了（写没写出来都算）； + - table_done(表名, ModelOutcome)：这张表的模型部分全部结束。""" + + call_started: Callable[[list[str]], Awaitable[None]] | None = None + usage: Callable[[int, int], Awaitable[None]] | None = None + part_done: Callable[[str], Awaitable[None]] | None = None + table_done: Callable[[str, ModelOutcome], Awaitable[None]] | None = None + + +def parts_for(columns: int) -> int: + """一张有这么多列的表分几段起草：不超过 MODEL_BATCH_COLUMNS 列是 1 段，否则是表级一段加列分段。""" + width = max(1, int(C.MODEL_BATCH_COLUMNS)) + return 1 if columns <= width else 1 + -(-columns // width) + + +async def run_model_draft(model: Any, source: Any, tables: Sequence[str], *, + existing: Mapping[str, Mapping[str, Any]] | None = None, + background: Background | None = None, + hooks: Hooks | None = None) -> dict[str, ModelOutcome]: + """用模型起草这些表（schema_cache["tables"] 的键；表结构里没有的跳过，不在结果里）。 + + model 要有 with_structured_output(schema, include_raw=…) 和 ainvoke(messages)；有 shrink() 的（catalog.DraftModel） + 在服务方嫌额度太大时换小额度。existing 是 {表名: 目录}：已确认、有确证的项作为「以此为准」的背景交给模型。 + """ + run = _Run(model, source, tables, existing or {}, background or Background(), hooks or Hooks()) + return await run.run() + + +# ========================================================================== +# 并发名额 +# ========================================================================== + +_GATES: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, tuple[int, asyncio.Semaphore]]" = \ + weakref.WeakKeyDictionary() + + +def model_gate() -> asyncio.Semaphore: + """全进程的模型调用名额（MODEL_CONCURRENCY 个）。按事件循环各一个：asyncio 的信号量一旦在某个循环里等过， + 就不能换循环用（测试每个用例一个循环）。""" + loop = asyncio.get_running_loop() + limit = max(1, int(C.MODEL_CONCURRENCY)) + held = _GATES.get(loop) + if held is None or held[0] != limit: + held = (limit, asyncio.Semaphore(limit)) + _GATES[loop] = held + return held[1] + + +# ========================================================================== +# 失败的种类 +# ========================================================================== + + +def _status_of(e: BaseException) -> int | None: + for obj in (e, getattr(e, "response", None)): + code = getattr(obj, "status_code", None) + if isinstance(code, int): + return code + return None + + +def classify_failure(e: BaseException) -> str: + """调用模型抛出来的异常属于哪一种：length 额度用完（LengthFinishReasonError）、budget 服务方嫌额度太大、 + unsupported 结构化输出不支持、timeout 超时、transient 网络错误 / 429 / 5xx、error 别的。""" + from app.core.errors import NETWORK, RATE_LIMIT, SERVER, TIMEOUT, classify + + names = {c.__name__ for c in type(e).__mro__} + if "LengthFinishReasonError" in names: + return "length" + status = _status_of(e) + text = str(e).lower() + if status in (400, 422) and any(w in text for w in BUDGET_WORDS): + return "budget" + # thinking 打开时 Anthropic 不让强制调用工具，langchain 改成不强制，模型没调用工具就抛 OutputParserException + if isinstance(e, NotImplementedError) or "OutputParserException" in names: + return "unsupported" + if status in (400, 422) and any(w in text for w in STRUCTURED_WORDS): + return "unsupported" + if isinstance(e, TimeoutError): + return "timeout" + kind = classify(e) + if kind == TIMEOUT: + return "timeout" + if kind in (RATE_LIMIT, SERVER, NETWORK): + return "transient" + return "error" + + +def _seconds(value: Any) -> Any: + number = float(value) + return int(number) if number.is_integer() else number + + +# ========================================================================== +# 读回复：结构化输出、正文里的 JSON、被截断的 JSON +# ========================================================================== + + +@dataclass +class _Reply: + #: 回复里的表对象（截断的只含已经完整的字段和列） + entries: list[dict[str, Any]] = field(default_factory=list) + #: JSON 没写完（捡回了一部分） + truncated: bool = False + #: 思考加输出用完了额度 + exhausted: bool = False + #: 这一段最终没写出来的原因（重试用完、别的错误、不是约定的格式） + error: str | None = None + + +class _Cut: + """被截断的容器（对象或数组）：value 里只有已经完整的成员，最后一个没写完的容器成员也包成 _Cut 留着，好往里捡。""" + + __slots__ = ("value",) + + def __init__(self, value: Any) -> None: + self.value = value + + +class _Bad(ValueError): + """不是 JSON（不是写到一半，是写错了）。""" + + +_NUMBER = re.compile(r"-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?") +_LITERALS = (("true", True), ("false", False), ("null", None)) +_MAX_DEPTH = 64 + + +def _skip(text: str, i: int) -> int: + n = len(text) + while i < n and text[i] in " \t\r\n": + i += 1 + return i + + +def _string(text: str, i: int) -> tuple[Any, int, bool]: + n = len(text) + j = i + 1 + while j < n: + ch = text[j] + if ch == "\\": + j += 2 + continue + if ch == '"': + try: + return json.loads(text[i:j + 1]), j + 1, True + except ValueError as e: + raise _Bad from e + j += 1 + return None, n, False + + +def _value(text: str, i: int, depth: int) -> tuple[Any, int, bool]: + """从 text[i] 起读一个 JSON 值：(值, 读到哪, 写完了没有)。读到结尾还没写完的容器返回 _Cut。""" + if depth > _MAX_DEPTH: + raise _Bad + n = len(text) + i = _skip(text, i) + if i >= n: + return None, i, False + ch = text[i] + if ch == "{": + return _object(text, i + 1, depth + 1) + if ch == "[": + return _array(text, i + 1, depth + 1) + if ch == '"': + return _string(text, i) + for word, value in _LITERALS: + if text.startswith(word, i): + return value, i + len(word), True + if word.startswith(text[i:]): + return None, n, False + m = _NUMBER.match(text, i) + if not m: + raise _Bad + j = m.end() + if j >= n: + return None, j, False # 数字写到了结尾，可能还没写完 + number = m.group() + return (float(number) if any(c in number for c in ".eE") else int(number)), j, True + + +def _object(text: str, i: int, depth: int) -> tuple[Any, int, bool]: + out: dict[str, Any] = {} + n = len(text) + while True: + i = _skip(text, i) + if i >= n: + return _Cut(out), i, False + ch = text[i] + if ch == "}": + return out, i + 1, True + if ch == ",": + i += 1 + continue + if ch != '"': + raise _Bad + key, i, ok = _string(text, i) + if not ok: + return _Cut(out), i, False + i = _skip(text, i) + if i >= n: + return _Cut(out), i, False + if text[i] != ":": + raise _Bad + value, i, ok = _value(text, i + 1, depth) + if not ok: + if isinstance(value, _Cut): + out[key] = value + return _Cut(out), i, False + out[key] = value + + +def _array(text: str, i: int, depth: int) -> tuple[Any, int, bool]: + out: list[Any] = [] + n = len(text) + while True: + i = _skip(text, i) + if i >= n: + return _Cut(out), i, False + ch = text[i] + if ch == "]": + return out, i + 1, True + if ch == ",": + i += 1 + continue + value, i, ok = _value(text, i, depth) + if not ok: + if isinstance(value, _Cut): + out.append(value) + return _Cut(out), i, False + out.append(value) + + +def _json_prefix(text: str) -> tuple[Any, bool] | None: + """正文里的第一个 JSON 值：(值, 写完了没有)。前面可以有一句话、围栏，后面可以有别的字；不是 JSON 返回 None。""" + for start in sorted(k for k in (text.find("{"), text.find("[")) if k >= 0): + try: + value, _, complete = _value(text, start, 0) + except (_Bad, RecursionError): + continue + return value, complete + return None + + +def _entry(item: Any) -> dict[str, Any] | None: + """一个表对象：只留写完了的字段；columns 留写完了的列对象。没有表名的不要。""" + if isinstance(item, _Cut): + item = item.value + if not isinstance(item, dict) or not isinstance(item.get("table"), str): + return None + out: dict[str, Any] = {} + for key, value in item.items(): + if key == "columns": + cols = value.value if isinstance(value, _Cut) else value + if isinstance(cols, list): + out["columns"] = [c for c in cols if isinstance(c, dict)] + elif not isinstance(value, _Cut): + out[key] = value + return out + + +def _tables_of(root: Any) -> list[dict[str, Any]] | None: + """解析出来的根 → 表对象列表。{"tables": […]}、[…]、单个表对象（实测见过不带 tables 外壳的）都认。""" + if isinstance(root, _Cut): + root = root.value + if isinstance(root, dict): + if "tables" in root: + items = root["tables"] + elif isinstance(root.get("table"), str): + items = [root] + elif not root: + return [] + else: + return None + else: + items = root + if isinstance(items, _Cut): + items = items.value + if not isinstance(items, list): + return None + return [e for e in (_entry(item) for item in items) if e is not None] + + +def _parse_text(text: str, cut: bool) -> _Reply | None: + found = _json_prefix(text) + if found is None: + return None + root, complete = found + entries = _tables_of(root) + if entries is None: + return None + if complete: + return _Reply(entries=entries) + return _Reply(entries=entries, truncated=True, exhausted=cut) + + +def _cut_off(raw: Any) -> bool: + """额度用完停下的：OpenAI 系是 finish_reason=length，Anthropic 是 stop_reason=max_tokens。""" + meta = getattr(raw, "response_metadata", None) or {} + return (meta.get("finish_reason") or meta.get("stop_reason")) in ("length", "max_tokens") + + +def _thought(raw: Any) -> bool: + """这次回复有没有思考：思考内容（Anthropic 的 thinking 块、OpenAI 兼容系的 reasoning_content），或者用量里有推理 token。""" + from app.engine.state import thinking_text + + if raw is None: + return False + if thinking_text(raw): + return True + details = (getattr(raw, "usage_metadata", None) or {}).get("output_token_details") or {} + return bool(details.get("reasoning")) + + +def _text_of(raw: Any) -> str: + from app.engine.state import message_text + + return message_text(raw) if raw is not None else "" + + +def _raw_texts(raw: Any) -> list[str]: + """结构化输出没解析出来时，原始回复里可能还有 JSON 的地方：工具调用的参数（截断的在 invalid_tool_calls 里，是一段 + 字符串）、正文。""" + out: list[str] = [] + for call in getattr(raw, "tool_calls", None) or []: + args = call.get("args") if isinstance(call, dict) else None + if isinstance(args, dict) and args: + out.append(json.dumps(args, ensure_ascii=False)) + for call in getattr(raw, "invalid_tool_calls", None) or []: + args = call.get("args") if isinstance(call, dict) else None + if isinstance(args, str) and args.strip(): + out.append(args) + text = _text_of(raw) + if text.strip(): + out.append(text) + return out + + +def _entries_of(parsed: Any) -> list[dict[str, Any]] | None: + if parsed is None: + return None + if isinstance(parsed, dict) and not parsed: + return [] + return C._model_tables(parsed) + + +def _structured_reply(raw: Any, parsed: Any) -> _Reply: + entries = _entries_of(parsed) + if entries is not None: + return _Reply(entries=entries) + cut = _cut_off(raw) + texts = _raw_texts(raw) + for text in texts: + reply = _parse_text(text, cut) + if reply is not None: + return reply + if cut or _thought(raw): + return _Reply(exhausted=True) + if not texts: + return _Reply() # 什么都没写：按空壳处理 + return _Reply(error=REASON_FORMAT) + + +def _text_reply(raw: Any) -> _Reply: + text = _text_of(raw) + cut = _cut_off(raw) + if text.strip(): + reply = _parse_text(text, cut) + if reply is not None: + return reply + if cut: + return _Reply(exhausted=True) # 写了字、额度就用完了，一个 JSON 对象都没捡回来 + return _Reply(error=REASON_FORMAT) + return _Reply(exhausted=cut or _thought(raw)) + + +def _unpack(out: Any) -> tuple[Any, Any]: + """with_structured_output(include_raw=True) 交回 {raw, parsed, parsing_error}；测试的假模型直接交回解析好的对象。""" + if isinstance(out, dict) and "raw" in out and ("parsed" in out or "parsing_error" in out): + return out.get("raw"), out.get("parsed") + return None, out + + +# ========================================================================== +# 一次起草 +# ========================================================================== + + +@dataclass(frozen=True) +class _Piece: + """一段里的一张表：要写哪些列、要不要表级的项、属于这张表的第几段（从 1 起）。""" + + table: str + columns: tuple[str, ...] + table_level: bool + part: int + + +@dataclass +class _Gain: + """一次回复里某张表拿到了什么。""" + + named: bool = False + items: int = 0 + table_items: int = 0 + answered: set[str] = field(default_factory=set) + + +@dataclass +class _Table: + name: str + meta: Mapping[str, Any] + columns: list[str] + existing: Mapping[str, Any] | None + parts: list[_Piece] = field(default_factory=list) + #: 已经拿到的模型起草（逐段合并） + notes: dict[str, Any] = field(default_factory=dict) + #: 模型的回复里出现过这张表 + named: bool = False + got_table: bool = False + got_columns: set[str] = field(default_factory=set) + #: 没写出来的列（表级的项记在 _TABLE）→ 原因 + failed: dict[str, str] = field(default_factory=dict) + parts_done: int = 0 + finished: bool = False + + def whole(self, piece: _Piece) -> bool: + return piece.table_level and len(piece.columns) == len(self.columns) + + +def _restrict(notes: Mapping[str, Any], piece: _Piece) -> dict[str, Any]: + """只留这一段请模型写的部分：没请写表级的项就不要表级的项，列只要请写的那几列。""" + out: dict[str, Any] = {} + if piece.table_level: + out.update({k: v for k, v in notes.items() if k != "columns"}) + wanted = set(piece.columns) + cols = {c: items for c, items in (notes.get("columns") or {}).items() if c in wanted} + if cols: + out["columns"] = cols + return out + + +def _halves(pieces: list[_Piece]) -> tuple[list[_Piece], list[_Piece]] | None: + """拆成两半：多表的段按表拆；单表的按列拆，两半都至少 MODEL_SPLIT_MIN_COLUMNS 列（表级的项跟着前一半）。拆不开返回 None。""" + if len(pieces) > 1: + k = (len(pieces) + 1) // 2 + return pieces[:k], pieces[k:] + piece = pieces[0] + n = len(piece.columns) + if n < 2 * max(1, int(C.MODEL_SPLIT_MIN_COLUMNS)): + return None + h = (n + 1) // 2 + return ([_Piece(piece.table, piece.columns[:h], piece.table_level, piece.part)], + [_Piece(piece.table, piece.columns[h:], False, piece.part)]) + + +class _Run: + def __init__(self, model: Any, source: Any, tables: Sequence[str], existing: Mapping[str, Mapping[str, Any]], + background: Background, hooks: Hooks) -> None: + self.model = model + self.hooks = hooks + #: structured 结构化输出；text 纯文本、从正文取 JSON。结构化输出用不了时这次起草后面都改走纯文本 + self.mode = "structured" + self.system = C.system_notes_source(source) + all_tables = (getattr(source, "schema_cache", None) or {}).get("tables") or {} + self.order = [t for t in dict.fromkeys(str(t) for t in tables) if isinstance(all_tables.get(t), dict)] + self.tables = { + name: _Table(name=name, meta=all_tables[name], existing=existing.get(name), + columns=[str(c["name"]) for c in all_tables[name].get("columns") or [] + if isinstance(c, dict) and c.get("name") is not None]) + for name in self.order + } + self.background = self._background(getattr(source, "description", None), background) + + # ---------------------------------------------------------------- 总流程 + + async def run(self) -> dict[str, ModelOutcome]: + await asyncio.gather(*(self._unit(unit) for unit in self._plan())) + return {name: self._outcome(t) for name, t in self.tables.items()} + + def _plan(self) -> list[list[str] | str]: + """分段：窄表按顺序拼成整表段（列表），宽表一个一项（表名）。各项按第一张表在 tables 里的位置排。""" + width = max(1, int(C.MODEL_BATCH_COLUMNS)) + most = max(1, int(C.MODEL_BATCH_TABLES)) + units: list[list[str] | str] = [] + pack: list[str] | None = None + used = 0 + for name in self.order: + st = self.tables[name] + n = len(st.columns) + if n > width: + st.parts = [_Piece(name, (), True, 1)] + [ + _Piece(name, tuple(st.columns[i:i + width]), False, k + 2) + for k, i in enumerate(range(0, n, width))] + units.append(name) + continue + st.parts = [_Piece(name, tuple(st.columns), True, 1)] + if pack is None or len(pack) >= most or used + n > width: + pack, used = [], 0 + units.append(pack) + pack.append(name) + used += n + return units + + async def _unit(self, unit: list[str] | str) -> None: + names = unit if isinstance(unit, list) else [unit] + if isinstance(unit, list): + await self._guarded([self.tables[n].parts[0] for n in unit]) + for name in unit: + await self._part_done(name) + else: + head, *rest = self.tables[unit].parts + await self._guarded([head]) # 表级那段先写：列分段拿它做背景 + await self._part_done(unit) + + async def column_part(piece: _Piece) -> None: + await self._guarded([piece]) + await self._part_done(unit) + + await asyncio.gather(*(column_part(p) for p in rest)) + for name in names: + await self._table_done(name) + + async def _guarded(self, pieces: list[_Piece]) -> None: + """起草一段；意料之外的错误只记这一段没写出来，不影响别的段和别的表。""" + try: + await self._segment(pieces, 0, False) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 + from app.core.errors import describe_exception, first_line + + logger.exception("起草数据目录时出错：%s", "、".join(p.table for p in pieces)) + self._fail(pieces, REASON_CRASH.format(line=first_line(e) or describe_exception(e))) + + # ---------------------------------------------------------------- 一段 + + async def _segment(self, pieces: list[_Piece], depth: int, followup: bool, *, mode: str | None = None, + last: bool = False) -> None: + """问一段，拿到的并进各表；按失败原因拆开、补缺、记下没写出来的。 + + depth：已经拆了几层（截断后只补剩下的也算一层）。followup：这一段是在补缺（补过就不再补，空着的算模型拿不准）。 + mode / last：拆不开的空壳换纯文本问的那一次（再空就记没写出来）。""" + reply = await self._ask(pieces, mode) + if reply.error is not None: + self._fail(pieces, reply.error) + return + gains = self._take(pieces, reply.entries) + anything = any(g.items for g in gains) + if last and anything and self.mode != "text": + # 结构化输出只给空壳、纯文本写出来了：这个模型的结构化输出靠不住，这次起草后面都走纯文本 + self.mode = "text" + logger.info("起草数据目录：结构化输出只给了空壳，纯文本写出来了，后面改走纯文本") + if reply.exhausted or reply.truncated: + rest = self._remainder(pieces, gains, truncated=True) + if not rest: + return + kind = "exhausted" if reply.exhausted else "truncated" + if not anything: + await self._split(rest, depth, kind, followup) + elif depth < int(C.MODEL_SPLIT_DEPTH): + await self._segment(rest, depth + 1, followup) # 捡回了一部分：只补剩下的 + else: + self._fail(rest, self._reason(kind, split=True)) + return + if not anything: + if followup: + return + if last: + self._fail(pieces, lambda p: self._reason("hollow", split=depth > 0, piece=p)) + return + await self._split(pieces, depth, "hollow", followup) + return + if followup: + return + rest = self._remainder(pieces, gains, truncated=False) + if rest: + await self._segment(rest, depth, True) + + async def _split(self, pieces: list[_Piece], depth: int, kind: str, followup: bool) -> None: + halves = _halves(pieces) if depth < int(C.MODEL_SPLIT_DEPTH) else None + if halves is not None: + first, second = halves + if (len(first) == len(second) == 1 and first[0].table == second[0].table and first[0].table_level): + # 同一张表按列拆、前一半带着表级的项：先写前一半，后一半拿它做背景 + await self._segment(first, depth + 1, followup) + await self._segment(second, depth + 1, followup) + return + results = await asyncio.gather(self._segment(first, depth + 1, followup), + self._segment(second, depth + 1, followup), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result + return + if kind == "hollow" and self.mode == "structured": + await self._segment(pieces, depth, followup, mode="text", last=True) + return + self._fail(pieces, lambda p: self._reason(kind, split=depth > 0, piece=p)) + + def _remainder(self, pieces: list[_Piece], gains: list[_Gain], *, truncated: bool) -> list[_Piece]: + """还没写出来的部分：回复里没有这张表、或者这张表一项可用内容都没有的整段重问；有内容的只问没给的列。 + 截断的回复里表级的项也没拿到的，表级的项一起重问（没截断时表级空着算模型拿不准）。""" + out: list[_Piece] = [] + for piece, gain in zip(pieces, gains): + if not gain.named or not gain.items: + out.append(piece) + continue + cols = tuple(c for c in piece.columns if c not in gain.answered) + table_level = piece.table_level and truncated and not gain.table_items + if cols or table_level: + out.append(_Piece(piece.table, cols, table_level, piece.part)) + return out + + # ---------------------------------------------------------------- 调用 + + async def _ask(self, pieces: list[_Piece], mode: str | None = None) -> _Reply: + """调一次模型（按原因重来：额度被拒、结构化输出不支持、网络错误……），交回读好的回复。""" + mode = mode or self.mode + retried = 0 + while True: + used = getattr(self.model, "budget", None) + try: + async with model_gate(): + await self._hook(self.hooks.call_started, [p.table for p in pieces]) + return await asyncio.wait_for(self._call(pieces, mode), timeout=C.MODEL_TIMEOUT_S) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 - 按原因分别处理 + kind = classify_failure(e) + if kind == "length": + return await self._length_reply(e) + await self._hook(self.hooks.usage, 0, 0) + if kind == "budget" and (getattr(self.model, "budget", None) != used or self._shrink()): + continue + if kind == "unsupported" and mode == "structured": + if self.mode == "structured": + logger.info("起草数据目录：结构化输出用不了（%s），后面改走纯文本", str(e)[:200]) + self.mode = mode = "text" + continue + if kind in ("transient", "timeout") and retried < int(C.MODEL_RETRIES): + delays = tuple(C.MODEL_RETRY_DELAYS_S) or (0.0,) + await _sleep(delays[min(retried, len(delays) - 1)]) + retried += 1 + continue + return _Reply(error=self._call_error(e, kind, retried)) + + async def _call(self, pieces: list[_Piece], mode: str) -> _Reply: + messages = self._messages(pieces, mode) + if mode == "structured": + runnable = self.model.with_structured_output(C.CATALOG_DRAFT_PROMPT_SCHEMA, include_raw=True) + raw, parsed = _unpack(await runnable.ainvoke(messages)) + await self._usage(raw) + return _structured_reply(raw, parsed) + raw = await self.model.ainvoke(messages) + await self._usage(raw) + return _text_reply(raw) + + async def _length_reply(self, e: BaseException) -> _Reply: + """LengthFinishReasonError：额度用完了。异常里带着那次的回复，能捡回来的照捡。""" + completion = getattr(e, "completion", None) + usage = getattr(completion, "usage", None) + await self._hook(self.hooks.usage, int(getattr(usage, "prompt_tokens", 0) or 0), + int(getattr(usage, "completion_tokens", 0) or 0)) + try: + text = completion.choices[0].message.content or "" # type: ignore[union-attr] + except (AttributeError, IndexError, TypeError): + text = "" + if text.strip(): + reply = _parse_text(text, True) + if reply is not None and reply.entries: + return _Reply(entries=reply.entries, truncated=True, exhausted=True) + return _Reply(exhausted=True) + + def _shrink(self) -> bool: + shrink = getattr(self.model, "shrink", None) + if callable(shrink) and shrink(): + logger.info("起草数据目录：服务方不接受这么大的输出额度，这次起草改用 %s", getattr(self.model, "budget", "")) + return True + return False + + def _budget(self) -> int: + return int(getattr(self.model, "budget", None) or C.MODEL_MAX_TOKENS) + + def _call_error(self, e: BaseException, kind: str, retried: int) -> str: + from app.core.errors import describe_exception, first_line + + tail = REASON_RETRIED.format(n=retried) if retried else "" + if kind == "timeout": + return REASON_TIMEOUT.format(seconds=_seconds(C.MODEL_TIMEOUT_S)) + tail + return REASON_CALL.format(line=first_line(e) or describe_exception(e)) + tail + + async def _usage(self, raw: Any) -> None: + meta = getattr(raw, "usage_metadata", None) or {} + await self._hook(self.hooks.usage, int(meta.get("input_tokens") or 0), int(meta.get("output_tokens") or 0)) + + # ---------------------------------------------------------------- 拿到的内容 + + def _match(self, pieces: list[_Piece], entries: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]: + """回复里的表对象按表名对到这一段的表上。表名照原样写的直接对上；带了 schema 前缀、大小写不一致的，只在恰好 + 对上一张时认。""" + names = [p.table for p in pieces] + out: dict[str, list[dict[str, Any]]] = {} + for entry in entries: + raw = entry.get("table") + if not isinstance(raw, str): + continue + name = raw if raw in names else None + if name is None: + short = raw.strip().rsplit(".", 1)[-1].strip('"`[] ') + found = [n for n in names if n == short or name_key(n) == name_key(short)] + name = found[0] if len(found) == 1 else None + if name is not None: + out.setdefault(name, []).append(entry) + return out + + def _take(self, pieces: list[_Piece], entries: list[dict[str, Any]]) -> list[_Gain]: + """回复里能用的并进各表（逐项校验，只留这一段请写的部分），返回每张表这次拿到了什么。""" + by_table = self._match(pieces, entries) + gains: list[_Gain] = [] + for piece in pieces: + st = self.tables[piece.table] + gain = _Gain() + for entry in by_table.get(piece.table, []): + gain.named = st.named = True + notes = _restrict(C._from_model(entry, st.meta, system=self.system), piece) + cols = entry.get("columns") if isinstance(entry.get("columns"), list) else [] + gain.answered |= {c.get("name") for c in cols if isinstance(c, dict) and c.get("name") in piece.columns} + slots = C._slots(notes) + gain.items += len(slots) + gain.table_items += sum(1 for key in slots if key[0] == "t") + self._merge(st, notes) + gains.append(gain) + return gains + + @staticmethod + def _merge(st: _Table, notes: Mapping[str, Any]) -> None: + for key, value in notes.items(): + if key == "columns": + cols = st.notes.setdefault("columns", {}) + for col, items in value.items(): + target = cols.setdefault(col, {}) + for name, item in items.items(): + target.setdefault(name, item) + st.got_columns.add(col) + else: + st.notes.setdefault(key, value) + st.got_table = True + + def _fail(self, pieces: list[_Piece], reason: str | Callable[[_Piece], str]) -> None: + """这几段没写出来：还没拿到内容的列（和表级的项）记下原因。""" + for piece in pieces: + st = self.tables[piece.table] + why = reason(piece) if callable(reason) else reason + for col in piece.columns: + if col not in st.got_columns: + st.failed.setdefault(col, why) + if piece.table_level and not st.got_table: + st.failed.setdefault(_TABLE, why) + + def _reason(self, kind: str, *, split: bool, piece: _Piece | None = None) -> str: + if kind == "exhausted": + base = REASON_EXHAUSTED.format(budget=self._budget()) + elif kind == "truncated": + base = REASON_TRUNCATED + else: + st = self.tables[piece.table] if piece is not None else None + if st is not None and st.whole(piece): + base = C.MODEL_EMPTY_REASON if st.named else REASON_MISSING + else: + base = REASON_HOLLOW + return base + (REASON_SPLIT if split else "") + + # ---------------------------------------------------------------- 收尾 + + def _outcome(self, st: _Table) -> ModelOutcome: + items = len(C._slots(st.notes)) + failed_cols = [c for c in st.columns if c in st.failed and c not in st.got_columns] + failed_table = _TABLE in st.failed and not st.got_table + total = len(st.parts) or 1 + if not items: + reasons = list(dict.fromkeys(st.failed.values())) + reason = "；".join(reasons[:2]) if reasons else (C.MODEL_EMPTY_REASON if st.named else REASON_MISSING) + return ModelOutcome("failed", {}, 0, reason, total, st.parts_done) + notes = copy.deepcopy(st.notes) + if failed_cols or failed_table: + return ModelOutcome("partial", notes, items, self._partial_text(st, failed_cols, failed_table), total, + st.parts_done) + return ModelOutcome("ok", notes, items, None, total, st.parts_done) + + @staticmethod + def _partial_text(st: _Table, failed_cols: list[str], failed_table: bool) -> str: + """只写出一部分时说清哪部分没写出来、为什么、怎么办。""" + keys = ([_TABLE] if failed_table else []) + failed_cols + reason = "；".join(list(dict.fromkeys(st.failed[k] for k in keys))[:2]) + n = len(st.columns) + if len(st.parts) > 1: + part_of = {c: p.part for p in st.parts for c in p.columns} + bad = sorted({part_of[c] for c in failed_cols if c in part_of} | ({1} if failed_table else set())) + head = f"这张表 {n} 列，分 {len(st.parts)} 段起草，第 {'、'.join(str(b) for b in bad)} 段没写出来" + head += f"（共 {len(failed_cols)} 列）" if failed_cols else "（表级的项）" + elif failed_cols and failed_table: + head = f"这张表 {n} 列，有表级的项和 {len(failed_cols)} 列没写出来" + elif failed_cols: + head = f"这张表 {n} 列，有 {len(failed_cols)} 列没写出来" + else: + head = f"这张表 {n} 列，表级的项没写出来" + return f"{head}：{reason}。可以只重试这张表" + + async def _part_done(self, name: str) -> None: + self.tables[name].parts_done += 1 + await self._hook(self.hooks.part_done, name) + + async def _table_done(self, name: str) -> None: + st = self.tables[name] + if st.finished: + return + st.finished = True + await self._hook(self.hooks.table_done, name, self._outcome(st)) + + @staticmethod + async def _hook(fn: Callable[..., Awaitable[None]] | None, *args: Any) -> None: + if fn is None: + return + try: + await fn(*args) + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 - 进度回调出错不影响起草 + logger.exception("起草数据目录的进度回调出错") + + # ---------------------------------------------------------------- 提示词 + + @staticmethod + def _background(description: Any, background: Background) -> str: + labels = C.CATALOG_DRAFT_PROMPT_LABELS + lines = [] + for key, text in (("description", description), ("guidance", background.guidance), ("note", background.note)): + text = text.strip() if isinstance(text, str) else "" + if text: + lines.append(labels[key].format(text=text)) + return C.CATALOG_DRAFT_BACKGROUND_PROMPT.format(lines="\n".join(lines)) if lines else "" + + def _messages(self, pieces: list[_Piece], mode: str) -> list[Any]: + from langchain_core.messages import HumanMessage, SystemMessage + + text = C.CATALOG_DRAFT_PROMPT.format(n=len(pieces), tables="\n\n".join(self._describe(p) for p in pieces)) + if self.background: + text = f"{self.background}\n\n{text}" + if mode == "text": + text = f"{text}\n\n{C.CATALOG_DRAFT_TEXT_PROMPT}" + return [SystemMessage(content=C.CATALOG_DRAFT_SYSTEM), HumanMessage(content=text)] + + def _describe(self, piece: _Piece) -> str: + """一段里一张表的结构描述：列、类型、主键、外键、唯一约束、注释；不含任何数据值。 + + 整表：照旧的完整描述。只写表级的项：全部列只给名字和类型。表级的项加几列：完整描述，表名下一行点名这几列。 + 只写几列：表注释、已起草的表级内容做背景，只列这几列。""" + labels = C.CATALOG_DRAFT_PROMPT_LABELS + st = self.tables[piece.table] + meta = st.meta + lines = [(labels["view"] if meta.get("is_view") else labels["table"]).format(name=st.name)] + brief = piece.table_level and not piece.columns + if not st.whole(piece): + if brief: + lines.append(labels["only_table"]) + elif piece.table_level: + lines.append(labels["table_and_columns"].format(n=len(piece.columns), cols="、".join(piece.columns))) + else: + lines.append(labels["only_columns"].format(n=len(piece.columns), total=len(st.columns))) + comment = (meta.get("comment") or "").strip() + if comment: + lines.append((labels["system"] if self.system else labels["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")) + if brief: + lines.append(labels["brief_columns"].format(n=len(st.columns))) + wanted = set(st.columns) + else: + if not piece.table_level: + drafted = C._notes_lines({k: v for k, v in st.notes.items() if k in C.TABLE_FIELDS}, meta=meta, + system_notes=False) + if drafted: + lines.append(labels["drafted"]) + lines += drafted + lines.append(labels["columns"]) + wanted = set(st.columns) if piece.table_level else set(piece.columns) + for col in meta.get("columns") or []: + if isinstance(col, dict) and col.get("name") in wanted: + lines.append(self._column_line(col, pk, fk_of, brief=brief)) + if piece.table_level: + for group in meta.get("unique") or []: + if isinstance(group, list) and group: + lines.append(labels["unique"].format(cols="、".join(group))) + confirmed = self._confirmed_lines(st, piece) + if confirmed: + lines.append(labels["confirmed"]) + lines += confirmed + return "\n".join(lines) + + def _column_line(self, col: Mapping[str, Any], pk: set[str], fk_of: Mapping[str, str], *, brief: bool) -> str: + labels = C.CATALOG_DRAFT_PROMPT_LABELS + marks = [] + if col["name"] in pk: + marks.append(labels["pk"]) + if not brief and not col.get("nullable", True): + marks.append(labels["not_null"]) + if col["name"] in fk_of: + marks.append(labels["fk"].format(target=fk_of[col["name"]])) + if not brief and (text := (col.get("comment") or "").strip()): + marks.append((labels["col_system"] if self.system else labels["col_comment"]).format(text=text)) + line = labels["column"].format(name=col["name"], type=col.get("type") or "") + return line + ("，" + "，".join(marks) if marks else "") + + @staticmethod + def _confirmed_lines(st: _Table, piece: _Piece) -> list[str]: + """已确认、有确证的项（以此为准的背景）：整表全给；其余给表级的项和这一段的列，关系只在要写表级的项时给。""" + confirmed = C._confirmed_only(st.existing) + if not st.whole(piece): + wanted = set(piece.columns) + kept = {key: item for key, item in C._slots(confirmed).items() + if key[0] == "t" or (key[0] == "c" and key[1] in wanted) or (key[0] == "r" and piece.table_level)} + confirmed = C._assemble(kept) + return C._notes_lines(confirmed, meta=st.meta, system_notes=False) + + +__all__ = ["BUDGET_WORDS", "Background", "Hooks", "ModelOutcome", "REASON_CALL", "REASON_EXHAUSTED", "REASON_FORMAT", + "REASON_HOLLOW", "REASON_MISSING", "REASON_RETRIED", "REASON_SPLIT", "REASON_TIMEOUT", "REASON_TRUNCATED", + "classify_failure", "model_gate", "parts_for", "run_model_draft"] diff --git a/backend/tests/fixtures/catalog/fake_model.py b/backend/tests/fixtures/catalog/fake_model.py index 42bfdec5..1ffd796d 100644 --- a/backend/tests/fixtures/catalog/fake_model.py +++ b/backend/tests/fixtures/catalog/fake_model.py @@ -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: diff --git a/backend/tests/test_catalog_draft.py b/backend/tests/test_catalog_draft.py index 96162564..85e98d74 100644 --- a/backend/tests/test_catalog_draft.py +++ b/backend/tests/test_catalog_draft.py @@ -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")} diff --git a/backend/tests/test_catalog_draft_engine.py b/backend/tests/test_catalog_draft_engine.py new file mode 100644 index 00000000..a1ac2272 --- /dev/null +++ b/backend/tests/test_catalog_draft_engine.py @@ -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