mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 04:31:58 +00:00
合入数据目录起草改造（后端）：起草模型设置组 catalog_draft、数据目录与配方 AI 起草共用取模型；新起草引擎保留思考、额度 32768 被拒退回 8192，40 列 / 4 张表分段，宽表表级加列分段，按失败原因拆小、补缺、截断捡回、改走纯文本、退避重试与止损；起草要求进提示词；后台起草任务与 draft-jobs 接口、启动标中断不续跑；目录清单 last_draft、active_job、guidance
This commit is contained in:
ThearchyHelios
+199
-14
@@ -7,11 +7,15 @@
|
||||
- POST /catalog/{table}/patch/preview 目录修改提案的预览（只读）：对着当前目录算改前、改后
|
||||
- POST /catalog/{table}/patch 保存目录修改提案（乐观锁），记为人工确认
|
||||
- GET /catalog/{table}/impact 影响面：引用这张表的已发布、受管模板（data/catalog_impact.py）
|
||||
- POST /catalog/draft 同步起草（注释、外键、命名推断，可选模型）
|
||||
- POST /catalog/draft 同步起草（注释、外键、命名推断，可选模型）。界面改用下面的后台起草任务，这个留着
|
||||
- PUT /catalog/guidance 数据源级的起草要求（使用者补充的背景和要求，起草时交给模型）
|
||||
- POST /catalog/draft-jobs 建后台起草任务（202）；GET /draft-jobs/active、/latest、/{job_id} 读任务；
|
||||
POST /draft-jobs/{job_id}/cancel 停止，/retry 只重试没写好的表（data/catalog_jobs.py）
|
||||
- POST /catalog/profile 同步剖析：对业务库发少量只读查询，核对关系、取码值候选（data/catalog_profile.py）
|
||||
|
||||
规则都在 data/catalog.py，这里只做取源、转换形状和把异常翻成状态码：版本不符 409，结构不合规、审阅路径
|
||||
不对 422，源或表不存在 404。所有写操作记录请求头 X-Actor 的署名（自报、未认证）。剖析没开启、正在剖析 409。
|
||||
起草任务的 409 带机读码（api/coded.py）：同一个源已有排队或进行中的任务、起草模型用不了、要停止的任务已经结束。
|
||||
|
||||
上传表格的源按当前快照的表结构回答（table_versions.resolve_source），和查询实际用的库是同一个版本。
|
||||
"""
|
||||
@@ -19,13 +23,17 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
from fastapi import APIRouter, Body, Depends, Header, HTTPException
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.api.coded import PROFILE_BUSY, SCHEMA_MISSING, SNAPSHOT_TAMPERED, CodedHTTPException
|
||||
from app.api.coded import (
|
||||
CATALOG_DRAFT_BUSY, CATALOG_DRAFT_FINISHED, CATALOG_DRAFT_MODEL_UNAVAILABLE, PROFILE_BUSY, SCHEMA_MISSING,
|
||||
SNAPSHOT_TAMPERED, CodedHTTPException,
|
||||
)
|
||||
from app.api.runs import actor_of
|
||||
from app.data import catalog, catalog_impact, catalog_profile, introspect, table_versions
|
||||
from app.data import catalog, catalog_impact, catalog_jobs, catalog_profile, introspect, table_versions
|
||||
from app.data.engine import SnapshotTampered
|
||||
from app.db.base import get_session
|
||||
from app.db.models import DataSource
|
||||
@@ -64,10 +72,26 @@ class CatalogPatchIn(CatalogPatchPreviewIn):
|
||||
class CatalogDraftIn(BaseModel):
|
||||
#: 要起草的表；不给或为空时按使用次数取前 20 张
|
||||
tables: list[str] | None = Field(default=None, max_length=200)
|
||||
#: 是否请助手的模型起草中文名、粒度、列的含义等
|
||||
#: 是否请起草模型（设置里的起草模型，没设沿用助手的模型）写中文名、粒度、列的含义等
|
||||
use_model: bool = False
|
||||
|
||||
|
||||
class DraftJobIn(BaseModel):
|
||||
#: 要起草的表；不给或为空时按使用次数取前 20 张。最多 500 张（超了 422，说明是中文）
|
||||
tables: list[str] | None = None
|
||||
#: 是否请起草模型写中文名、粒度、列的含义等
|
||||
use_model: bool = False
|
||||
#: 这一次的补充说明，和数据源级的起草要求一起交给模型（最多 1000 字）
|
||||
note: str | None = None
|
||||
|
||||
|
||||
#: PUT /catalog/guidance 的请求 {guidance: string}：数据源级的起草要求，空串（或 null）清除，最多 2000 字。必须写
|
||||
#: 这一项：表名恰好叫 guidance 时，整份提交那张表目录的请求也落到这个路径上，不能把它当成「清除起草要求」。
|
||||
#: 没写、不是文字都回 422，说明是中文的一句话（同超长）
|
||||
GUIDANCE_MISSING = "请求里要写起草要求（guidance 一项，一段文字；要清除起草要求请传空串）"
|
||||
GUIDANCE_NOT_TEXT = "起草要求要是一段文字（要清除起草要求请传空串）"
|
||||
|
||||
|
||||
class CatalogProfileIn(BaseModel):
|
||||
#: 要剖析的表；不给或为空时按使用次数取前 10 张目录里有待核实关系的表
|
||||
tables: list[str] | None = Field(default=None, max_length=200)
|
||||
@@ -100,7 +124,7 @@ def _item_value(notes: dict[str, Any], name: str) -> tuple[Any, str | None]:
|
||||
|
||||
|
||||
def _summary_row(name: str, meta: dict[str, Any] | None, entry: catalog.CatalogEntry | None,
|
||||
usage: int) -> dict[str, Any]:
|
||||
usage: int, last_draft: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
notes = entry.notes if entry else {}
|
||||
label, label_status = _item_value(notes, "label")
|
||||
kind, _ = _item_value(notes, "kind")
|
||||
@@ -120,6 +144,8 @@ def _summary_row(name: str, meta: dict[str, Any] | None, entry: catalog.CatalogE
|
||||
"version": entry.version if entry else 0,
|
||||
"updated_at": _iso(entry.updated_at) if entry else None,
|
||||
"updated_by": entry.updated_by if entry else None,
|
||||
# 这张表最近一次起草任务里的那一项：{status, job_id, finished_at, error, model_error}，没起草过是 null
|
||||
"last_draft": last_draft,
|
||||
}
|
||||
|
||||
|
||||
@@ -180,15 +206,18 @@ async def list_catalog(source_id: str, session: AsyncSession = Depends(get_sessi
|
||||
说明有没有截断，schema_total 是数据库里一共几张表（含视图），界面据此说明「只探查了前 200 张」。老缓存没记
|
||||
总数时按探查到的张数算。
|
||||
"""
|
||||
_, source = await _resolved(session, source_id)
|
||||
row, source = await _resolved(session, source_id)
|
||||
cache = source.schema_cache or {}
|
||||
tables = cache.get("tables") or {}
|
||||
entries = await catalog.read_catalog(session, source.id)
|
||||
usage = await catalog.table_usage(session, source)
|
||||
drafts = await catalog_jobs.last_drafts(session, source.id)
|
||||
order = {name: i for i, name in enumerate(tables)}
|
||||
rows = [_summary_row(name, meta, entries.get(name), usage.get(name, 0)) for name, meta in tables.items()]
|
||||
rows = [_summary_row(name, meta, entries.get(name), usage.get(name, 0), drafts.get(name))
|
||||
for name, meta in tables.items()]
|
||||
rows.sort(key=lambda r: (-r["usage"], order[r["table_name"]]))
|
||||
rows += [_summary_row(name, None, entry, 0) for name, entry in sorted(entries.items()) if name not in tables]
|
||||
rows += [_summary_row(name, None, entry, 0, drafts.get(name))
|
||||
for name, entry in sorted(entries.items()) if name not in tables]
|
||||
return {
|
||||
"tables": rows,
|
||||
"system_notes": catalog.system_notes_source(source),
|
||||
@@ -196,6 +225,10 @@ async def list_catalog(source_id: str, session: AsyncSession = Depends(get_sessi
|
||||
"schema_note": None if tables else introspect.why_empty(cache),
|
||||
"schema_truncated": bool(cache.get("truncated")),
|
||||
"schema_total": max(_int(cache.get("total")), len(tables)),
|
||||
# 数据源级的起草要求（没有是 null）
|
||||
"guidance": catalog_jobs.guidance_of(row),
|
||||
# 排队或进行中的起草任务 {id, status, total, done}，界面据此显示进度；没有是 null
|
||||
"active_job": await catalog_jobs.job_brief(session, await catalog_jobs.active_job(session, source.id)),
|
||||
}
|
||||
|
||||
|
||||
@@ -207,6 +240,156 @@ async def get_catalog_table(source_id: str, table: str,
|
||||
return await _detail(session, source, key, entry)
|
||||
|
||||
|
||||
async def _source_row(session: AsyncSession, source_id: str) -> DataSource:
|
||||
row = await session.get(DataSource, source_id)
|
||||
if row is None:
|
||||
raise HTTPException(404, "数据源不存在，可能已被删除")
|
||||
return row
|
||||
|
||||
|
||||
@router.put("/{source_id}/catalog/guidance")
|
||||
async def put_catalog_guidance(source_id: str, body: dict[str, Any] = Body(...),
|
||||
x_actor: str | None = Header(default=None),
|
||||
session: AsyncSession = Depends(get_session)) -> Any:
|
||||
"""存数据源级的起草要求：去首尾空白，空串清除，最多 2000 字（超了 422）。上传表格的源存在它自己的数据源记录上。
|
||||
|
||||
要登记在 PUT /catalog/{table} 前面：不然「guidance」会被当成表名。反过来，表名恰好叫 guidance 时，整份提交那张表
|
||||
目录的请求（{notes, if_version}，没有 guidance 这一项）也落到这里：照旧当作那张表的整份提交。"""
|
||||
if "guidance" not in body and "notes" in body:
|
||||
try:
|
||||
table_payload = CatalogPutIn.model_validate(body)
|
||||
except ValidationError as e:
|
||||
raise RequestValidationError(e.errors()) from e
|
||||
return await put_catalog_table(source_id, "guidance", table_payload, x_actor=x_actor, session=session)
|
||||
if "guidance" not in body:
|
||||
raise HTTPException(422, GUIDANCE_MISSING)
|
||||
guidance = body["guidance"]
|
||||
if guidance is not None and not isinstance(guidance, str):
|
||||
raise HTTPException(422, GUIDANCE_NOT_TEXT)
|
||||
row = await _source_row(session, source_id)
|
||||
try:
|
||||
value = await catalog_jobs.save_guidance(session, row, guidance)
|
||||
except catalog_jobs.TextTooLong as e:
|
||||
raise HTTPException(422, str(e)) from e
|
||||
return {"guidance": value}
|
||||
|
||||
|
||||
async def _ensure_model(session: AsyncSession) -> None:
|
||||
"""起草模型用不了（没有接入、停用、只有演示接入……）：409，不建任务。说法同同步起草时的 model_error。"""
|
||||
try:
|
||||
await catalog_jobs.check_model(session)
|
||||
except catalog.CatalogModelUnavailable as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_MODEL_UNAVAILABLE) from e
|
||||
|
||||
|
||||
def _no_schema(cache: dict[str, Any]) -> HTTPException:
|
||||
return HTTPException(409, f"{introspect.why_empty(cache)}，无法起草数据目录。请先在数据源卡片上点「探查结构」")
|
||||
|
||||
|
||||
@router.post("/{source_id}/catalog/draft-jobs", status_code=202)
|
||||
async def create_draft_job(source_id: str, payload: DraftJobIn, x_actor: str | None = Header(default=None),
|
||||
session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""建后台起草任务，马上返回（202，任务是排队中）。同一个源已有排队或进行中的任务 409 catalog_draft_busy（界面改去
|
||||
GET /draft-jobs/active 显示那一个）；用模型而起草模型用不了 409 catalog_draft_model_unavailable；没有表结构 409。
|
||||
表结构里没有的表照样建进任务、状态 skipped。"""
|
||||
if payload.tables and len(payload.tables) > catalog_jobs.MAX_TABLES:
|
||||
raise HTTPException(422, f"一次最多起草 {catalog_jobs.MAX_TABLES} 张表，当前 {len(payload.tables)} 张，"
|
||||
"请分几次起草")
|
||||
try:
|
||||
note = catalog_jobs.clean_text(payload.note, catalog_jobs.NOTE_MAX, "本次说明")
|
||||
except catalog_jobs.TextTooLong as e:
|
||||
raise HTTPException(422, str(e)) from e
|
||||
row, source = await _resolved(session, source_id)
|
||||
try:
|
||||
await catalog_jobs.ensure_idle(session, row)
|
||||
except catalog_jobs.DraftBusy as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_BUSY) from e
|
||||
cache = source.schema_cache or {}
|
||||
if not cache.get("tables"):
|
||||
raise _no_schema(cache)
|
||||
if payload.use_model:
|
||||
await _ensure_model(session)
|
||||
try:
|
||||
job, tasks = await catalog_jobs.create_job(session, row, source, tables=payload.tables or None,
|
||||
use_model=payload.use_model, note=note, actor=actor_of(x_actor))
|
||||
except catalog_jobs.DraftBusy as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_BUSY) from e
|
||||
catalog_jobs.manager.start(job.id)
|
||||
return {"job": catalog_jobs.job_dict(job, tasks)}
|
||||
|
||||
|
||||
@router.get("/{source_id}/catalog/draft-jobs/active")
|
||||
async def active_draft_job(source_id: str, session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""排队或进行中的那个起草任务，没有是 null。"""
|
||||
row = await _source_row(session, source_id)
|
||||
return {"job": await catalog_jobs.job_out(session, await catalog_jobs.active_job(session, row.id))}
|
||||
|
||||
|
||||
@router.get("/{source_id}/catalog/draft-jobs/latest")
|
||||
async def latest_draft_job(source_id: str, session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""最近建的一个起草任务（任何状态），没有是 null。"""
|
||||
row = await _source_row(session, source_id)
|
||||
return {"job": await catalog_jobs.job_out(session, await catalog_jobs.latest_job(session, row.id))}
|
||||
|
||||
|
||||
async def _job_or_404(session: AsyncSession, source_id: str, job_id: str) -> tuple[DataSource, Any]:
|
||||
row = await _source_row(session, source_id)
|
||||
job = await catalog_jobs.load(session, row.id, job_id)
|
||||
if job is None:
|
||||
raise HTTPException(404, "起草任务不存在，可能属于别的数据源，或者已随数据源删除")
|
||||
return row, job
|
||||
|
||||
|
||||
@router.get("/{source_id}/catalog/draft-jobs/{job_id}")
|
||||
async def get_draft_job(source_id: str, job_id: str, session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
_, job = await _job_or_404(session, source_id, job_id)
|
||||
return {"job": await catalog_jobs.job_out(session, job)}
|
||||
|
||||
|
||||
@router.post("/{source_id}/catalog/draft-jobs/{job_id}/cancel")
|
||||
async def cancel_draft_job(source_id: str, job_id: str,
|
||||
session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""停止：排队的直接停止；进行中的停掉正在进行的模型调用，没开始和没写完的表记停止，已经写进目录的不回滚。
|
||||
已经结束的 409 catalog_draft_finished。返回停止之后的任务。"""
|
||||
row, job = await _job_or_404(session, source_id, job_id)
|
||||
owner = row.id
|
||||
try:
|
||||
await catalog_jobs.cancel(session, job)
|
||||
except catalog_jobs.DraftFinished as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_FINISHED) from e
|
||||
session.expire_all() # 后台在别的会话里收的尾：重新读
|
||||
return {"job": await catalog_jobs.job_out(session, await catalog_jobs.load(session, owner, job_id))}
|
||||
|
||||
|
||||
@router.post("/{source_id}/catalog/draft-jobs/{job_id}/retry", status_code=202)
|
||||
async def retry_draft_job(source_id: str, job_id: str, x_actor: str | None = Header(default=None),
|
||||
session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""重试：新任务，只含原任务里没写好（failed、partial）和停止了（cancelled）的表；用不用模型、本次说明沿用，起草要求
|
||||
取数据源当前的。有排队或进行中的任务 409 catalog_draft_busy；没有可重试的表 422；起草模型用不了 409。"""
|
||||
_, old = await _job_or_404(session, source_id, job_id)
|
||||
row, source = await _resolved(session, source_id)
|
||||
try:
|
||||
await catalog_jobs.ensure_idle(session, row)
|
||||
except catalog_jobs.DraftBusy as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_BUSY) from e
|
||||
try:
|
||||
tables = await catalog_jobs.retry_tables(session, old)
|
||||
except catalog_jobs.NothingToRetry as e:
|
||||
raise HTTPException(422, str(e)) from e
|
||||
cache = source.schema_cache or {}
|
||||
if not cache.get("tables"):
|
||||
raise _no_schema(cache)
|
||||
if old.use_model:
|
||||
await _ensure_model(session)
|
||||
try:
|
||||
job, tasks = await catalog_jobs.create_job(session, row, source, tables=tables, use_model=old.use_model,
|
||||
note=old.note, actor=actor_of(x_actor), retry_of=old.id)
|
||||
except catalog_jobs.DraftBusy as e:
|
||||
raise CodedHTTPException(409, str(e), CATALOG_DRAFT_BUSY) from e
|
||||
catalog_jobs.manager.start(job.id)
|
||||
return {"job": catalog_jobs.job_dict(job, tasks)}
|
||||
|
||||
|
||||
@router.put("/{source_id}/catalog/{table}")
|
||||
async def put_catalog_table(source_id: str, table: str, payload: CatalogPutIn,
|
||||
x_actor: str | None = Header(default=None),
|
||||
@@ -228,17 +411,19 @@ async def put_catalog_table(source_id: str, table: str, payload: CatalogPutIn,
|
||||
async def draft_catalog(source_id: str, payload: CatalogDraftIn,
|
||||
x_actor: str | None = Header(default=None),
|
||||
session: AsyncSession = Depends(get_session)) -> dict[str, Any]:
|
||||
"""同步起草。返回每张表新增、更新、删除了几项；模型用不了或某张表的模型起草失败，照实写在结果里。
|
||||
"""同步起草（界面改用后台起草任务，这个接口留着，返回形状不变）。引擎同后台任务（分段、按失败原因处理），起草要求
|
||||
取数据源级的那份。返回每张表新增、更新、删除了几项；模型用不了或某张表的模型起草失败，照实写在结果里。
|
||||
|
||||
model_items：模型给了这张表几项可用内容（没用模型、模型这部分失败时为 null）。界面据此区分「模型给了内容、
|
||||
只是和现有目录一致」（model_items > 0、三个数都是 0）和「模型没给内容」（model_error 有值）。
|
||||
"""
|
||||
_, source = await _resolved(session, source_id)
|
||||
row, source = await _resolved(session, source_id)
|
||||
cache = source.schema_cache or {}
|
||||
if not cache.get("tables"):
|
||||
raise HTTPException(409, f"{introspect.why_empty(cache)}，无法起草数据目录。请先在数据源卡片上点「探查结构」")
|
||||
raise _no_schema(cache)
|
||||
report = await catalog.draft_catalog(session, source, tables=payload.tables or None,
|
||||
use_model=payload.use_model, actor=actor_of(x_actor))
|
||||
use_model=payload.use_model, actor=actor_of(x_actor),
|
||||
guidance=catalog_jobs.guidance_of(row))
|
||||
rows = [{"table_name": r.table, "added": r.added, "updated": r.updated, "removed": r.removed,
|
||||
"version": r.version, "error": r.error, "model_error": r.model_error, "model_items": r.model_items}
|
||||
for r in report.tables]
|
||||
|
||||
@@ -19,6 +19,11 @@ DATASOURCE_INACTIVE = "datasource_inactive"
|
||||
SCHEMA_MISSING = "schema_missing"
|
||||
PROFILE_BUSY = "profile_busy"
|
||||
SNAPSHOT_TAMPERED = "snapshot_tampered"
|
||||
#: 数据目录的后台起草任务（api/catalog.py 的 /catalog/draft-jobs）：同一个源已有排队或进行中的任务（409，界面改去
|
||||
#: GET /draft-jobs/active 显示那一个）、起草模型用不了（409，不建任务）、要停止的任务已经结束了（409）
|
||||
CATALOG_DRAFT_BUSY = "catalog_draft_busy"
|
||||
CATALOG_DRAFT_MODEL_UNAVAILABLE = "catalog_draft_model_unavailable"
|
||||
CATALOG_DRAFT_FINISHED = "catalog_draft_finished"
|
||||
|
||||
|
||||
class CodedHTTPException(HTTPException):
|
||||
|
||||
@@ -27,8 +27,8 @@ from app.data.catalog_profile import profile_settings_rejection
|
||||
from app.data import introspect as introspect_mod
|
||||
from app.data import table_versions
|
||||
from app.data.engine import (
|
||||
MASK_COLUMNS_OPTION, MAX_QUERY_TIMEOUT_S, QUERY_TIMEOUT_OPTION, SUPPORTED_KINDS, SnapshotTampered, build_url,
|
||||
engine_args, engines, mask_columns_problem, query_timeout_problem,
|
||||
CATALOG_DRAFT_GUIDANCE_OPTION, MASK_COLUMNS_OPTION, MAX_QUERY_TIMEOUT_S, QUERY_TIMEOUT_OPTION, SUPPORTED_KINDS,
|
||||
SnapshotTampered, build_url, engine_args, engines, mask_columns_problem, query_timeout_problem,
|
||||
)
|
||||
from app.db.base import get_session, new_id
|
||||
from app.db.models import DataSource, ImportStaging, SourceSnapshot, TableImport, TableRecipe
|
||||
@@ -124,6 +124,22 @@ def _import_mode(row: DataSource) -> str | None:
|
||||
return "simple" if _is_upload(row) else None
|
||||
|
||||
|
||||
#: options 里由别的接口维护的项：数据源表单不显示、不提交它们，整份回写 options 时原样留着（同模型接入 extra 里的
|
||||
#: last_check）。数据目录的起草要求只经 PUT /catalog/guidance 改（api/catalog.py）：表单把认不得的项放进「高级参数」
|
||||
#: 的单行输入框，多行的要求会被压成一行
|
||||
_MANAGED_OPTIONS = (CATALOG_DRAFT_GUIDANCE_OPTION,)
|
||||
|
||||
|
||||
def _shown_options(options: dict[str, Any] | None) -> dict[str, Any]:
|
||||
return {k: v for k, v in (options or {}).items() if k not in _MANAGED_OPTIONS}
|
||||
|
||||
|
||||
def _with_managed(options: dict[str, Any], saved: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""表单提交的 options + 存着的、由别的接口维护的项（提交里带了也不认）。"""
|
||||
kept = {k: v for k, v in (saved or {}).items() if k in _MANAGED_OPTIONS}
|
||||
return {**_shown_options(options), **kept}
|
||||
|
||||
|
||||
def _to_out(row: DataSource, snapshot: dict[str, Any] | None = None,
|
||||
recipe_info: dict[str, Any] | None = None) -> DataSourceOut:
|
||||
cache = row.schema_cache or {}
|
||||
@@ -137,7 +153,7 @@ def _to_out(row: DataSource, snapshot: dict[str, Any] | None = None,
|
||||
schema_error=str(cache.get("error") or "") if cache.get("failed") else "",
|
||||
available_schemas=list(cache.get("available_schemas") or []),
|
||||
id=row.id, name=row.name, kind=row.kind, host=row.host, port=row.port,
|
||||
database=row.database, username=row.username, options=row.options or {},
|
||||
database=row.database, username=row.username, options=_shown_options(row.options),
|
||||
readonly=row.readonly, description=row.description, enabled=row.enabled,
|
||||
password_masked=mask(row.password),
|
||||
has_password=bool(row.password),
|
||||
@@ -429,7 +445,8 @@ async def create_source(
|
||||
_refuse_reserved(payload.kind, payload.database)
|
||||
|
||||
row = DataSource(
|
||||
**payload.model_dump(exclude={"password"}),
|
||||
**payload.model_dump(exclude={"password", "options"}),
|
||||
options=_with_managed(payload.options, None),
|
||||
password=encrypt(payload.password),
|
||||
)
|
||||
_settle_oracle(row, None, None)
|
||||
@@ -472,6 +489,7 @@ async def update_source(
|
||||
raise HTTPException(422, _UPLOAD_DIR_READONLY)
|
||||
if "options" in data:
|
||||
_refuse_bad_options(data["options"])
|
||||
data["options"] = _with_managed(data["options"] or {}, row.options)
|
||||
if "password" in data:
|
||||
# 空串表示清空，非空表示换新的
|
||||
data["password"] = encrypt(data["password"]) if data["password"] else None
|
||||
|
||||
@@ -14,6 +14,7 @@ from app.api.coded import CodedHTTPException
|
||||
from app.core.errors import explain, not_configured, raw
|
||||
from app.core.config import settings as app_settings
|
||||
from app.core.crypto import encrypt, mask
|
||||
from app.data.draft_model import CATALOG_DRAFT_SETTING_KEY, setting_pair
|
||||
from app.db.base import get_session
|
||||
from app.db.models import Provider, Setting
|
||||
from app.engine.judge import JUDGE_DEFAULTS, MAX_TOKENS as JUDGE_MAX_TOKENS
|
||||
@@ -569,6 +570,39 @@ async def verify_model_spec(session: AsyncSession, node_config: dict[str, Any] |
|
||||
return ModelSpec(provider=provider, model=model, max_tokens=JUDGE_MAX_TOKENS, thinking="off")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 起草模型：数据目录的「用模型起草」和按配方导入的 AI 起草用哪个模型
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
#: 设置里 catalog_draft 这一组：{provider, model}。都留空时沿用助手的模型，再没有就是默认接入（取值顺序见
|
||||
#: data/draft_model.py，两处起草共用）。只放接入和模型：额度、思考由两处起草各自定
|
||||
CATALOG_DRAFT_GROUP = CATALOG_DRAFT_SETTING_KEY
|
||||
CATALOG_DRAFT_DEFAULTS: dict[str, Any] = {"provider": None, "model": None}
|
||||
#: 设置页「起草模型」一节各项的叫法：报错时写它，不写键名
|
||||
CATALOG_DRAFT_LABEL = {"provider": "起草模型 · 接入", "model": "起草模型 · 模型"}
|
||||
|
||||
|
||||
def catalog_draft_problem(value: Any) -> str | None:
|
||||
"""设置页写进来的 catalog_draft 这一组有什么问题，没问题返回 None。点名的接入在不在由 _provider_problem 查。"""
|
||||
if not isinstance(value, dict):
|
||||
return f"「起草模型」设置格式有误，需要是一个对象，当前为「{value}」"
|
||||
for key, item in value.items():
|
||||
if key not in CATALOG_DRAFT_DEFAULTS:
|
||||
return (f"「起草模型」设置中不支持「{key}」，可设置的项："
|
||||
+ "、".join(f"「{CATALOG_DRAFT_LABEL[k]}」" for k in CATALOG_DRAFT_DEFAULTS))
|
||||
if item is not None and not isinstance(item, str):
|
||||
# 「跟随助手的模型」是设置页接入下拉框空选项的原话
|
||||
what = "模型接入的名称" if key == "provider" else "模型 ID"
|
||||
return f"「{CATALOG_DRAFT_LABEL[key]}」需要填写{what}，留空表示跟随助手的模型，当前为「{item}」"
|
||||
return None
|
||||
|
||||
|
||||
def catalog_draft_values(stored: Any) -> dict[str, Any]:
|
||||
"""库里存的 catalog_draft 一组 → 补齐默认、规范化（去掉首尾空白，空串和存坏了的当没填）。"""
|
||||
provider, model = setting_pair(stored)
|
||||
return {"provider": provider, "model": model}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 通知：定时任务推送 Webhook 时拼运行链接用的对外访问地址
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -661,6 +695,8 @@ DEFAULT_SETTINGS: dict[str, Any] = {
|
||||
VERIFY_GROUP: verify_values(None),
|
||||
# 通知：定时任务的 Webhook 里拼运行链接用的对外访问地址
|
||||
NOTIFY_GROUP: notify_values(None),
|
||||
# 起草模型：数据目录起草、按配方导入的 AI 起草用的模型，留空沿用助手的模型
|
||||
CATALOG_DRAFT_GROUP: catalog_draft_values(None),
|
||||
}
|
||||
|
||||
|
||||
@@ -717,6 +753,7 @@ async def get_settings(session: AsyncSession = Depends(get_session)) -> dict[str
|
||||
merged[JUDGE_GROUP] = judge_values(stored.get(JUDGE_GROUP))
|
||||
merged[VERIFY_GROUP] = verify_values(stored.get(VERIFY_GROUP))
|
||||
merged[NOTIFY_GROUP] = notify_values(stored.get(NOTIFY_GROUP))
|
||||
merged[CATALOG_DRAFT_GROUP] = catalog_draft_values(stored.get(CATALOG_DRAFT_GROUP))
|
||||
for key, value in stored.items():
|
||||
merged.setdefault(key, value)
|
||||
return merged
|
||||
@@ -731,8 +768,8 @@ PROVIDER_INVALID = "settings_provider_invalid"
|
||||
|
||||
|
||||
def _provider_fields() -> dict[tuple[str, str], tuple[str, str | None]]:
|
||||
"""设置里点名模型接入的项：(组, 键) → (这一项在界面上的叫法, 同一组模型那一项的叫法)。门控、裁判、核查的叫法和
|
||||
设置页一致；助手的模型由「编排」页助手栏的「助手使用的模型」单独设定（PUT /api/copilot/model），只有这一个下拉框
|
||||
"""设置里点名模型接入的项：(组, 键) → (这一项在界面上的叫法, 同一组模型那一项的叫法)。门控、裁判、核查、起草模型的
|
||||
叫法和设置页一致；助手的模型由「编排」页助手栏的「助手使用的模型」单独设定（PUT /api/copilot/model），只有这一个下拉框
|
||||
（模型那一项为 None），这里只防经设置接口写进去的那一份。"""
|
||||
from app.api.copilot import COPILOT_SETTING_KEY # 在这里导入：copilot 模块很大，设置模块被引擎各处导入
|
||||
from app.providers.factory import ASSISTANT_MODEL_FIELD
|
||||
@@ -741,6 +778,7 @@ def _provider_fields() -> dict[tuple[str, str], tuple[str, str | None]]:
|
||||
("run", "tool_gate_provider"): ("门控模型 · 接入", "门控模型 · 模型"),
|
||||
(JUDGE_GROUP, "provider"): (JUDGE_LABEL["provider"], JUDGE_LABEL["model"]),
|
||||
(VERIFY_GROUP, "provider"): (VERIFY_SETTING_LABEL["provider"], VERIFY_SETTING_LABEL["model"]),
|
||||
(CATALOG_DRAFT_GROUP, "provider"): (CATALOG_DRAFT_LABEL["provider"], CATALOG_DRAFT_LABEL["model"]),
|
||||
(COPILOT_SETTING_KEY, "provider"): (ASSISTANT_MODEL_FIELD, None),
|
||||
}
|
||||
|
||||
@@ -798,6 +836,9 @@ async def put_settings(
|
||||
raise HTTPException(422, problem)
|
||||
if VERIFY_GROUP in payload.values and (problem := verify_problem(payload.values[VERIFY_GROUP])):
|
||||
raise HTTPException(422, problem)
|
||||
if CATALOG_DRAFT_GROUP in payload.values and (
|
||||
problem := catalog_draft_problem(payload.values[CATALOG_DRAFT_GROUP])):
|
||||
raise HTTPException(422, problem)
|
||||
if NOTIFY_GROUP in payload.values:
|
||||
if problem := notify_problem(payload.values[NOTIFY_GROUP]):
|
||||
raise HTTPException(422, problem)
|
||||
|
||||
+225
-228
@@ -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:
|
||||
@@ -1401,15 +1405,28 @@ def draft_structure(source: Any, table: str) -> TableDraft:
|
||||
# 起草：模型（可选，按表分批，结构化输出）
|
||||
# ==========================================================================
|
||||
|
||||
#: 每批最多几张表、合计多少列。一批太大模型容易漏表、输出被截断；太小则调用次数多
|
||||
#: 每次调用最多几张表、合计多少列（data/catalog_draft.py 按它分段）。实测每列约 50 个输出 token（不算思考），
|
||||
#: 会思考的模型思考也算在输出额度里：一段太大，想完就没额度写正文。列数超过 MODEL_BATCH_COLUMNS 的宽表拆成
|
||||
#: 「表级一段 + 列分段」，每个列分段也不超过这个列数
|
||||
MODEL_BATCH_TABLES = 4
|
||||
MODEL_BATCH_COLUMNS = 120
|
||||
#: 同时发出的批数
|
||||
MODEL_BATCH_COLUMNS = 40
|
||||
#: 全进程同时进行的模型调用最多几个（所有起草任务、同步起草共用）
|
||||
MODEL_CONCURRENCY = 3
|
||||
#: 一批最多等多久
|
||||
MODEL_TIMEOUT_S = 120
|
||||
#: 输出额度。一批四张表的目录远用不完；给太大的额度有的服务方会直接拒收
|
||||
MODEL_MAX_TOKENS = 8192
|
||||
#: 单次调用最多等多久（秒）。会思考的模型想得久，一段 40 列常常要几分钟
|
||||
MODEL_TIMEOUT_S = 420
|
||||
#: 输出额度。会思考的模型（经 OpenAI 兼容网关的推理模型等）思考也算在里面：8192 时上百列的表常常想完就没额度
|
||||
#: 写正文。保留思考（不关），加大额度；服务方嫌额度太大时退回 MODEL_FALLBACK_MAX_TOKENS（同 copilot、recipe_ai）
|
||||
MODEL_MAX_TOKENS = 32768
|
||||
MODEL_FALLBACK_MAX_TOKENS = 8192
|
||||
#: 一段没写出来（额度用完、只给空壳）时拆成两半再问：最多拆几层、拆到最少几列为止
|
||||
MODEL_SPLIT_DEPTH = 3
|
||||
MODEL_SPLIT_MIN_COLUMNS = 5
|
||||
#: 止损：连续这么多段拆到底仍一项没写出来，这次起草后面的段不再拆小重试（一段 40 列拆到底要 15 次调用，模型明显
|
||||
#: 用不了时不该照这个倍数烧额度）
|
||||
MODEL_GIVE_UP_AFTER = 2
|
||||
#: 网络错误、429、5xx、超时最多重试几次，每次重试前等多久（秒，依次取用）
|
||||
MODEL_RETRIES = 2
|
||||
MODEL_RETRY_DELAYS_S = (10.0, 30.0)
|
||||
#: 不指定表时起草使用次数最多的前几张
|
||||
DRAFT_DEFAULT_TABLES = 20
|
||||
|
||||
@@ -1435,11 +1452,27 @@ status 状态；attribute 属性）。
|
||||
规矩：
|
||||
- 只写数据事实，不写计算公式，不写指标怎么算。
|
||||
- 拿不准的字段留空，不要猜。列名、表名照原样写，不要翻译或改写。
|
||||
- 请你写的每一列都在 columns 里写一项，拿不准的字段可以留空，但不要漏掉列：漏掉的列会被再问一次。
|
||||
- 标着「系统说明」的内容是导入时按核对结果生成的，以它为准：不要重复，也不要与之矛盾。
|
||||
- 标着「已确认」的目录项由人工确认过，以它为准。"""
|
||||
- 标着「已确认」的目录项由人工确认过，以它为准。
|
||||
- 表名下一行写着「这次只写……」的表，只写它说的那部分：只写表级的项时不写列，只写列出的几列时不写表级的项、\
|
||||
也不写别的列。
|
||||
- 使用者补充的背景和要求帮你理解业务、统一用词；它和这里的规矩冲突时，以这里的规矩为准。
|
||||
- 思考时不必把要输出的内容逐列草拟一遍，想清楚后直接按格式输出：思考和输出共用同一份长度额度，额度用完时还没输出的\
|
||||
内容就全丢了。"""
|
||||
|
||||
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": "表类型", "keys": ["列名"], \
|
||||
"business_date": {"column": "列名", "rule": "一句规则", "timezone": "时区"}, \
|
||||
"columns": [{"name": "列名", "label": "中文名", "meaning": "含义", "unit": "单位", "measure": "度量类型"}]}]}
|
||||
表类型写 fact、dimension、snapshot、log、config 之一；度量类型写 flow、stock、ratio、identifier、status、attribute 之一。"""
|
||||
|
||||
#: 拼表结构描述用的固定片段（写给模型）
|
||||
CATALOG_DRAFT_PROMPT_LABELS = {
|
||||
"table": "表 {name}",
|
||||
@@ -1455,6 +1488,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] = {
|
||||
@@ -1499,99 +1542,68 @@ 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）。一次起草（一个起草任务、一次同步起草）共用一个，退回之后不再先失败一次。
|
||||
|
||||
只转发起草用到的两个方法：with_structured_output 每次都取当前那个模型的，ainvoke 同理。"""
|
||||
|
||||
def __init__(self, big: Any, small: Any | None, model_name: str = "") -> None:
|
||||
self.model, self._small = big, small
|
||||
self.model_name = model_name
|
||||
self.budget = MODEL_MAX_TOKENS
|
||||
|
||||
def shrink(self) -> bool:
|
||||
"""换成小额度的模型。已经换过、没有小的，返回 False。"""
|
||||
if self._small is None or self.model is self._small:
|
||||
return False
|
||||
self.model, self.budget = self._small, MODEL_FALLBACK_MAX_TOKENS
|
||||
return True
|
||||
|
||||
def with_structured_output(self, schema: Any, **kwargs: Any) -> Any:
|
||||
return self.model.with_structured_output(schema, **kwargs)
|
||||
|
||||
async def ainvoke(self, messages: list[Any]) -> Any:
|
||||
return await self.model.ainvoke(messages)
|
||||
|
||||
|
||||
async def resolve_draft_model(session: AsyncSession) -> tuple[Any, str]:
|
||||
"""(模型, 模型 id)：用设置里助手的模型，关掉思考、温度 0。没有可用的模型抛 CatalogModelUnavailable。
|
||||
"""(模型, 模型 id)：用哪个接入、哪个模型按 data/draft_model.py 的顺序（设置里的起草模型 > 助手的模型 > 默认接入），
|
||||
和按配方导入的 AI 起草共用。没有可用的模型抛 CatalogModelUnavailable。
|
||||
|
||||
不关思考（spec 里不传 thinking，按接入默认）、温度 0、额度 MODEL_MAX_TOKENS，服务方嫌大时退回
|
||||
MODEL_FALLBACK_MAX_TOKENS（DraftModel）。单次调用的时限 MODEL_TIMEOUT_S 也交给 SDK，免得它先按全局的 5 分钟断开。
|
||||
SDK 自己不重试（max_retries=0）：429、5xx、网络错误、超时由起草引擎按失败原因重试（最多 MODEL_RETRIES 次），
|
||||
不然两层叠起来一段最多发 9 次请求，报错却写「已重试 2 次」，网关正在限流时还会放大流量。
|
||||
只建模型对象，不发请求。
|
||||
|
||||
和 recipe_ai.resolve_draft_model 一样不用 get_chat_model：它在没有接入时静默退回演示模型，拿演示模型
|
||||
起草出来的目录会被当成真的写进库里。数据层不导入接口层，设置键取 recipe_ai 里那份（与助手的同值，
|
||||
那边有测试核对）。
|
||||
起草出来的目录会被当成真的写进库里。
|
||||
"""
|
||||
from app.data.recipe_ai import COPILOT_SETTING_KEY
|
||||
from app.db.models import Setting
|
||||
from app.providers.factory import (
|
||||
ModelSpec, ProviderNotConfigured, assistant_unavailable, build_chat_model, resolve_provider,
|
||||
)
|
||||
from app.data.draft_model import DraftModelUnavailable, draft_unavailable, resolve_draft_provider
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, build_chat_model
|
||||
|
||||
row = await session.get(Setting, COPILOT_SETTING_KEY)
|
||||
saved = (row.value if row else None) or {}
|
||||
spec = ModelSpec(provider=saved.get("provider") or None, model=saved.get("model") or None,
|
||||
max_tokens=MODEL_MAX_TOKENS, temperature=0, thinking="off")
|
||||
try:
|
||||
# 存的接入名不存在、已停用，只存了模型名、提供它的接入已停用：resolve_provider 都会报出来
|
||||
provider = await resolve_provider(session, spec.provider, spec.model)
|
||||
except ProviderNotConfigured as e:
|
||||
raise CatalogModelUnavailable(assistant_unavailable(e)) from e
|
||||
if provider is None:
|
||||
raise CatalogModelUnavailable("未配置模型接入，无法由助手起草数据目录。请到「设置 → 模型接入」添加")
|
||||
if provider.kind == "mock":
|
||||
raise CatalogModelUnavailable("当前只有演示用的模型接入，无法用于起草数据目录。"
|
||||
"请到「设置 → 模型接入」添加真实的模型接入")
|
||||
provider, choice = await resolve_draft_provider(
|
||||
session, none_reason="未配置模型接入，无法由模型起草数据目录。请到「设置 → 模型接入」添加",
|
||||
mock_reason="当前只有演示用的模型接入，无法用于起草数据目录。请到「设置 → 模型接入」添加真实的模型接入")
|
||||
except DraftModelUnavailable as e:
|
||||
raise CatalogModelUnavailable(str(e)) from e
|
||||
spec = ModelSpec(provider=choice.provider, model=choice.model, max_tokens=MODEL_MAX_TOKENS, temperature=0,
|
||||
timeout=MODEL_TIMEOUT_S, max_retries=0)
|
||||
try:
|
||||
model = build_chat_model(provider, spec)
|
||||
big = build_chat_model(provider, spec)
|
||||
small = build_chat_model(provider, spec.model_copy(update={"max_tokens": MODEL_FALLBACK_MAX_TOKENS}))
|
||||
except ProviderNotConfigured as e:
|
||||
raise CatalogModelUnavailable(assistant_unavailable(e)) from e
|
||||
return model, spec.model or provider.default_model or ""
|
||||
raise CatalogModelUnavailable(draft_unavailable(e, choice.origin)) from e
|
||||
model_id = spec.model or provider.default_model or ""
|
||||
return DraftModel(big, small, model_id), model_id
|
||||
|
||||
|
||||
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 = "模型没有给出这张表的可用内容"
|
||||
@@ -1621,39 +1633,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
|
||||
@@ -1708,65 +1687,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
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
@@ -1776,7 +1720,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
|
||||
@@ -1785,9 +1730,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
|
||||
@@ -1802,35 +1752,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:
|
||||
@@ -1840,41 +1857,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
|
||||
|
||||
|
||||
@@ -2239,12 +2232,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
@@ -0,0 +1,645 @@
|
||||
"""数据目录的后台起草任务，以及数据源级的起草要求。
|
||||
|
||||
以前起草是一次同步长请求（POST /catalog/draft，界面每 12 张表一个请求）：一个请求出错整个对话框就停，关掉对话框
|
||||
之后没有任何记录哪些表失败、为什么。现在：
|
||||
|
||||
- **建任务**（create_job）：同一个数据源同时只有一个排队或进行中的任务（DraftBusy；部分唯一索引兜底）。不指定表时
|
||||
按使用次数取前 20 张；表结构里没有的表照样建进任务、状态 skipped。起草要求（数据源级）和本次说明随任务记下快照。
|
||||
- **在后台起草**（manager.start）：进程内的 asyncio 任务，引擎是 data/catalog_draft.py（分段、按失败原因处理）。
|
||||
每张表一结束就并入目录（catalog.write_draft，规则和同步起草一样：人工确认、驳回过的不动，乐观锁冲突重读重并）；
|
||||
逐表的状态、到第几段、计数、用量随时落库，接口读的是库里的状态。
|
||||
- **停止**（cancel）：排队的直接 cancelled；进行中的停掉正在进行的模型调用，没开始和没写完的表记 cancelled，已经写进
|
||||
目录的不回滚，任务最后是 cancelled。
|
||||
- **重试**（retry_tables + create_job）：新任务，只含原任务里 failed、partial、cancelled 的表；用不用模型、本次说明沿用，
|
||||
起草要求取数据源当前的。
|
||||
- **服务重启**：启动时（startup）把排队、进行中的任务标成 interrupted，没结束的表记 cancelled、写明「服务重启，起草
|
||||
中断」。不自动续跑：复制出来的数据目录（检查、验收）启动时不能自己去调模型；界面让人点「重试」续上。关停时
|
||||
（manager.shutdown）停掉进行中的任务，同样标成 interrupted。
|
||||
|
||||
任务的库写入一律经 _JobRun._write：一条一条来，shield 起来——停止任务时不会打断写到一半的库操作，收尾前等它们写完。
|
||||
|
||||
**起草要求**存在数据源 options 的 catalog_draft_guidance 里（上传表格的源存在它自己的 DataSource 行上），一段文字，
|
||||
去首尾空白、最多 GUIDANCE_MAX 字，空串等于没有。只经 PUT /catalog/guidance 改：数据源接口不显示它、整份回写 options
|
||||
时留着（api/datasources.py），拼连接串时摘掉（data/engine.py）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.data import catalog, catalog_draft
|
||||
from app.data.engine import CATALOG_DRAFT_GUIDANCE_OPTION
|
||||
from app.db.base import SessionLocal, new_id, utcnow
|
||||
from app.db.models import CatalogDraftJob, CatalogDraftTask, DataSource
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: 任务还没结束的状态（同一个源只能有一个）、结束了的状态
|
||||
ACTIVE = ("queued", "running")
|
||||
FINISHED = ("succeeded", "partial", "failed", "cancelled", "interrupted")
|
||||
#: 逐表的状态，以及其中算「已结束」的、重试时要带上的
|
||||
TASK_STATUSES = ("pending", "running", "ok", "partial", "failed", "skipped", "cancelled")
|
||||
TASK_DONE = ("ok", "partial", "failed", "skipped", "cancelled")
|
||||
RETRYABLE = ("failed", "partial", "cancelled")
|
||||
_OPEN = ("pending", "running")
|
||||
|
||||
#: 一个任务最多几张表；起草要求、本次说明最多几个字
|
||||
MAX_TABLES = 500
|
||||
GUIDANCE_MAX = 2000
|
||||
NOTE_MAX = 1000
|
||||
#: 目录清单上「这张表最近一次起草」只看最近这么多个任务
|
||||
LAST_DRAFT_JOBS = 100
|
||||
#: 停止进行中的任务时最多等它收尾多久（秒）；关停服务时另给（shutdown 的 grace）
|
||||
STOP_GRACE_S = 10.0
|
||||
|
||||
INTERRUPTED_ERROR = "服务重启，起草中断"
|
||||
UNFINISHED_ERROR = "起草结束时这张表还没有写完，请重试这张表"
|
||||
WRITE_ERROR = "写入数据目录时出错：{line}"
|
||||
JOB_ERROR = "起草任务出错：{line}"
|
||||
SOURCE_GONE = "数据源不存在，可能已被删除"
|
||||
|
||||
|
||||
class DraftBusy(Exception):
|
||||
"""这个数据源已经有排队或进行中的起草任务。message 给人看。"""
|
||||
|
||||
|
||||
class DraftFinished(Exception):
|
||||
"""要停止的任务已经结束了。"""
|
||||
|
||||
|
||||
class NothingToRetry(Exception):
|
||||
"""原任务里没有需要重试的表。"""
|
||||
|
||||
|
||||
class TextTooLong(ValueError):
|
||||
"""起草要求、本次说明超长。message 给人看。"""
|
||||
|
||||
|
||||
class _JobFailed(Exception):
|
||||
"""任务整体出错（数据源没了、快照没了、模型用不了）：message 就是任务的 error。"""
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# 起草要求、本次说明
|
||||
# ==========================================================================
|
||||
|
||||
|
||||
def clean_text(value: Any, limit: int, label: str) -> str | None:
|
||||
"""去首尾空白；空串（和不是字符串的）等于没有；超过 limit 个字抛 TextTooLong。"""
|
||||
text = value.strip() if isinstance(value, str) else ""
|
||||
if len(text) > limit:
|
||||
raise TextTooLong(f"{label}最多 {limit} 字，当前 {len(text)} 字，请删减后再保存")
|
||||
return text or None
|
||||
|
||||
|
||||
def guidance_of(row: Any) -> str | None:
|
||||
"""数据源级的起草要求（没有是 None）。"""
|
||||
raw = (getattr(row, "options", None) or {}).get(CATALOG_DRAFT_GUIDANCE_OPTION)
|
||||
return raw.strip() or None if isinstance(raw, str) else None
|
||||
|
||||
|
||||
async def save_guidance(session: AsyncSession, row: DataSource, value: Any) -> str | None:
|
||||
"""存起草要求，返回存下的（清除了是 None）。超长抛 TextTooLong，什么都不改。"""
|
||||
text = clean_text(value, GUIDANCE_MAX, "起草要求")
|
||||
options = {k: v for k, v in (row.options or {}).items() if k != CATALOG_DRAFT_GUIDANCE_OPTION}
|
||||
if text:
|
||||
options[CATALOG_DRAFT_GUIDANCE_OPTION] = text
|
||||
row.options = options
|
||||
await session.commit()
|
||||
return text
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# 读任务
|
||||
# ==========================================================================
|
||||
|
||||
|
||||
def _iso(value: Any) -> str | None:
|
||||
return value.isoformat() if value is not None else None
|
||||
|
||||
|
||||
async def active_job(session: AsyncSession, source_id: str) -> CatalogDraftJob | None:
|
||||
"""排队或进行中的那个（同一个源最多一个）。"""
|
||||
return (await session.execute(
|
||||
select(CatalogDraftJob).where(CatalogDraftJob.source_id == source_id, CatalogDraftJob.status.in_(ACTIVE))
|
||||
.order_by(CatalogDraftJob.created_at.desc()).limit(1)
|
||||
)).scalar_one_or_none()
|
||||
|
||||
|
||||
async def latest_job(session: AsyncSession, source_id: str) -> CatalogDraftJob | None:
|
||||
"""最近建的一个（任何状态）。"""
|
||||
return (await session.execute(
|
||||
select(CatalogDraftJob).where(CatalogDraftJob.source_id == source_id)
|
||||
.order_by(CatalogDraftJob.created_at.desc()).limit(1)
|
||||
)).scalar_one_or_none()
|
||||
|
||||
|
||||
async def load(session: AsyncSession, source_id: str, job_id: str) -> CatalogDraftJob | None:
|
||||
"""这个源的这个任务；不存在、不属于这个源都是 None。"""
|
||||
job = await session.get(CatalogDraftJob, job_id)
|
||||
return job if job is not None and job.source_id == source_id else None
|
||||
|
||||
|
||||
async def tasks_of(session: AsyncSession, job_id: str) -> list[CatalogDraftTask]:
|
||||
return list((await session.execute(
|
||||
select(CatalogDraftTask).where(CatalogDraftTask.job_id == job_id).order_by(CatalogDraftTask.seq)
|
||||
)).scalars())
|
||||
|
||||
|
||||
def task_dict(task: CatalogDraftTask) -> dict[str, Any]:
|
||||
return {
|
||||
"table_name": task.table_name,
|
||||
"status": task.status,
|
||||
"columns": task.column_count,
|
||||
"parts_total": task.parts_total,
|
||||
"parts_done": task.parts_done,
|
||||
"added": task.added,
|
||||
"updated": task.updated,
|
||||
"removed": task.removed,
|
||||
"version": task.version,
|
||||
"model_items": task.model_items,
|
||||
"error": task.error,
|
||||
"model_error": task.model_error,
|
||||
"started_at": _iso(task.started_at),
|
||||
"finished_at": _iso(task.finished_at),
|
||||
}
|
||||
|
||||
|
||||
def _counts(statuses: list[str]) -> dict[str, int]:
|
||||
counts = dict.fromkeys(TASK_STATUSES, 0)
|
||||
for status in statuses:
|
||||
counts[status] = counts.get(status, 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def job_dict(job: CatalogDraftJob, tasks: list[CatalogDraftTask]) -> dict[str, Any]:
|
||||
"""接口返回的任务（CatalogDraftJob）：逐表的状态按起草顺序；计数、合计都按逐表的算。"""
|
||||
counts = _counts([t.status for t in tasks])
|
||||
return {
|
||||
"id": job.id,
|
||||
"source_id": job.source_id,
|
||||
"status": job.status,
|
||||
"use_model": bool(job.use_model),
|
||||
"model": job.model,
|
||||
"guidance": job.guidance,
|
||||
"note": job.note,
|
||||
"created_at": _iso(job.created_at),
|
||||
"started_at": _iso(job.started_at),
|
||||
"finished_at": _iso(job.finished_at),
|
||||
"created_by": job.created_by,
|
||||
"retry_of": job.retry_of,
|
||||
"total": len(tasks),
|
||||
"done": sum(counts.get(s, 0) for s in TASK_DONE),
|
||||
"counts": counts,
|
||||
"added": sum(t.added or 0 for t in tasks),
|
||||
"updated": sum(t.updated or 0 for t in tasks),
|
||||
"removed": sum(t.removed or 0 for t in tasks),
|
||||
"usage": {"calls": job.calls or 0, "input_tokens": job.input_tokens or 0,
|
||||
"output_tokens": job.output_tokens or 0},
|
||||
"error": job.error,
|
||||
"tasks": [task_dict(t) for t in tasks],
|
||||
}
|
||||
|
||||
|
||||
async def job_out(session: AsyncSession, job: CatalogDraftJob | None) -> dict[str, Any] | None:
|
||||
return None if job is None else job_dict(job, await tasks_of(session, job.id))
|
||||
|
||||
|
||||
async def job_brief(session: AsyncSession, job: CatalogDraftJob | None) -> dict[str, Any] | None:
|
||||
"""目录清单顶上的进度：{id, status, total, done}。"""
|
||||
if job is None:
|
||||
return None
|
||||
rows = (await session.execute(
|
||||
select(CatalogDraftTask.status, func.count()).where(CatalogDraftTask.job_id == job.id)
|
||||
.group_by(CatalogDraftTask.status)
|
||||
)).all()
|
||||
counts = {status: n for status, n in rows}
|
||||
return {"id": job.id, "status": job.status, "total": sum(counts.values()),
|
||||
"done": sum(counts.get(s, 0) for s in TASK_DONE)}
|
||||
|
||||
|
||||
async def last_drafts(session: AsyncSession, source_id: str) -> dict[str, dict[str, Any]]:
|
||||
"""每张表最近一次起草任务里的那一项：{表名: {status, job_id, finished_at, error, model_error}}。任务还在进行，
|
||||
就是进行中的状态。只看最近 LAST_DRAFT_JOBS 个任务；每张表取哪一项在库里挑好（清单页起草时会反复取）。"""
|
||||
recent = (select(CatalogDraftJob.id).where(CatalogDraftJob.source_id == source_id)
|
||||
.order_by(CatalogDraftJob.created_at.desc()).limit(LAST_DRAFT_JOBS)).scalar_subquery()
|
||||
rank = func.row_number().over(partition_by=CatalogDraftTask.table_name,
|
||||
order_by=(CatalogDraftJob.created_at.desc(), CatalogDraftTask.seq)).label("rank")
|
||||
inner = (select(CatalogDraftTask.table_name, CatalogDraftTask.status, CatalogDraftTask.job_id,
|
||||
CatalogDraftTask.finished_at, CatalogDraftTask.error, CatalogDraftTask.model_error, rank)
|
||||
.join(CatalogDraftJob, CatalogDraftJob.id == CatalogDraftTask.job_id)
|
||||
.where(CatalogDraftJob.id.in_(recent))).subquery()
|
||||
rows = (await session.execute(select(inner).where(inner.c.rank == 1))).mappings()
|
||||
return {row["table_name"]: {"status": row["status"], "job_id": row["job_id"],
|
||||
"finished_at": _iso(row["finished_at"]), "error": row["error"],
|
||||
"model_error": row["model_error"]} for row in rows}
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# 建任务、重试、停止
|
||||
# ==========================================================================
|
||||
|
||||
|
||||
def busy_message(row: Any) -> str:
|
||||
return f"数据源「{getattr(row, 'name', '')}」正在起草数据目录，请等这次起草结束（或先停止它）再开始新的起草"
|
||||
|
||||
|
||||
async def ensure_idle(session: AsyncSession, row: Any) -> None:
|
||||
"""这个源没有排队或进行中的任务；有就抛 DraftBusy。"""
|
||||
if await active_job(session, row.id) is not None:
|
||||
raise DraftBusy(busy_message(row))
|
||||
|
||||
|
||||
async def check_model(session: AsyncSession) -> None:
|
||||
"""起草模型能不能用：不能用抛 catalog.CatalogModelUnavailable（只建模型对象，不发请求）。"""
|
||||
await catalog.resolve_draft_model(session)
|
||||
|
||||
|
||||
async def create_job(session: AsyncSession, row: Any, source: Any, *, tables: list[str] | None, use_model: bool,
|
||||
note: str | None, actor: str | None,
|
||||
retry_of: str | None = None) -> tuple[CatalogDraftJob, list[CatalogDraftTask]]:
|
||||
"""建一个排队的任务（还没开始，调用方接着 manager.start）。row 是数据源记录（起草要求取它的），source 是读表结构
|
||||
用的源（上传源是绑定当前快照的视图）。有排队或进行中的任务抛 DraftBusy。"""
|
||||
await ensure_idle(session, row)
|
||||
cache = getattr(source, "schema_cache", None) or {}
|
||||
all_tables = cache.get("tables") or {}
|
||||
wanted = list(tables) if tables else await catalog.default_tables(session, source)
|
||||
job = CatalogDraftJob(id=new_id(), source_id=row.id, status="queued", use_model=bool(use_model),
|
||||
guidance=guidance_of(row), note=note, created_by=actor, retry_of=retry_of)
|
||||
tasks: list[CatalogDraftTask] = []
|
||||
seen: set[str] = set()
|
||||
now = utcnow()
|
||||
for name in dict.fromkeys(str(t) for t in wanted):
|
||||
key = catalog.resolve_table_name(cache, name)
|
||||
if key is None:
|
||||
tasks.append(CatalogDraftTask(job_id=job.id, seq=len(tasks), table_name=name, status="skipped",
|
||||
error=catalog.missing_table_error(name), finished_at=now))
|
||||
continue
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
columns = len((all_tables.get(key) or {}).get("columns") or [])
|
||||
tasks.append(CatalogDraftTask(job_id=job.id, seq=len(tasks), table_name=key, status="pending",
|
||||
column_count=columns,
|
||||
parts_total=catalog_draft.parts_for(columns) if use_model else None))
|
||||
session.add(job)
|
||||
try:
|
||||
await session.flush()
|
||||
session.add_all(tasks)
|
||||
await session.commit()
|
||||
except IntegrityError as e:
|
||||
# 部分唯一索引：另一个请求同时给这个源建了任务
|
||||
await session.rollback()
|
||||
raise DraftBusy(busy_message(row)) from e
|
||||
return job, tasks
|
||||
|
||||
|
||||
async def retry_tables(session: AsyncSession, job: CatalogDraftJob) -> list[str]:
|
||||
"""重试要带上的表：原任务里 failed、partial、cancelled 的，按原来的顺序。没有抛 NothingToRetry。"""
|
||||
names = [t.table_name for t in await tasks_of(session, job.id) if t.status in RETRYABLE]
|
||||
if not names:
|
||||
raise NothingToRetry("这次起草里没有失败、只写出一部分或停止了的表，不需要重试")
|
||||
return names
|
||||
|
||||
|
||||
async def cancel(session: AsyncSession, job: CatalogDraftJob) -> None:
|
||||
"""停止一个任务。排队的直接 cancelled；进行中的交给 manager 停（停掉模型调用、没写完的表记 cancelled）。
|
||||
已经结束的抛 DraftFinished。"""
|
||||
if job.status in FINISHED:
|
||||
raise DraftFinished("这次起草已经结束了，不需要停止")
|
||||
if job.status == "queued":
|
||||
now = utcnow()
|
||||
claimed = await session.execute(
|
||||
update(CatalogDraftJob).where(CatalogDraftJob.id == job.id, CatalogDraftJob.status == "queued")
|
||||
.values(status="cancelled", finished_at=now).execution_options(synchronize_session=False))
|
||||
if claimed.rowcount == 1:
|
||||
await session.execute(
|
||||
update(CatalogDraftTask).where(CatalogDraftTask.job_id == job.id, CatalogDraftTask.status.in_(_OPEN))
|
||||
.values(status="cancelled", finished_at=now).execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
return
|
||||
await session.rollback() # 刚被后台认领、开始了：按进行中的停
|
||||
if await manager.stop(job.id, "user", grace=STOP_GRACE_S):
|
||||
return
|
||||
# 库里说在进行，这个进程里却没有在跑它（记录对不上）：直接记停止
|
||||
await _close(session, job.id, "cancelled", None, "cancelled", None)
|
||||
|
||||
|
||||
async def _close(session: AsyncSession, job_id: str, job_status: str, job_error: str | None, task_status: str,
|
||||
task_error: str | None) -> None:
|
||||
"""收尾：还没结束的表记 task_status，任务（还没结束的话）记 job_status。"""
|
||||
now = utcnow()
|
||||
await session.execute(
|
||||
update(CatalogDraftTask).where(CatalogDraftTask.job_id == job_id, CatalogDraftTask.status.in_(_OPEN))
|
||||
.values(status=task_status, error=task_error, finished_at=now).execution_options(synchronize_session=False))
|
||||
await session.execute(
|
||||
update(CatalogDraftJob).where(CatalogDraftJob.id == job_id, CatalogDraftJob.status.in_(ACTIVE))
|
||||
.values(status=job_status, error=job_error, finished_at=now).execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def startup() -> int:
|
||||
"""服务启动时：排队、进行中的任务标成 interrupted，没结束的表记 cancelled、写明服务重启。不自动续跑。返回标了几个。"""
|
||||
async with SessionLocal() as session:
|
||||
ids = [job_id for job_id in (await session.execute(
|
||||
select(CatalogDraftJob.id).where(CatalogDraftJob.status.in_(ACTIVE)))).scalars()
|
||||
if not manager.running(job_id)]
|
||||
for job_id in ids:
|
||||
await _close(session, job_id, "interrupted", INTERRUPTED_ERROR, "cancelled", INTERRUPTED_ERROR)
|
||||
if ids:
|
||||
logger.warning("服务重启：%d 个数据目录起草任务没有做完，已标为中断，可在数据目录页重试", len(ids))
|
||||
return len(ids)
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# 在后台起草
|
||||
# ==========================================================================
|
||||
|
||||
|
||||
class _JobRun:
|
||||
"""一个任务在后台的一次执行。"""
|
||||
|
||||
def __init__(self, job_id: str) -> None:
|
||||
self.job_id = job_id
|
||||
self.task: asyncio.Task[None] | None = None
|
||||
#: 为什么停：user 使用者点了停止，shutdown 服务关停（标中断）
|
||||
self.stop_reason: str | None = None
|
||||
self._writes: set[asyncio.Task[Any]] = set()
|
||||
self._lock = asyncio.Lock()
|
||||
self._started: set[str] = set()
|
||||
self.source: Any = None
|
||||
self.actor: str | None = None
|
||||
self.at = ""
|
||||
|
||||
# ---------------------------------------------------------------- 库写入
|
||||
|
||||
async def _write(self, fn: Callable[..., Awaitable[Any]], *args: Any) -> Any:
|
||||
"""一条库写入：同一个任务的写入一条一条来；shield 起来，停止任务时不会打断写到一半的库操作。"""
|
||||
async def go() -> Any:
|
||||
async with self._lock:
|
||||
async with SessionLocal() as session:
|
||||
return await fn(session, *args)
|
||||
|
||||
task = asyncio.ensure_future(go())
|
||||
self._writes.add(task)
|
||||
task.add_done_callback(self._writes.discard)
|
||||
return await asyncio.shield(task)
|
||||
|
||||
async def _settle(self) -> None:
|
||||
"""等手上的库写入都写完（停止、出错收尾之前）。"""
|
||||
while self._writes:
|
||||
await asyncio.gather(*list(self._writes), return_exceptions=True)
|
||||
|
||||
# ---------------------------------------------------------------- 主流程
|
||||
|
||||
async def run(self) -> None:
|
||||
try:
|
||||
await self._run()
|
||||
except asyncio.CancelledError:
|
||||
current = asyncio.current_task()
|
||||
if current is not None:
|
||||
current.uncancel()
|
||||
if self.stop_reason == "user":
|
||||
await self._end("cancelled", None, "cancelled", None)
|
||||
else:
|
||||
await self._end("interrupted", INTERRUPTED_ERROR, "cancelled", INTERRUPTED_ERROR)
|
||||
except _JobFailed as e:
|
||||
await self._end("failed", str(e), "failed", str(e))
|
||||
except Exception as e: # noqa: BLE001 - 任务整体出错：照实记下，不留在「起草中」
|
||||
from app.core.errors import describe_exception, first_line
|
||||
|
||||
logger.exception("数据目录起草任务 %s 出错", self.job_id)
|
||||
reason = JOB_ERROR.format(line=first_line(e) or describe_exception(e))
|
||||
await self._end("failed", reason, "failed", reason)
|
||||
|
||||
async def _end(self, job_status: str, job_error: str | None, task_status: str, task_error: str | None) -> None:
|
||||
"""停止、出错时收尾：等手上的库写入写完，再把没结束的表和任务记下来。收尾本身写不进去只记日志（下次启动时
|
||||
startup 会把它标成中断）。"""
|
||||
try:
|
||||
await self._settle()
|
||||
await self._write(_close, self.job_id, job_status, job_error, task_status, task_error)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("数据目录起草任务 %s 收尾时写库出错", self.job_id)
|
||||
|
||||
async def _run(self) -> None:
|
||||
if not await self._write(self._claim):
|
||||
return # 开始之前就被停止了
|
||||
async with SessionLocal() as session:
|
||||
job = await session.get(CatalogDraftJob, self.job_id)
|
||||
row = await session.get(DataSource, job.source_id) if job is not None else None
|
||||
if job is None or row is None:
|
||||
raise _JobFailed(SOURCE_GONE)
|
||||
from app.data import table_versions
|
||||
|
||||
try:
|
||||
source = await table_versions.resolve_source(session, row)
|
||||
except table_versions.SnapshotMissing as e:
|
||||
raise _JobFailed(str(e)) from e
|
||||
pending = [t.table_name for t in await tasks_of(session, job.id) if t.status == "pending"]
|
||||
use_model, background = bool(job.use_model), catalog_draft.Background(guidance=job.guidance, note=job.note)
|
||||
self.actor = job.created_by
|
||||
model: Any = None
|
||||
existing: dict[str, Any] = {}
|
||||
if use_model and pending:
|
||||
try:
|
||||
model, model_id = await catalog.resolve_draft_model(session)
|
||||
except catalog.CatalogModelUnavailable as e:
|
||||
raise _JobFailed(str(e)) from e
|
||||
existing = {t: e.notes for t, e in (await catalog.read_catalog(session, source.id)).items()}
|
||||
self.source, self.at = source, catalog.now_iso()
|
||||
all_tables = (getattr(source, "schema_cache", None) or {}).get("tables") or {}
|
||||
names = [n for n in pending if isinstance(all_tables.get(n), dict)]
|
||||
gone = [n for n in pending if n not in names]
|
||||
if gone:
|
||||
# 建任务之后表结构变了（重新探查过）：这几张表跳过
|
||||
await self._write(self._skip, gone)
|
||||
if use_model and names:
|
||||
await self._write(self._set_model, model_id)
|
||||
hooks = catalog_draft.Hooks(call_started=self._call_started, usage=self._usage,
|
||||
part_done=self._part_done, table_done=self._table_done)
|
||||
await catalog_draft.run_model_draft(model, source, names, existing=existing, background=background,
|
||||
hooks=hooks)
|
||||
else:
|
||||
for name in names:
|
||||
await self._call_started([name])
|
||||
await self._table_done(name, None)
|
||||
await self._settle()
|
||||
await self._write(self._finish)
|
||||
|
||||
# ---------------------------------------------------------------- 进度回调（引擎调）
|
||||
|
||||
async def _call_started(self, tables: list[str]) -> None:
|
||||
fresh = [t for t in dict.fromkeys(tables) if t not in self._started]
|
||||
if fresh:
|
||||
self._started.update(fresh)
|
||||
await self._write(self._mark_running, fresh)
|
||||
|
||||
async def _usage(self, input_tokens: int, output_tokens: int) -> None:
|
||||
await self._write(self._add_usage, input_tokens, output_tokens)
|
||||
|
||||
async def _part_done(self, table: str) -> None:
|
||||
await self._write(self._bump_part, table)
|
||||
|
||||
async def _table_done(self, table: str, outcome: catalog_draft.ModelOutcome | None) -> None:
|
||||
await self._write(self._write_table, table, outcome)
|
||||
|
||||
# ---------------------------------------------------------------- 库操作
|
||||
|
||||
def _task(self, *where: Any) -> Any:
|
||||
return update(CatalogDraftTask).where(CatalogDraftTask.job_id == self.job_id, *where) \
|
||||
.execution_options(synchronize_session=False)
|
||||
|
||||
async def _claim(self, session: AsyncSession) -> bool:
|
||||
result = await session.execute(
|
||||
update(CatalogDraftJob).where(CatalogDraftJob.id == self.job_id, CatalogDraftJob.status == "queued")
|
||||
.values(status="running", started_at=utcnow()).execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
return result.rowcount == 1
|
||||
|
||||
async def _set_model(self, session: AsyncSession, model_id: str) -> None:
|
||||
await session.execute(update(CatalogDraftJob).where(CatalogDraftJob.id == self.job_id)
|
||||
.values(model=model_id or None).execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
|
||||
async def _skip(self, session: AsyncSession, names: list[str]) -> None:
|
||||
now = utcnow()
|
||||
for name in names:
|
||||
await session.execute(self._task(CatalogDraftTask.table_name == name, CatalogDraftTask.status == "pending")
|
||||
.values(status="skipped", error=catalog.missing_table_error(name), finished_at=now))
|
||||
await session.commit()
|
||||
|
||||
async def _mark_running(self, session: AsyncSession, names: list[str]) -> None:
|
||||
await session.execute(self._task(CatalogDraftTask.table_name.in_(names), CatalogDraftTask.status == "pending")
|
||||
.values(status="running", started_at=utcnow()))
|
||||
await session.commit()
|
||||
|
||||
async def _add_usage(self, session: AsyncSession, input_tokens: int, output_tokens: int) -> None:
|
||||
await session.execute(
|
||||
update(CatalogDraftJob).where(CatalogDraftJob.id == self.job_id)
|
||||
.values(calls=CatalogDraftJob.calls + 1, input_tokens=CatalogDraftJob.input_tokens + int(input_tokens),
|
||||
output_tokens=CatalogDraftJob.output_tokens + int(output_tokens))
|
||||
.execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
|
||||
async def _bump_part(self, session: AsyncSession, table: str) -> None:
|
||||
await session.execute(self._task(CatalogDraftTask.table_name == table)
|
||||
.values(parts_done=CatalogDraftTask.parts_done + 1))
|
||||
await session.commit()
|
||||
|
||||
async def _write_table(self, session: AsyncSession, table: str,
|
||||
outcome: catalog_draft.ModelOutcome | None) -> None:
|
||||
"""一张表结束：并入目录、写入，记下这张表的结果。"""
|
||||
try:
|
||||
result = await catalog.write_draft(session, self.source, table, outcome, at=self.at, actor=self.actor)
|
||||
except Exception as e: # noqa: BLE001 - 这一张写不进去不影响别的表
|
||||
from app.core.errors import describe_exception, first_line
|
||||
|
||||
logger.exception("起草任务 %s 写入表 %s 的数据目录时出错", self.job_id, table)
|
||||
await session.rollback()
|
||||
result = catalog.TableDraftResult(table=table, status="failed",
|
||||
error=WRITE_ERROR.format(line=first_line(e) or describe_exception(e)),
|
||||
model_error=outcome.error if outcome is not None else None)
|
||||
values: dict[str, Any] = {
|
||||
"status": result.status, "added": result.added, "updated": result.updated, "removed": result.removed,
|
||||
"version": result.version, "model_items": result.model_items, "error": result.error,
|
||||
"model_error": result.model_error, "finished_at": utcnow(),
|
||||
}
|
||||
if outcome is not None:
|
||||
values["parts_done"] = outcome.parts_done
|
||||
await session.execute(self._task(CatalogDraftTask.table_name == table, CatalogDraftTask.status.in_(_OPEN))
|
||||
.values(**values))
|
||||
await session.commit()
|
||||
|
||||
async def _finish(self, session: AsyncSession) -> None:
|
||||
"""起草完：万一还有没收尾的表（进度回调写库失败之类）记失败；有 failed、partial、cancelled 的表是 partial，
|
||||
否则 succeeded。"""
|
||||
now = utcnow()
|
||||
await session.execute(self._task(CatalogDraftTask.status.in_(_OPEN))
|
||||
.values(status="failed", error=UNFINISHED_ERROR, finished_at=now))
|
||||
statuses = list((await session.execute(
|
||||
select(CatalogDraftTask.status).where(CatalogDraftTask.job_id == self.job_id))).scalars())
|
||||
status = "partial" if any(s in RETRYABLE for s in statuses) else "succeeded"
|
||||
await session.execute(
|
||||
update(CatalogDraftJob).where(CatalogDraftJob.id == self.job_id, CatalogDraftJob.status == "running")
|
||||
.values(status=status, finished_at=now).execution_options(synchronize_session=False))
|
||||
await session.commit()
|
||||
|
||||
|
||||
class DraftJobs:
|
||||
"""进程内的起草任务（照 app/scheduler 的写法：一个 asyncio 任务一件事，停的时候给宽限）。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._runs: dict[str, _JobRun] = {}
|
||||
|
||||
def _alive(self, job_id: str) -> _JobRun | None:
|
||||
run = self._runs.get(job_id)
|
||||
if run is None or run.task is None or run.task.done():
|
||||
return None
|
||||
try:
|
||||
if run.task.get_loop() is not asyncio.get_running_loop():
|
||||
self._runs.pop(job_id, None) # 别的事件循环留下的（测试每个用例一个循环）
|
||||
return None
|
||||
except RuntimeError:
|
||||
return None
|
||||
return run
|
||||
|
||||
def start(self, job_id: str) -> None:
|
||||
"""开始在后台起草这个（排队中的）任务。"""
|
||||
run = _JobRun(job_id)
|
||||
run.task = asyncio.get_running_loop().create_task(run.run(), name=f"catalog-draft-{job_id}")
|
||||
self._runs[job_id] = run
|
||||
|
||||
def forget(_: asyncio.Task[None]) -> None:
|
||||
if self._runs.get(job_id) is run:
|
||||
self._runs.pop(job_id, None)
|
||||
|
||||
run.task.add_done_callback(forget)
|
||||
|
||||
def running(self, job_id: str) -> bool:
|
||||
return self._alive(job_id) is not None
|
||||
|
||||
async def stop(self, job_id: str, reason: str, *, grace: float = STOP_GRACE_S) -> bool:
|
||||
"""停掉这个进程里在跑的任务，等它收尾（最多 grace 秒）。不在这个进程里跑返回 False。"""
|
||||
run = self._alive(job_id)
|
||||
if run is None:
|
||||
return False
|
||||
if run.stop_reason is None:
|
||||
run.stop_reason = reason
|
||||
run.task.cancel() # type: ignore[union-attr]
|
||||
await asyncio.wait([run.task], timeout=grace) # type: ignore[list-item]
|
||||
return True
|
||||
|
||||
async def wait(self, job_id: str, timeout: float | None = None) -> None:
|
||||
"""等这个任务做完（测试、调试用）。不在跑直接返回。"""
|
||||
run = self._alive(job_id)
|
||||
if run is not None:
|
||||
await asyncio.wait([run.task], timeout=timeout) # type: ignore[list-item]
|
||||
|
||||
async def shutdown(self, *, grace: float = 3.0) -> None:
|
||||
"""关停：停掉进行中的任务（标中断），最多等 grace 秒收尾。"""
|
||||
runs = [run for job_id in list(self._runs) if (run := self._alive(job_id)) is not None]
|
||||
for run in runs:
|
||||
if run.stop_reason is None:
|
||||
run.stop_reason = "shutdown"
|
||||
run.task.cancel() # type: ignore[union-attr]
|
||||
if runs:
|
||||
await asyncio.wait([run.task for run in runs], timeout=grace) # type: ignore[misc]
|
||||
logger.info("已停止 %d 个数据目录起草任务", len(runs))
|
||||
|
||||
|
||||
manager = DraftJobs()
|
||||
|
||||
|
||||
__all__ = ["ACTIVE", "DraftBusy", "DraftFinished", "DraftJobs", "FINISHED", "GUIDANCE_MAX", "INTERRUPTED_ERROR",
|
||||
"MAX_TABLES", "NOTE_MAX", "NothingToRetry", "RETRYABLE", "TASK_DONE", "TASK_STATUSES", "TextTooLong",
|
||||
"active_job", "cancel", "check_model", "clean_text", "create_job", "ensure_idle", "guidance_of", "job_brief",
|
||||
"job_dict", "job_out", "last_drafts", "latest_job", "load", "manager", "retry_tables", "save_guidance",
|
||||
"startup", "tasks_of"]
|
||||
@@ -0,0 +1,114 @@
|
||||
"""起草用哪个模型：数据目录的「用模型起草」和按配方导入的 AI 起草共用这一处取值。
|
||||
|
||||
顺序（哪一级写了就整组用哪一级，不跨级拼接入和模型）：
|
||||
|
||||
1. 设置里的起草模型（catalog_draft 一组：{provider, model}），填了接入或模型就用它；
|
||||
2. 否则沿用助手的模型设置（copilot 一组，「编排」页助手栏的「助手使用的模型」）；
|
||||
3. 再没有就是默认接入（providers/factory.resolve_provider 的回退）。
|
||||
|
||||
这里只决定「用哪个接入、哪个模型」。额度、思考、温度由两边各自定：数据目录起草保留思考、额度 32768
|
||||
（data/catalog.resolve_draft_model），配方 AI 起草照旧关思考（data/recipe_ai.resolve_draft_model）。
|
||||
|
||||
**为什么不用 get_chat_model。** 它在没有接入时静默退回演示模型，拿演示模型起草出来的东西会被当成真的。
|
||||
两边都自己解析接入：没有、停用、只有演示接入、缺 key，都在调用模型之前判定为不可用，说清是哪一级的设置出了问题。
|
||||
|
||||
数据层不导入接口层：设置键在这里定义，api/settings.py 从这里取（助手那一项和 api/copilot.py 同值，有测试核对）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
#: 设置里起草模型那一组的键
|
||||
CATALOG_DRAFT_SETTING_KEY = "catalog_draft"
|
||||
#: 助手的模型设置键（与 api/copilot.py 的 COPILOT_SETTING_KEY 同值）
|
||||
COPILOT_SETTING_KEY = "copilot"
|
||||
#: 起草模型在设置页上的位置。报错里让人去这里改；和前端设置页放这一组的地方一致
|
||||
DRAFT_MODEL_PLACE = "「设置 → 偏好设置」的「起草模型」"
|
||||
|
||||
#: 取自哪一级：draft 设置里的起草模型、assistant 助手的模型、default 默认接入
|
||||
ORIGINS = ("draft", "assistant", "default")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DraftModelChoice:
|
||||
"""起草用的接入和模型（都可以为 None：接入为 None 时按模型名找、再退回默认接入；模型为 None 时用接入的默认模型）。"""
|
||||
|
||||
provider: str | None
|
||||
model: str | None
|
||||
origin: str
|
||||
|
||||
|
||||
def _text(value: Any) -> str | None:
|
||||
"""设置里的一项：字符串去掉首尾空白，空串和别的类型（存坏了的）都当没填。"""
|
||||
return value.strip() or None if isinstance(value, str) else None
|
||||
|
||||
|
||||
def setting_pair(stored: Any) -> tuple[str | None, str | None]:
|
||||
"""库里存的一组 → (接入, 模型)，规范化同 _text。"""
|
||||
stored = stored if isinstance(stored, dict) else {}
|
||||
return _text(stored.get("provider")), _text(stored.get("model"))
|
||||
|
||||
|
||||
async def draft_model_choice(session: AsyncSession) -> DraftModelChoice:
|
||||
"""按上面的顺序取。只读设置，不查接入在不在（resolve_provider 去查）。"""
|
||||
from app.db.models import Setting
|
||||
|
||||
for key, origin in ((CATALOG_DRAFT_SETTING_KEY, "draft"), (COPILOT_SETTING_KEY, "assistant")):
|
||||
row = await session.get(Setting, key)
|
||||
provider, model = setting_pair(row.value if row is not None else None)
|
||||
if provider or model:
|
||||
return DraftModelChoice(provider, model, origin)
|
||||
return DraftModelChoice(None, None, "default")
|
||||
|
||||
|
||||
def draft_unavailable(e: BaseException, origin: str) -> str:
|
||||
"""取到的接入用不了（点名的不存在、已停用，缺 key、没有模型……）时给人看的那句话。
|
||||
|
||||
沿用助手的模型（或默认接入）时，和助手取不到模型时同一句（factory.assistant_unavailable）；用的是设置里的起草
|
||||
模型时，照实说是起草模型那一项、去设置页哪里改。"""
|
||||
from app.core.errors import not_configured
|
||||
from app.providers.factory import NAMED_FIX, assistant_unavailable, model_fact
|
||||
|
||||
if origin != "draft":
|
||||
return assistant_unavailable(e)
|
||||
code = getattr(e, "code", "")
|
||||
if code in NAMED_FIX:
|
||||
said = f"起草模型使用的{getattr(e, 'reason', '') or e}。"
|
||||
if owner := getattr(e, "model_owner", ""):
|
||||
return (f"{said}{NAMED_FIX[code]}；{model_fact(getattr(e, 'name', ''), owner)}，"
|
||||
f"也可在{DRAFT_MODEL_PLACE}中把它填进模型一栏")
|
||||
return f"{said}{NAMED_FIX[code]}，或在{DRAFT_MODEL_PLACE}中重新选择"
|
||||
return (f"起草模型尚未配置完成：{not_configured(e, with_hint=False)}。请到「设置 → 模型接入」检查，"
|
||||
f"或在{DRAFT_MODEL_PLACE}中重新选择")
|
||||
|
||||
|
||||
class DraftModelUnavailable(RuntimeError):
|
||||
"""没有可用于起草的模型。message 给人看；在调用模型之前判定。"""
|
||||
|
||||
|
||||
async def resolve_draft_provider(session: AsyncSession, *, none_reason: str, mock_reason: str) -> tuple[Any, DraftModelChoice]:
|
||||
"""(接入, 取值)。用不了时抛 DraftModelUnavailable：没有任何启用的接入说 none_reason，只有演示接入说 mock_reason
|
||||
（两边各有各的说法），点名的接入不存在、已停用按取自哪一级说（draft_unavailable）。
|
||||
|
||||
只解析接入，不建模型对象：建的时候缺 key 之类的问题由调用方经 draft_unavailable 说。"""
|
||||
from app.providers.factory import ProviderNotConfigured, resolve_provider
|
||||
|
||||
choice = await draft_model_choice(session)
|
||||
try:
|
||||
# 存的接入名不存在、已停用，只存了模型名、提供它的接入已停用：resolve_provider 都会报出来
|
||||
provider = await resolve_provider(session, choice.provider, choice.model)
|
||||
except ProviderNotConfigured as e:
|
||||
raise DraftModelUnavailable(draft_unavailable(e, choice.origin)) from e
|
||||
if provider is None:
|
||||
raise DraftModelUnavailable(none_reason)
|
||||
if provider.kind == "mock":
|
||||
raise DraftModelUnavailable(mock_reason)
|
||||
return provider, choice
|
||||
|
||||
|
||||
__all__ = ["CATALOG_DRAFT_SETTING_KEY", "COPILOT_SETTING_KEY", "DRAFT_MODEL_PLACE", "DraftModelChoice",
|
||||
"DraftModelUnavailable", "draft_model_choice", "draft_unavailable", "resolve_draft_provider",
|
||||
"setting_pair"]
|
||||
@@ -62,12 +62,18 @@ MASK_COLUMNS_OPTION = "mask_columns"
|
||||
#: data/catalog_profile.py；放在这里只为让拼连接串时摘掉它
|
||||
CATALOG_PROFILE_OPTION = "catalog_profile"
|
||||
|
||||
#: 数据源 options 里的数据目录起草要求：使用者补充的背景和要求，起草时交给模型（一段文字，最多 2000 字）。
|
||||
#: 只经 PUT /api/datasources/{id}/catalog/guidance 维护，见 data/catalog_jobs.py；放在这里只为让拼连接串时摘掉它
|
||||
CATALOG_DRAFT_GUIDANCE_OPTION = "catalog_draft_guidance"
|
||||
|
||||
# options 里这些 key 是 AgentLab 自己的配置，不是驱动参数，拼 URL 时要摘掉。
|
||||
# schema：探查哪个 schema（企业库里只读账号名下常常什么都没有，数据在别处）
|
||||
# query_timeout_s：查询时限，由数据层按语句下发给数据库，见 _server_deadline
|
||||
# mask_columns：证据面板遮罩的列，驱动不认识它，拼进连接串会被当成未知参数拒掉
|
||||
# catalog_profile：数据剖析设置，同上；它还是个对象，拼进去连 URL 都不成形
|
||||
_NON_DRIVER_OPTIONS = frozenset({"schema", QUERY_TIMEOUT_OPTION, MASK_COLUMNS_OPTION, CATALOG_PROFILE_OPTION})
|
||||
# catalog_draft_guidance：数据目录的起草要求，一段中文，拼进去连接直接失败
|
||||
_NON_DRIVER_OPTIONS = frozenset({"schema", QUERY_TIMEOUT_OPTION, MASK_COLUMNS_OPTION, CATALOG_PROFILE_OPTION,
|
||||
CATALOG_DRAFT_GUIDANCE_OPTION})
|
||||
|
||||
_MASK_SPLIT = re.compile(r"[,，、;；\n]+")
|
||||
|
||||
|
||||
@@ -6,7 +6,8 @@
|
||||
必须人工确认之后才能启用。每月重传按已确认的配方重放，不再调用模型。
|
||||
|
||||
**为什么不用 get_chat_model。** 它在没有接入时会静默退回演示模型；拿演示模型起草等于假装起草了。这里
|
||||
自己按助手的模型设置解析接入，没有、停用、只有演示接入、缺 key 都在调用模型之前判定为不可用。
|
||||
自己解析接入（取值顺序和数据目录起草共用 data/draft_model.py：设置里的起草模型 > 助手的模型 > 默认接入），
|
||||
没有、停用、只有演示接入、缺 key 都在调用模型之前判定为不可用。
|
||||
|
||||
**调用参数照 c837cdb 修过的问题来**：关思考（thinking="off"）、额度 32768、服务方嫌额度大时退回 8192；
|
||||
正文一律经 message_text 抽取（回复可能是块列表）；输出被截断单独记一种失败，不当成「不是 JSON」去修订。
|
||||
@@ -42,6 +43,7 @@ from app.data.recipe_types import (
|
||||
GridCell,
|
||||
Recipe,
|
||||
)
|
||||
from app.data.draft_model import COPILOT_SETTING_KEY as _DRAFT_COPILOT_KEY
|
||||
from app.data.tabular import EXCEL_ERRORS
|
||||
from app.data.xlsx_scan import WorkbookScan, cell_ref, col_letter
|
||||
|
||||
@@ -58,8 +60,9 @@ class AiTooLarge(ValueError):
|
||||
|
||||
# 没有 AiFailed：调用之后的任何失败都体现在返回值里（draft=None、error、usage 照记）
|
||||
|
||||
#: 助手的模型设置键（与 app/api/copilot.py 的 COPILOT_SETTING_KEY 同值；数据层不导入接口层，测试断言两者相等）
|
||||
COPILOT_SETTING_KEY = "copilot"
|
||||
#: 助手的模型设置键（与 app/api/copilot.py 的 COPILOT_SETTING_KEY 同值；数据层不导入接口层，测试断言两者相等）。
|
||||
#: 取模型时由 data/draft_model.py 读，这里留着这个名字给测试和老代码用
|
||||
COPILOT_SETTING_KEY = _DRAFT_COPILOT_KEY
|
||||
#: 1 次起草 + 最多 2 次修订
|
||||
AI_MAX_ATTEMPTS = 3
|
||||
#: 单次调用的超时（秒）；测试 monkeypatch 调小
|
||||
@@ -252,37 +255,27 @@ class _BudgetFallback:
|
||||
return await self.model.ainvoke(messages)
|
||||
|
||||
|
||||
def _unconfigured(e: BaseException) -> str:
|
||||
"""和 copilot._unconfigured 同一句（providers/factory.assistant_unavailable）。"""
|
||||
from app.providers.factory import assistant_unavailable
|
||||
|
||||
return assistant_unavailable(e)
|
||||
|
||||
|
||||
async def resolve_draft_model(session: AsyncSession) -> tuple[Any | None, str, AiAvailability]:
|
||||
"""(可调用的模型, 模型 id, 可用性)；不可用时模型为 None。只建模型对象，不发请求。"""
|
||||
from app.db.models import Setting
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, build_chat_model, resolve_provider
|
||||
"""(可调用的模型, 模型 id, 可用性)；不可用时模型为 None。只建模型对象，不发请求。
|
||||
|
||||
用哪个接入、哪个模型和数据目录起草共用一处取值（data/draft_model.py：起草模型 > 助手的模型 > 默认接入）；
|
||||
额度、关思考照旧是这里自己的。"""
|
||||
from app.data.draft_model import DraftModelUnavailable, draft_unavailable, resolve_draft_provider
|
||||
from app.providers.factory import ModelSpec, ProviderNotConfigured, build_chat_model
|
||||
|
||||
row = await session.get(Setting, COPILOT_SETTING_KEY)
|
||||
saved = (row.value if row else None) or {}
|
||||
spec = ModelSpec(provider=saved.get("provider") or None, model=saved.get("model") or None,
|
||||
max_tokens=AI_MAX_TOKENS, temperature=0, thinking="off")
|
||||
try:
|
||||
provider = await resolve_provider(session, spec.provider, spec.model)
|
||||
except ProviderNotConfigured as e:
|
||||
# 存的接入名不存在、已停用，只存了模型名、提供它的接入已停用：resolve_provider 都会报出来
|
||||
return None, "", AiAvailability(False, reason=_unconfigured(e))
|
||||
if provider is None:
|
||||
return None, "", AiAvailability(False, reason="未配置模型接入，无法使用 AI 起草。请到「设置 → 模型接入」添加")
|
||||
if provider.kind == "mock":
|
||||
return None, "", AiAvailability(False, reason="当前只有演示用的模型接入，无法用于 AI 起草。"
|
||||
"请到「设置 → 模型接入」添加真实的模型接入")
|
||||
provider, choice = await resolve_draft_provider(
|
||||
session, none_reason="未配置模型接入，无法使用 AI 起草。请到「设置 → 模型接入」添加",
|
||||
mock_reason="当前只有演示用的模型接入，无法用于 AI 起草。请到「设置 → 模型接入」添加真实的模型接入")
|
||||
except DraftModelUnavailable as e:
|
||||
return None, "", AiAvailability(False, reason=str(e))
|
||||
spec = ModelSpec(provider=choice.provider, model=choice.model, max_tokens=AI_MAX_TOKENS, temperature=0,
|
||||
thinking="off")
|
||||
try:
|
||||
big = build_chat_model(provider, spec)
|
||||
small = build_chat_model(provider, spec.model_copy(update={"max_tokens": AI_FALLBACK_MAX_TOKENS}))
|
||||
except ProviderNotConfigured as e:
|
||||
return None, "", AiAvailability(False, reason=_unconfigured(e))
|
||||
return None, "", AiAvailability(False, reason=draft_unavailable(e, choice.origin))
|
||||
model_id = spec.model or provider.default_model or ""
|
||||
return (_BudgetFallback(big, small, model_id), model_id,
|
||||
AiAvailability(True, model=model_id, provider=provider.name))
|
||||
|
||||
@@ -4,7 +4,7 @@ from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import (
|
||||
Float, ForeignKey, Index, Integer, LargeBinary, String, Text, UniqueConstraint,
|
||||
Float, ForeignKey, Index, Integer, LargeBinary, String, Text, UniqueConstraint, text,
|
||||
)
|
||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||
|
||||
@@ -348,6 +348,76 @@ class CatalogNote(Base, TimestampMixin):
|
||||
updated_by: Mapped[str | None] = mapped_column(String(100), default=None)
|
||||
|
||||
|
||||
class CatalogDraftJob(Base, TimestampMixin):
|
||||
"""数据目录的一次后台起草任务（data/catalog_jobs.py）。逐表的状态在 catalog_draft_tasks。
|
||||
|
||||
以前起草是一次同步长请求，出错整个对话框就停、关掉之后没有任何记录哪些表失败、为什么；现在建任务后在后台起草，
|
||||
进度（每张表到第几段、计数、用量）随时落库，接口读的是库里的状态。
|
||||
|
||||
同一个数据源同时只有一个排队或进行中的任务：接口先查一遍，部分唯一索引兜底（两个请求同时建，后提交的撞上）。
|
||||
服务重启时还没结束的标成 interrupted，不自动续跑。挂外键、级联删除：任务是数据源的附属记录，源删掉了留着没有意义。
|
||||
"""
|
||||
|
||||
__tablename__ = "catalog_draft_jobs"
|
||||
__table_args__ = (
|
||||
Index("ix_catalog_draft_jobs_source_created", "source_id", "created_at"),
|
||||
Index("ux_catalog_draft_jobs_active", "source_id", unique=True,
|
||||
sqlite_where=text("status IN ('queued', 'running')")),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id)
|
||||
source_id: Mapped[str] = mapped_column(ForeignKey("data_sources.id", ondelete="CASCADE"))
|
||||
# queued | running | succeeded | partial | failed | cancelled | interrupted
|
||||
status: Mapped[str] = mapped_column(String(20), default="queued")
|
||||
use_model: Mapped[bool] = mapped_column(default=False)
|
||||
# 实际用的模型 id（用模型时开始起草后写上）
|
||||
model: Mapped[str | None] = mapped_column(String(200), default=None)
|
||||
# 建任务时的快照：数据源级的起草要求、这一次的补充说明
|
||||
guidance: Mapped[str | None] = mapped_column(Text, default=None)
|
||||
note: Mapped[str | None] = mapped_column(Text, default=None)
|
||||
# 请求头 X-Actor 的署名（自报、未认证）
|
||||
created_by: Mapped[str | None] = mapped_column(String(100), default=None)
|
||||
started_at: Mapped[datetime | None] = mapped_column(UTCDateTime, default=None)
|
||||
finished_at: Mapped[datetime | None] = mapped_column(UTCDateTime, default=None)
|
||||
# 用量：调用次数、输入和输出 token（拿不到用量的调用按 0 算）
|
||||
calls: Mapped[int] = mapped_column(Integer, default=0)
|
||||
input_tokens: Mapped[int] = mapped_column(Integer, default=0)
|
||||
output_tokens: Mapped[int] = mapped_column(Integer, default=0)
|
||||
# 任务整体出错的原因（逐表的原因在 catalog_draft_tasks）
|
||||
error: Mapped[str | None] = mapped_column(Text, default=None)
|
||||
# 重试建的任务记下是重试哪一个
|
||||
retry_of: Mapped[str | None] = mapped_column(String(32), default=None)
|
||||
|
||||
|
||||
class CatalogDraftTask(Base):
|
||||
"""起草任务里的一张表。按 seq（起草顺序）排。状态：pending 排队、running 起草中、ok、partial 模型只写出一部分、
|
||||
failed、skipped 表结构里没有、cancelled（停止或服务重启）。"""
|
||||
|
||||
__tablename__ = "catalog_draft_tasks"
|
||||
__table_args__ = (UniqueConstraint("job_id", "seq"),)
|
||||
|
||||
id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id)
|
||||
job_id: Mapped[str] = mapped_column(ForeignKey("catalog_draft_jobs.id", ondelete="CASCADE"), index=True)
|
||||
seq: Mapped[int] = mapped_column(Integer)
|
||||
# 表结构里的键；表结构里没有的（skipped）是请求里的原样名字
|
||||
table_name: Mapped[str] = mapped_column(String(255))
|
||||
status: Mapped[str] = mapped_column(String(20), default="pending")
|
||||
column_count: Mapped[int] = mapped_column(Integer, default=0)
|
||||
# 模型分几段起草（不用模型时为 NULL）、已经结束了几段
|
||||
parts_total: Mapped[int | None] = mapped_column(Integer, default=None)
|
||||
parts_done: Mapped[int] = mapped_column(Integer, default=0)
|
||||
added: Mapped[int] = mapped_column(Integer, default=0)
|
||||
updated: Mapped[int] = mapped_column(Integer, default=0)
|
||||
removed: Mapped[int] = mapped_column(Integer, default=0)
|
||||
version: Mapped[int] = mapped_column(Integer, default=0)
|
||||
model_items: Mapped[int | None] = mapped_column(Integer, default=None)
|
||||
# error：整张表没起草（表结构里没有、写入一直冲突、服务重启中断）；model_error：模型那部分没写出来的原因和办法
|
||||
error: Mapped[str | None] = mapped_column(Text, default=None)
|
||||
model_error: Mapped[str | None] = mapped_column(Text, default=None)
|
||||
started_at: Mapped[datetime | None] = mapped_column(UTCDateTime, default=None)
|
||||
finished_at: Mapped[datetime | None] = mapped_column(UTCDateTime, default=None)
|
||||
|
||||
|
||||
class Setting(Base, TimestampMixin):
|
||||
"""键值形式的用户设置，前端 Settings 页直接读写。"""
|
||||
|
||||
|
||||
+11
-1
@@ -70,6 +70,15 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("上传表格的版本整理失败，服务照常启动")
|
||||
|
||||
# 上次没做完的数据目录起草任务：标成中断（没结束的表记停止），不自动续跑——拷来的数据目录（检查脚本、验收副本）
|
||||
# 启动时不能替人去调模型。界面上点「重试」接着起草没写好的表
|
||||
from app.data import catalog_jobs
|
||||
|
||||
try:
|
||||
await catalog_jobs.startup()
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("数据目录起草任务的启动整理失败，服务照常启动")
|
||||
|
||||
await run_manager.setup()
|
||||
from app.sandbox.manager import sandbox_manager as _sbm
|
||||
|
||||
@@ -84,8 +93,9 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# 先停调度，再停运行：停运行的途中不该再有新运行被发起
|
||||
# 先停调度，再停运行：停运行的途中不该再有新运行被发起。起草任务停下来标成中断，下次启动不会再续
|
||||
await scheduler.stop()
|
||||
await catalog_jobs.manager.shutdown()
|
||||
await run_manager.shutdown()
|
||||
from app.sandbox.manager import sandbox_manager
|
||||
|
||||
|
||||
@@ -28,6 +28,9 @@ class ModelSpec(BaseModel):
|
||||
thinking: str | None = None # off | adaptive
|
||||
effort: str | None = None # low | medium | high | xhigh | max
|
||||
timeout: int | None = None
|
||||
# SDK 自己遇到 429、5xx、连接错误、超时时重试几次；None 照 SDK 默认（openai、anthropic 都是 2 次）。
|
||||
# 自己按失败原因重试的调用方（数据目录起草）设 0，免得两层重试叠在一起、实际次数和报错写的对不上
|
||||
max_retries: int | None = None
|
||||
extra: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
@@ -227,6 +230,8 @@ def build_chat_model(provider: Provider, spec: ModelSpec) -> BaseChatModel:
|
||||
# 以前只有节点显式配了才传，而没有任何地方给 spec.timeout 赋过值——所有调用
|
||||
# 都走 SDK 默认的 10 分钟
|
||||
kwargs["timeout"] = spec.timeout or settings.model_timeout_seconds
|
||||
if spec.max_retries is not None:
|
||||
kwargs["max_retries"] = spec.max_retries
|
||||
|
||||
# 采样参数：能力允许才传
|
||||
if spec.temperature is not None and catalog.supports_sampling(model):
|
||||
@@ -314,7 +319,8 @@ def build_chat_model(provider: Provider, spec: ModelSpec) -> BaseChatModel:
|
||||
# 且 api_key=None 时不会发出 x-api-key。
|
||||
import anthropic
|
||||
|
||||
client_args = {"base_url": base_url, "api_key": None, "auth_token": api_key}
|
||||
client_args = {"base_url": base_url, "api_key": None, "auth_token": api_key,
|
||||
"max_retries": chat.max_retries}
|
||||
chat._client = anthropic.Anthropic(**client_args)
|
||||
chat._async_client = anthropic.AsyncAnthropic(**client_args)
|
||||
return chat
|
||||
|
||||
+93
-8
@@ -3,10 +3,13 @@
|
||||
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 里这几列）；
|
||||
- 照提示词的要求，请它写的每一列都列出来：replies 里没写的列只给列名（{"name": 列名}，字段留空，是拿不准），
|
||||
不整列漏掉。list_all=False 时只给 replies 里有的列（模拟漏列的模型：漏掉的列会被补问一次，补问后仍没给的算没写出来）；
|
||||
- 直接 ainvoke 是纯文本那条路：同样的内容写成一段正文（AIMessage），前面带一句话；
|
||||
- fail 里的表所在的那一批抛异常，用来验证「一批失败只影响这一批」；
|
||||
- fail 里的表所在的那一段抛异常，用来验证「一段失败只影响这一段」；
|
||||
- calls 记下每次收到的提示词全文，用来核对交给模型的内容（只有结构，没有数据）；modes 记下每次走的是
|
||||
structured 还是 text。
|
||||
|
||||
@@ -22,12 +25,90 @@ 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
|
||||
|
||||
|
||||
def listed(text: str) -> dict[str, list[str]]:
|
||||
"""提示词里每张表的结构描述列出了哪些列（「- 列名 类型」那几行）。"""
|
||||
out: dict[str, list[str]] = {}
|
||||
for chunk in text.split("\n\n"):
|
||||
lines = chunk.split("\n")
|
||||
if not lines[0].startswith("表 "):
|
||||
continue
|
||||
name = lines[0][2:]
|
||||
if name.endswith("（视图）"):
|
||||
name = name[:-len("（视图）")]
|
||||
out[name] = [line[2:].split(" ")[0] for line in lines if line.startswith("- ")]
|
||||
return out
|
||||
|
||||
|
||||
def wanted_columns(text: str, name: str, part: str, columns: list[str] | None) -> list[str]:
|
||||
"""这一段请它写这张表的哪几列：整张表是结构描述里的全部列；只写表级的项时没有。"""
|
||||
if part == "full":
|
||||
return listed(text).get(name, [])
|
||||
if part in ("columns", "table+columns"):
|
||||
return list(columns or [])
|
||||
return []
|
||||
|
||||
|
||||
def with_every_column(entry: dict, wanted: list[str]) -> dict:
|
||||
"""请写的列都列出来：没写的只给列名（字段留空）。"""
|
||||
have = {c.get("name") for c in entry.get("columns") or [] if isinstance(c, dict)}
|
||||
missing = [{"name": c} for c in wanted if c not in have]
|
||||
if not missing:
|
||||
return entry
|
||||
return {**entry, "columns": list(entry.get("columns") or []) + missing}
|
||||
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self, replies: dict[str, dict], fail: set[str] = frozenset(), *,
|
||||
structured: str = "tables", text: str = "tables") -> None:
|
||||
structured: str = "tables", text: str = "tables", list_all: bool = True) -> None:
|
||||
self.replies, self.fail, self.calls = replies, set(fail), []
|
||||
self.structured, self.text = structured, text
|
||||
self.list_all = list_all
|
||||
self.modes: list[str] = []
|
||||
|
||||
def with_structured_output(self, schema: Any, **_: Any) -> "_Structured":
|
||||
@@ -44,17 +125,21 @@ 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 self.list_all:
|
||||
answers = [with_every_column(a, wanted_columns(text, t, *asked[t])) for a, t in zip(answers, 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,11 +274,15 @@ async def test_model_single_table_object_without_wrapper_is_recognized():
|
||||
|
||||
|
||||
async def test_hollow_structured_output_retries_with_text():
|
||||
"""结构化输出只给了 {"tables":[{"table":"visits"}]} 这样的空壳：改走纯文本再问一次，用纯文本给的内容。"""
|
||||
model = FakeModel({"visits": {"label": "入园记录", "kind": "fact"}}, structured="hollow", text="tables")
|
||||
"""结构化输出只给了 {"tables":[{"table":"visits"}]} 这样的空壳、这一段又拆不开（不到 10 列）：改走纯文本再问一次，
|
||||
用纯文本给的内容。纯文本只给了表级的项、一列都没列出来：再补问一次没给的列（只问缺的，data/catalog_draft.py）；
|
||||
补问后仍没给，这几列算没写出来，表是只写出一部分（旧的模型起草项不删）。"""
|
||||
model = FakeModel({"visits": {"label": "入园记录", "kind": "fact"}}, structured="hollow", text="tables",
|
||||
list_all=False)
|
||||
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")}
|
||||
assert got["visits"].covered == frozenset() and "有 5 列没写出来" in got["visits"].error
|
||||
|
||||
|
||||
async def test_hollow_structured_then_single_text_object():
|
||||
|
||||
@@ -0,0 +1,800 @@
|
||||
"""数据目录模型起草的引擎（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) == 3
|
||||
assert model.calls[1].asked == {"b_tab": ("full", None), "c_tab": ("full", None)}
|
||||
assert model.calls[2].asked == {"c_tab": ("full", None)} and model.calls[2].mode == "text"
|
||||
assert _status(got) == {"a_tab": "ok", "b_tab": "ok", "c_tab": "failed"}
|
||||
assert got["c_tab"].error == "模型的回复里没有这张表"
|
||||
|
||||
|
||||
async def test_a_table_still_missing_after_the_followup_is_split_like_a_hollow_shell():
|
||||
"""补问时整张表还是一项都没有：照第一次的规则拆小（这里按列拆成两半），不直接记失败。"""
|
||||
tables = {"a_tab": _table(4), "b_tab": _table(12)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
reply = ok(call)
|
||||
if call.n < 2:
|
||||
reply.obj["tables"] = [e for e in reply.obj["tables"] if e["table"] == "a_tab"]
|
||||
return reply
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["a_tab", "b_tab"])
|
||||
assert _status(got) == {"a_tab": "ok", "b_tab": "ok"}
|
||||
cols = [f"c{i:03d}" for i in range(1, 13)]
|
||||
assert [c.asked for c in model.calls[2:]] == [{"b_tab": ("table+columns", cols[:6])},
|
||||
{"b_tab": ("columns", cols[6:])}]
|
||||
|
||||
|
||||
async def test_columns_still_missing_after_the_followup_are_not_written():
|
||||
"""补问过仍然没列出来的列算没写出来：表是 partial，说清哪部分没写出来（不算拿不准，旧的模型起草项不删）。"""
|
||||
tables = {"t8": _table(8)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
reply = ok(call)
|
||||
for entry in reply.obj["tables"]:
|
||||
entry["columns"] = [c for c in entry.get("columns", []) if c["name"] not in ("c007", "c008")]
|
||||
return reply
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["t8"])
|
||||
assert len(model.calls) == 2 and model.calls[1].asked == {"t8": ("columns", ["c007", "c008"])}
|
||||
o = got["t8"]
|
||||
assert o.status == "partial"
|
||||
assert o.error == "这张表 8 列，有 2 列没写出来：模型的回复里漏了这些列，补问一次仍没有给出。可以只重试这张表"
|
||||
assert set(o.notes["columns"]) == {f"c{i:03d}" for i in range(1, 7)}
|
||||
compat = (await catalog.draft_with_model(Scripted(handle), _source(tables), ["t8"]))["t8"]
|
||||
assert compat.covered == frozenset(), "没写全的不算完整起草过"
|
||||
|
||||
|
||||
async def test_a_followup_that_comes_back_as_an_empty_shell_leaves_the_table_partial():
|
||||
"""40 列的表第一次只回了 20 列，补问只给一个空壳：没写出来的 20 列照实记下，表是 partial。"""
|
||||
tables = {"t40": _table(40)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
if call.n == 0:
|
||||
reply = ok(call)
|
||||
reply.obj["tables"][0]["columns"] = reply.obj["tables"][0]["columns"][:20]
|
||||
return reply
|
||||
return Reply(obj={"tables": [{"table": "t40"}]})
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["t40"])
|
||||
o = got["t40"]
|
||||
assert len(model.calls) == 2
|
||||
assert model.calls[1].asked == {"t40": ("columns", [f"c{i:03d}" for i in range(21, 41)])}
|
||||
assert o.status == "partial" and len(o.notes["columns"]) == 20
|
||||
assert o.error == "这张表 40 列，有 20 列没写出来：模型的回复里漏了这些列，补问一次仍没有给出。可以只重试这张表"
|
||||
|
||||
|
||||
async def test_columns_listed_with_empty_fields_are_left_to_the_model():
|
||||
"""列出了、字段留空的列是模型拿不准：不补问，也不算没写出来。"""
|
||||
tables = {"t8": _table(8)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
reply = ok(call)
|
||||
for entry in reply.obj["tables"]:
|
||||
entry["columns"] = [c if c["name"] not in ("c007", "c008") else {"name": c["name"]}
|
||||
for c in entry.get("columns", [])]
|
||||
return reply
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["t8"])
|
||||
assert len(model.calls) == 1
|
||||
assert got["t8"].status == "ok" and got["t8"].error is None
|
||||
assert set(got["t8"].notes["columns"]) == {f"c{i:03d}" for i in range(1, 7)}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 额度被拒、结构化输出不支持、网络错误、超时
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_budget_rejection_falls_back_to_8192_for_the_rest():
|
||||
tables = {"wide45": _table(45)}
|
||||
|
||||
def refuse(call: Call) -> Reply:
|
||||
raise _Status(400, "max_tokens: 32768 > 8192, which is the maximum allowed number of output tokens")
|
||||
|
||||
big, small = Scripted(refuse), Scripted(answering(tables))
|
||||
model = catalog.DraftModel(big, small, "m")
|
||||
got = await D.run_model_draft(model, _source(tables), ["wide45"])
|
||||
assert got["wide45"].status == "ok"
|
||||
assert len(big.calls) == 1 and len(small.calls) == 3, "退回之后，表级那段重来、两个列分段都用小额度"
|
||||
assert model.budget == 8192
|
||||
|
||||
|
||||
async def test_exhausted_message_names_the_budget_actually_used():
|
||||
def refuse(call: Call) -> Reply:
|
||||
raise _Status(400, "max_tokens is too large")
|
||||
|
||||
model = catalog.DraftModel(Scripted(refuse), Scripted(lambda c: Reply(text="", finish="length")), "m")
|
||||
got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"])
|
||||
assert got["t4"].error == "模型的思考和输出用完了额度（8192 个 token），这一段没写出来"
|
||||
|
||||
|
||||
async def test_structured_output_unsupported_switches_to_text_for_the_rest():
|
||||
tables = {"wide45": _table(45)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
if call.mode == "structured":
|
||||
raise _Status(400, "This response_format type is unavailable now")
|
||||
return ok(call)
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["wide45"])
|
||||
assert got["wide45"].status == "ok"
|
||||
assert [c.mode for c in model.calls] == ["structured", "text", "text", "text"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("thinking", ["", "先想想每一列是什么……"], ids=["plain", "with-thinking"])
|
||||
async def test_structured_output_that_comes_back_as_prose_falls_back_to_text(thinking):
|
||||
"""网关接受了结构化输出的参数却没照做（模型写成了一段说明文字）：换纯文本问，不当成额度用完去拆，也不直接判失败。"""
|
||||
tables = {"t40": _table(40)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
if call.mode == "structured":
|
||||
return Reply(text="好的。这张表记录的是一笔笔业务，各列的含义如下：第一列是编号……", thinking=thinking)
|
||||
return ok(call)
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["t40"])
|
||||
assert got["t40"].status == "ok" and len(got["t40"].notes["columns"]) == 40
|
||||
assert [c.mode for c in model.calls] == ["structured", "text"]
|
||||
assert catalog.CATALOG_DRAFT_TEXT_PROMPT in model.calls[1].text
|
||||
|
||||
|
||||
async def test_prose_on_both_paths_is_a_format_failure_not_a_budget_one():
|
||||
def handle(call: Call) -> Reply:
|
||||
return Reply(text="这张表看起来是业务明细，具体含义需要结合业务确认。", thinking="……")
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source({"t40": _table(40)}), ["t40"])
|
||||
assert [c.mode for c in model.calls] == ["structured", "text"], "格式不对不拆小重试"
|
||||
assert got["t40"].status == "failed" and got["t40"].error == "模型的回复不是约定的格式"
|
||||
|
||||
|
||||
async def test_a_model_without_structured_output_goes_to_text():
|
||||
class TextOnly(Scripted):
|
||||
def with_structured_output(self, schema, **_):
|
||||
raise NotImplementedError
|
||||
|
||||
tables = {"t4": _table(4)}
|
||||
model = TextOnly(answering(tables))
|
||||
got = await D.run_model_draft(model, _source(tables), ["t4"])
|
||||
assert got["t4"].status == "ok" and [c.mode for c in model.calls] == ["text"]
|
||||
|
||||
|
||||
async def test_network_errors_are_retried_with_backoff(_no_wait):
|
||||
tables = {"t4": _table(4)}
|
||||
ok = answering(tables)
|
||||
|
||||
def handle(call: Call) -> Reply:
|
||||
if call.n < 2:
|
||||
raise ConnectError("All connection attempts failed")
|
||||
return ok(call)
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source(tables), ["t4"])
|
||||
assert got["t4"].status == "ok" and len(model.calls) == 3
|
||||
assert _no_wait == list(catalog.MODEL_RETRY_DELAYS_S)
|
||||
|
||||
|
||||
async def test_server_errors_give_up_after_two_retries():
|
||||
def handle(call: Call) -> Reply:
|
||||
raise _Status(503, "Service Unavailable")
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"])
|
||||
assert len(model.calls) == 3
|
||||
assert got["t4"].status == "failed"
|
||||
assert got["t4"].error == "调用模型出错：Service Unavailable（已重试 2 次）"
|
||||
|
||||
|
||||
async def test_other_errors_are_not_retried():
|
||||
def handle(call: Call) -> Reply:
|
||||
raise RuntimeError("网关拒绝了这次请求")
|
||||
|
||||
model = Scripted(handle)
|
||||
got = await D.run_model_draft(model, _source({"t4": _table(4)}), ["t4"])
|
||||
assert len(model.calls) == 1 and got["t4"].error == "调用模型出错：网关拒绝了这次请求"
|
||||
|
||||
|
||||
async def test_a_call_past_the_time_limit_is_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"),
|
||||
])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 报错去重、止损、几次起草轮流用模型
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _hollow(call: Call) -> Reply:
|
||||
return Reply(obj={"tables": [{"table": t} for t in call.asked]})
|
||||
|
||||
|
||||
async def test_the_same_reason_is_said_once_even_when_some_parts_were_split(monkeypatch):
|
||||
"""宽表的表级那段拆不开、列分段拆过：原因只说一次，拆过的在最后补一句。"""
|
||||
monkeypatch.setattr(catalog, "MODEL_GIVE_UP_AFTER", 99)
|
||||
model = Scripted(_hollow)
|
||||
got = await D.run_model_draft(model, _source({"wide45": _table(45)}), ["wide45"])
|
||||
o = got["wide45"]
|
||||
assert o.status == "failed"
|
||||
assert o.error == "模型只给了空壳，没有写出可用的内容；已拆小重试仍没写出"
|
||||
|
||||
|
||||
def test_reasons_are_joined_by_their_base():
|
||||
W = D._Why
|
||||
exhausted = D.REASON_EXHAUSTED.format(budget=32768)
|
||||
assert D._join([W(D.REASON_HOLLOW), W(D.REASON_HOLLOW, split=True)]) == D.REASON_HOLLOW + D.REASON_SPLIT
|
||||
assert D._join([W(exhausted, split=True), W(D.REASON_HOLLOW)]) == \
|
||||
f"{exhausted}；{D.REASON_HOLLOW}{D.REASON_SPLIT}"
|
||||
assert D._join([W("甲"), W("乙"), W("丙", split=True)]) == "甲；乙", "没说出来的原因不补后缀"
|
||||
assert D._join([W(D.REASON_HOLLOW, halted=True), W(D.REASON_HOLLOW)]) == \
|
||||
D.REASON_HOLLOW + D.REASON_HALTED.format(n=catalog.MODEL_GIVE_UP_AFTER)
|
||||
|
||||
|
||||
def _by_table(model: Scripted) -> dict[str, int]:
|
||||
out: dict[str, int] = {}
|
||||
for call in model.calls:
|
||||
for name in call.asked:
|
||||
out[name] = out.get(name, 0) + 1
|
||||
return out
|
||||
|
||||
|
||||
async def test_after_two_segments_in_a_row_come_to_nothing_the_rest_are_not_split(monkeypatch):
|
||||
"""模型明显用不了（思考总是吃满额度）：连续 2 段拆到底仍一项没写出，后面的段只问一次、不再拆，报错说明已停止拆小重试。
|
||||
有一段写出来了就恢复拆小。一次只起草一段（并发 1），先后是确定的。"""
|
||||
monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1)
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
assert catalog.MODEL_GIVE_UP_AFTER == 2
|
||||
tables = {f"t{i}": _table(10) for i in range(1, 5)}
|
||||
exhausted = Reply(text="", finish="length", thinking="……")
|
||||
|
||||
model = Scripted(lambda call: exhausted)
|
||||
got = await D.run_model_draft(model, _source(tables), list(tables))
|
||||
assert _by_table(model) == {"t1": 3, "t2": 3, "t3": 1, "t4": 1}
|
||||
budget = D.REASON_EXHAUSTED.format(budget=32768)
|
||||
assert got["t1"].error == budget + D.REASON_SPLIT
|
||||
assert got["t3"].error == budget + "；模型连续 2 段都没写出来，已停止拆小重试，建议更换起草模型"
|
||||
|
||||
ok = answering(tables)
|
||||
model = Scripted(lambda call: ok(call) if "t3" in call.asked else exhausted)
|
||||
got = await D.run_model_draft(model, _source(tables), list(tables))
|
||||
assert _by_table(model) == {"t1": 3, "t2": 3, "t3": 1, "t4": 3}, "t3 写出来了，t4 照常拆小"
|
||||
assert got["t3"].status == "ok" and got["t4"].error == budget + D.REASON_SPLIT
|
||||
|
||||
|
||||
async def test_the_column_parts_of_a_wide_table_stop_splitting_too(monkeypatch):
|
||||
"""宽表的列分段一段一段排队起草（一段拆到底再起草下一段），止损在宽表里同样起作用。
|
||||
|
||||
95 列：表级那段拆不开（结构化、纯文本各一次）；第 2 段 40 列拆到底（15 次结构化 + 8 次纯文本）；这时已经连续 2 段
|
||||
没写出来，第 3、4 段只问一次、换纯文本一次。"""
|
||||
monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1)
|
||||
model = Scripted(_hollow)
|
||||
got = await D.run_model_draft(model, _source({"wide95": _table(95)}), ["wide95"])
|
||||
assert len(model.calls) == 2 + 23 + 2 + 2
|
||||
assert got["wide95"].error == (D.REASON_HOLLOW + D.REASON_SPLIT
|
||||
+ "；模型连续 2 段都没写出来，已停止拆小重试，建议更换起草模型")
|
||||
|
||||
|
||||
async def test_a_wide_table_that_never_comes_back_is_bounded_by_the_first_wave():
|
||||
"""默认并发（3 段同时起草）：294 列一直只给空壳，最多是表级那段 2 次、头 3 个列分段拆到底（各 23 次）、其余 5 段
|
||||
各 2 次。以前所有列分段一下子都排上去、一层一层地拆，要 168 次。"""
|
||||
model = Scripted(_hollow)
|
||||
got = await D.run_model_draft(model, _source({"wide294": _table(294)}), ["wide294"])
|
||||
assert got["wide294"].status == "failed"
|
||||
assert len(model.calls) <= 2 + 3 * 23 + 5 * 2
|
||||
|
||||
|
||||
async def test_a_big_draft_does_not_hold_the_model_ahead_of_a_later_one(monkeypatch):
|
||||
"""全进程的调用名额先来先放：一次起草不把它的所有段一下子都排上去，后开始的那次起草（另一个数据源）能轮上。"""
|
||||
monkeypatch.setattr(catalog, "MODEL_CONCURRENCY", 1)
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
big = {f"a{i}": _table(4) for i in range(1, 5)}
|
||||
small = {"b1": _table(4)}
|
||||
order: list[str] = []
|
||||
release = asyncio.Event()
|
||||
|
||||
def handler(tables: dict) -> Callable[[Call], Any]:
|
||||
ok = answering(tables)
|
||||
|
||||
async def handle(call: Call) -> Reply:
|
||||
order.append(next(iter(call.asked)))
|
||||
if len(order) == 1:
|
||||
await release.wait()
|
||||
return ok(call)
|
||||
return handle
|
||||
|
||||
first = asyncio.create_task(D.run_model_draft(Scripted(handler(big)), _source(big), list(big)))
|
||||
while not order:
|
||||
await asyncio.sleep(0)
|
||||
second = asyncio.create_task(D.run_model_draft(Scripted(handler(small)), _source(small), ["b1"]))
|
||||
for _ in range(20):
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
got_big, got_small = await asyncio.gather(first, second)
|
||||
assert set(_status(got_big).values()) == {"ok"} and _status(got_small) == {"b1": "ok"}
|
||||
assert order == ["a1", "b1", "a2", "a3", "a4"]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 起草要求、用量
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 从正文里取 JSON
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text, tables, truncated", [
|
||||
# 前面有一句话、带括号的说明、围栏
|
||||
('按 {table} 的格式写好了：\n```json\n{"tables": [{"table": "a", "label": "甲"}]}\n```', ["a"], False),
|
||||
('[说明] 草稿：{"tables": [{"table": "a"}]} 以上', ["a"], False),
|
||||
# 外面多包了一层
|
||||
('{"result": {"tables": [{"table": "a"}, {"table": "b"}]}}', ["a", "b"], False),
|
||||
# 截断在字符串、数字、字面量、键名中间：只收写完了的
|
||||
('{"tables": [{"table": "a", "label": "甲"}, {"table": "b", "label": "乙', ["a", "b"], True),
|
||||
('{"tables": [{"table": "a", "columns": [{"name": "x", "label": "某"}, {"name": "y", "len": 12', ["a"], True),
|
||||
('{"tables": [{"table": "a", "keys": ["x", "y"', ["a"], True),
|
||||
('{"tables": [{"table": "a"}, {"tab', ["a"], True),
|
||||
('{"tables": [{"table": "a", "flag": tr', ["a"], True),
|
||||
])
|
||||
def test_json_is_taken_from_the_text(text, tables, truncated):
|
||||
reply = D._parse_text(text, cut=False)
|
||||
assert reply is not None and [e["table"] for e in reply.entries] == tables
|
||||
assert reply.truncated is truncated
|
||||
|
||||
|
||||
def test_cut_values_are_dropped_and_complete_ones_kept():
|
||||
text = '{"tables": [{"table": "a", "label": "甲", "keys": ["x", "y"], "columns": [{"name": "x", "label": "某"}, ' \
|
||||
'{"name": "y", "label": "半'
|
||||
entry = D._parse_text(text, cut=True).entries[0]
|
||||
assert entry == {"table": "a", "label": "甲", "keys": ["x", "y"], "columns": [{"name": "x", "label": "某"}]}
|
||||
cut_keys = D._parse_text('{"tables": [{"table": "a", "label": "甲", "keys": ["x", "y"', cut=True).entries[0]
|
||||
assert cut_keys == {"table": "a", "label": "甲"}, "没写完的业务主键整项不要，不能只留前半"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["好的，我来写", "{坏的 json}", ""])
|
||||
def test_not_json_is_not_a_reply(text):
|
||||
assert D._parse_text(text, cut=False) is None
|
||||
@@ -0,0 +1,604 @@
|
||||
"""数据目录的后台起草任务（data/catalog_jobs.py、/api/datasources/{id}/catalog/draft-jobs）与起草要求。
|
||||
|
||||
- 建任务 202；同一个源已有排队或进行中的任务 409 catalog_draft_busy；起草模型用不了 409
|
||||
catalog_draft_model_unavailable（不建任务）；没有表结构 409；表结构里没有的表照样建进任务、状态 skipped
|
||||
- 在后台起草，每张表一结束就并入目录；进度、计数、用量落库，接口读的是库里的状态
|
||||
- 停止：排队的直接 cancelled；进行中的停掉正在进行的模型调用，没写完的表记 cancelled，已经写进目录的不回滚；已结束的 409
|
||||
- 服务启动时把没结束的任务标成 interrupted（没结束的表记 cancelled、写明服务重启），不自动续跑；关停时同样
|
||||
- 重试：只含原任务里 failed、partial、cancelled 的表，本次说明沿用、起草要求取数据源当前的；没有可重试的 422
|
||||
- 目录清单：每行 last_draft，顶层 guidance、active_job
|
||||
- 起草要求存在数据源 options 里（上传表格的源存在它自己那一行），数据源接口不显示、整份回写 options 时留着，拼连接串时摘掉
|
||||
|
||||
表名、列名、说明文字都是虚构的。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import quote
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.data import catalog, catalog_jobs
|
||||
from app.data.engine import CATALOG_DRAFT_GUIDANCE_OPTION, build_url
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import CatalogDraftJob, CatalogDraftTask, DataSource
|
||||
from app.main import app as _fastapi_app
|
||||
from tests.fixtures.catalog.fake_model import FakeModel, requested
|
||||
from tests.fixtures.sources import drop_source
|
||||
|
||||
ACTOR = {"X-Actor": quote("王敏")}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 夹具
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _col(name: str, type_: str = "INTEGER", comment: str | None = None) -> dict:
|
||||
return {"name": name, "type": type_, "nullable": True, "comment": comment}
|
||||
|
||||
|
||||
def _tables() -> dict:
|
||||
return {
|
||||
"visits": {"qualified": "visits", "columns": [_col("id"), _col("ticket_no", "TEXT"), _col("visit_time", "TEXT"),
|
||||
_col("channel_id")],
|
||||
"primary_key": ["id"], "comment": None, "is_view": False,
|
||||
"foreign_keys": [{"columns": ["channel_id"], "to_table": "channels", "to_columns": ["id"]}]},
|
||||
"channels": {"qualified": "channels", "columns": [_col("id"), _col("name", "TEXT")], "primary_key": ["id"],
|
||||
"comment": "售票渠道", "is_view": False, "foreign_keys": []},
|
||||
"gates": {"qualified": "gates", "columns": [_col("id"), _col("gate_code", "TEXT")], "primary_key": ["id"],
|
||||
"comment": None, "is_view": False, "foreign_keys": []},
|
||||
}
|
||||
|
||||
|
||||
REPLIES = {
|
||||
"visits": {"label": "入园记录", "kind": "fact", "columns": [{"name": "ticket_no", "label": "票号"}]},
|
||||
"channels": {"label": "渠道", "kind": "dimension"},
|
||||
"gates": {"label": "闸口", "kind": "dimension"},
|
||||
}
|
||||
|
||||
_CREATED: list[str] = []
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def _cleanup():
|
||||
yield
|
||||
# 后台任务先收完（停不下来的停掉），再删源：删源会级联删掉任务
|
||||
await catalog_jobs.manager.shutdown(grace=5)
|
||||
while _CREATED:
|
||||
await drop_source(_CREATED.pop())
|
||||
|
||||
|
||||
async def _source(tables: dict | None = None, **kw) -> str:
|
||||
async with SessionLocal() as session:
|
||||
row = DataSource(name=f"cdj_{uuid.uuid4().hex[:10]}", kind="sqlite", database=":memory:",
|
||||
schema_cache={"tables": tables if tables is not None else _tables()}, **kw)
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
_CREATED.append(row.id)
|
||||
return row.id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client():
|
||||
async with AsyncClient(transport=ASGITransport(app=_fastapi_app), base_url="http://test") as c:
|
||||
yield c
|
||||
|
||||
|
||||
def _base(sid: str) -> str:
|
||||
return f"/api/datasources/{sid}/catalog"
|
||||
|
||||
|
||||
def use(monkeypatch, model) -> None:
|
||||
"""起草模型换成假模型（建任务时查一次能不能用，开始起草时再取一次）。"""
|
||||
async def resolve(session):
|
||||
return model, "fake-model"
|
||||
|
||||
monkeypatch.setattr(catalog, "resolve_draft_model", resolve)
|
||||
|
||||
|
||||
async def finished(job_id: str) -> None:
|
||||
await catalog_jobs.manager.wait(job_id, timeout=10)
|
||||
|
||||
|
||||
async def get_job(client, sid: str, job_id: str) -> dict:
|
||||
r = await client.get(f"{_base(sid)}/draft-jobs/{job_id}")
|
||||
assert r.status_code == 200, r.text
|
||||
return r.json()["job"]
|
||||
|
||||
|
||||
class Held:
|
||||
"""假模型：请求里有 hold 里的表时，等 release 才回；被停止时记下这次调用被取消了。"""
|
||||
|
||||
def __init__(self, replies: dict, hold: set[str]) -> None:
|
||||
self.inner, self.hold = FakeModel(replies), set(hold)
|
||||
self.release = asyncio.Event()
|
||||
self.cancelled = 0
|
||||
self.model_name = "held-model"
|
||||
|
||||
def with_structured_output(self, schema, **_):
|
||||
outer = self
|
||||
|
||||
class _S:
|
||||
async def ainvoke(self, messages):
|
||||
text = "\n".join(str(m.content) for m in messages)
|
||||
if outer.hold & set(requested(text)):
|
||||
try:
|
||||
await outer.release.wait()
|
||||
except asyncio.CancelledError:
|
||||
outer.cancelled += 1
|
||||
raise
|
||||
return outer.inner._reply(messages, "structured")
|
||||
|
||||
return _S()
|
||||
|
||||
async def ainvoke(self, messages):
|
||||
return await self.inner.ainvoke(messages)
|
||||
|
||||
|
||||
async def until(predicate, timeout: float = 5.0) -> None:
|
||||
for _ in range(int(timeout / 0.01)):
|
||||
if await predicate():
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("等不到")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 建任务、在后台起草
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_a_job_drafts_in_the_background_and_writes_each_table(client, monkeypatch):
|
||||
sid = await _source()
|
||||
model = FakeModel(REPLIES)
|
||||
use(monkeypatch, model)
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "金额一律按元写单位"})
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", headers=ACTOR,
|
||||
json={"tables": ["visits", "nope", "channels"], "use_model": True, "note": " 看渠道 "})
|
||||
assert r.status_code == 202, r.text
|
||||
job = r.json()["job"]
|
||||
assert job["status"] in ("queued", "running") and job["source_id"] == sid
|
||||
assert job["use_model"] is True and job["guidance"] == "金额一律按元写单位" and job["note"] == "看渠道"
|
||||
assert job["created_by"] == "王敏" and job["total"] == 3
|
||||
assert [t["table_name"] for t in job["tasks"]] == ["visits", "nope", "channels"], "按起草顺序"
|
||||
await finished(job["id"])
|
||||
|
||||
job = await get_job(client, sid, job["id"])
|
||||
assert job["status"] == "succeeded" and job["error"] is None and job["model"] == "fake-model"
|
||||
assert job["done"] == 3 and job["counts"] == {"pending": 0, "running": 0, "ok": 2, "partial": 0, "failed": 0,
|
||||
"skipped": 1, "cancelled": 0}
|
||||
assert job["started_at"] and job["finished_at"] and job["usage"]["calls"] >= 1
|
||||
by = {t["table_name"]: t for t in job["tasks"]}
|
||||
assert by["nope"]["status"] == "skipped" and by["nope"]["error"] == "表结构里没有 nope，请先重新探查结构"
|
||||
visits = by["visits"]
|
||||
assert visits["status"] == "ok" and visits["parts_total"] == 1 and visits["parts_done"] == 1
|
||||
assert visits["model_items"] == 3 and visits["added"] == 4 and visits["version"] == 1 and visits["columns"] == 4
|
||||
assert visits["started_at"] and visits["finished_at"] and visits["error"] is None
|
||||
assert job["added"] == sum(t["added"] for t in job["tasks"])
|
||||
# 起草要求和本次说明都交给了模型
|
||||
assert "起草要求：金额一律按元写单位" in model.calls[0] and "本次说明：看渠道" in model.calls[0]
|
||||
|
||||
detail = (await client.get(f"{_base(sid)}/visits")).json()
|
||||
assert detail["notes"]["label"]["value"] == "入园记录" and detail["updated_by"] == "王敏"
|
||||
|
||||
|
||||
async def test_a_job_without_the_model_writes_structure_only(client):
|
||||
sid = await _source()
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"tables": [], "use_model": False})
|
||||
assert r.status_code == 202, r.text
|
||||
job = r.json()["job"]
|
||||
assert job["total"] == 3, "不给表时按使用次数取前 20 张"
|
||||
await finished(job["id"])
|
||||
job = await get_job(client, sid, job["id"])
|
||||
assert job["status"] == "succeeded" and job["model"] is None and job["usage"]["calls"] == 0
|
||||
assert all(t["parts_total"] is None and t["model_items"] is None for t in job["tasks"])
|
||||
assert {t["table_name"]: t["added"] for t in job["tasks"]}["visits"] == 1 # 外键一条
|
||||
|
||||
|
||||
async def test_a_table_whose_model_part_failed_is_failed_but_structure_is_written(client, monkeypatch):
|
||||
sid = await _source()
|
||||
use(monkeypatch, FakeModel(REPLIES, fail={"visits"}))
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["visits", "channels"],
|
||||
"use_model": True})).json()["job"]
|
||||
await finished(job["id"])
|
||||
job = await get_job(client, sid, job["id"])
|
||||
by = {t["table_name"]: t for t in job["tasks"]}
|
||||
assert job["status"] == "partial"
|
||||
assert by["visits"]["status"] == "failed" and "网关超时" in by["visits"]["model_error"]
|
||||
assert by["visits"]["error"] is None and by["visits"]["added"] == 1, "外键照常写入"
|
||||
assert by["channels"]["status"] == "ok"
|
||||
|
||||
|
||||
class HalfThenHollow:
|
||||
"""假模型：第一次只给前一半列，补问时只给一个空壳（表名对上、一项都没有）。"""
|
||||
|
||||
model_name = "half-model"
|
||||
|
||||
def __init__(self, table: str, reply: dict) -> None:
|
||||
self.table, self.reply, self.calls = table, reply, 0
|
||||
|
||||
def with_structured_output(self, schema, **_):
|
||||
outer = self
|
||||
|
||||
class _S:
|
||||
async def ainvoke(self, messages):
|
||||
outer.calls += 1
|
||||
if outer.calls == 1:
|
||||
cols = outer.reply["columns"]
|
||||
return {"tables": [{"table": outer.table, **outer.reply, "columns": cols[:len(cols) // 2]}]}
|
||||
return {"tables": [{"table": outer.table}]}
|
||||
|
||||
return _S()
|
||||
|
||||
async def ainvoke(self, messages):
|
||||
raise AssertionError("这里不该走纯文本")
|
||||
|
||||
|
||||
async def test_a_followup_that_comes_back_empty_keeps_the_earlier_model_items(client, monkeypatch):
|
||||
"""补问只给空壳：表是 partial、任务是 partial，以前起草过的列（待确认的模型起草项）不因这次没写全被删掉。"""
|
||||
cols = [_col(f"c{i:02d}") for i in range(1, 9)]
|
||||
sid = await _source({"tickets": {"qualified": "tickets", "columns": cols, "primary_key": ["c01"], "comment": None,
|
||||
"is_view": False, "foreign_keys": []}})
|
||||
full = {"label": "门票", "kind": "fact",
|
||||
"columns": [{"name": f"c{i:02d}", "label": f"第{i}列"} for i in range(1, 9)]}
|
||||
use(monkeypatch, FakeModel({"tickets": full}))
|
||||
first = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["tickets"], "use_model": True})).json()
|
||||
await finished(first["job"]["id"])
|
||||
before = (await client.get(f"{_base(sid)}/tickets")).json()["notes"]["columns"]
|
||||
assert {c for c, items in before.items() if "label" in items} == {f"c{i:02d}" for i in range(1, 9)}
|
||||
|
||||
model = HalfThenHollow("tickets", full)
|
||||
use(monkeypatch, model)
|
||||
second = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["tickets"], "use_model": True})).json()
|
||||
await finished(second["job"]["id"])
|
||||
job = await get_job(client, sid, second["job"]["id"])
|
||||
task = job["tasks"][0]
|
||||
assert model.calls == 2
|
||||
assert job["status"] == "partial" and task["status"] == "partial" and task["removed"] == 0
|
||||
assert "有 4 列没写出来" in task["model_error"] and "可以只重试这张表" in task["model_error"]
|
||||
after = (await client.get(f"{_base(sid)}/tickets")).json()["notes"]["columns"]
|
||||
assert {c for c, items in after.items() if "label" in items} == {f"c{i:02d}" for i in range(1, 9)}
|
||||
|
||||
|
||||
async def test_a_job_that_cannot_start_is_failed_with_the_reason(client, monkeypatch):
|
||||
"""建任务时起草模型还能用，开始起草时已经用不了（接入刚被停用）：任务整体 failed，error 写原因，没开始的表也记 failed。"""
|
||||
sid = await _source()
|
||||
calls: list[int] = []
|
||||
|
||||
async def flaky(session):
|
||||
calls.append(1)
|
||||
if len(calls) > 1:
|
||||
raise catalog.CatalogModelUnavailable("起草模型使用的模型接入「虚构接入」已停用")
|
||||
return FakeModel(REPLIES), "fake-model"
|
||||
|
||||
monkeypatch.setattr(catalog, "resolve_draft_model", flaky)
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["visits", "gates"],
|
||||
"use_model": True})).json()["job"]
|
||||
await finished(job["id"])
|
||||
got = await get_job(client, sid, job["id"])
|
||||
assert got["status"] == "failed" and got["error"] == "起草模型使用的模型接入「虚构接入」已停用"
|
||||
assert {t["status"] for t in got["tasks"]} == {"failed"} and got["finished_at"]
|
||||
assert all(t["error"] == got["error"] for t in got["tasks"])
|
||||
retry = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/retry")
|
||||
assert retry.status_code == 409 and retry.json()["code"] == "catalog_draft_model_unavailable"
|
||||
|
||||
|
||||
async def test_busy_source_is_409_and_the_active_job_is_reachable(client, monkeypatch):
|
||||
sid = await _source()
|
||||
model = Held(REPLIES, hold={"visits"})
|
||||
use(monkeypatch, model)
|
||||
first = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["visits"], "use_model": True})).json()
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["gates"], "use_model": False})
|
||||
assert r.status_code == 409 and r.json()["code"] == "catalog_draft_busy", r.text
|
||||
active = (await client.get(f"{_base(sid)}/draft-jobs/active")).json()["job"]
|
||||
assert active["id"] == first["job"]["id"]
|
||||
listed = (await client.get(_base(sid))).json()
|
||||
assert listed["active_job"]["id"] == first["job"]["id"] and listed["active_job"]["total"] == 1
|
||||
assert listed["active_job"]["status"] in ("queued", "running") and listed["active_job"]["done"] == 0
|
||||
model.release.set()
|
||||
await finished(first["job"]["id"])
|
||||
assert (await client.get(f"{_base(sid)}/draft-jobs/active")).json()["job"] is None
|
||||
latest = (await client.get(f"{_base(sid)}/draft-jobs/latest")).json()["job"]
|
||||
assert latest["id"] == first["job"]["id"] and latest["status"] == "succeeded"
|
||||
assert (await client.get(_base(sid))).json()["active_job"] is None
|
||||
|
||||
|
||||
async def test_unavailable_model_is_409_and_no_job_is_created(client, monkeypatch):
|
||||
sid = await _source()
|
||||
|
||||
async def unavailable(session):
|
||||
raise catalog.CatalogModelUnavailable("未配置模型接入，无法由模型起草数据目录。请到「设置 → 模型接入」添加")
|
||||
|
||||
monkeypatch.setattr(catalog, "resolve_draft_model", unavailable)
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"use_model": True})
|
||||
assert r.status_code == 409, r.text
|
||||
assert r.json()["code"] == "catalog_draft_model_unavailable" and "模型接入" in r.json()["detail"]
|
||||
assert (await client.get(f"{_base(sid)}/draft-jobs/latest")).json()["job"] is None
|
||||
|
||||
|
||||
async def test_without_schema_or_with_bad_input_nothing_is_created(client):
|
||||
sid = await _source(tables={})
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"use_model": False})
|
||||
assert r.status_code == 409 and "探查结构" in r.json()["detail"]
|
||||
sid = await _source()
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"use_model": False, "note": "说" * 1001})
|
||||
assert r.status_code == 422 and "1000" in r.json()["detail"]
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs", json={"tables": [f"t{i}" for i in range(501)]})
|
||||
assert r.status_code == 422 and "500" in r.json()["detail"]
|
||||
assert (await client.get(f"{_base(sid)}/draft-jobs/latest")).json()["job"] is None
|
||||
assert (await client.get(f"{_base('missing-source')}/draft-jobs/latest")).status_code == 404
|
||||
assert (await client.get(f"{_base(sid)}/draft-jobs/not-a-job")).status_code == 404
|
||||
|
||||
|
||||
async def test_a_job_of_another_source_is_404(client):
|
||||
a, b = await _source(), await _source()
|
||||
job = (await client.post(f"{_base(a)}/draft-jobs", json={"tables": ["gates"]})).json()["job"]
|
||||
await finished(job["id"])
|
||||
assert (await client.get(f"{_base(b)}/draft-jobs/{job['id']}")).status_code == 404
|
||||
assert (await client.post(f"{_base(b)}/draft-jobs/{job['id']}/cancel")).status_code == 404
|
||||
assert (await client.post(f"{_base(b)}/draft-jobs/{job['id']}/retry")).status_code == 404
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 停止、服务重启
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_cancelling_a_queued_job(client, monkeypatch):
|
||||
sid = await _source()
|
||||
monkeypatch.setattr(catalog_jobs.manager, "start", lambda job_id: None) # 不让它开始
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["gates", "channels"]})).json()["job"]
|
||||
assert job["status"] == "queued"
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/cancel")
|
||||
assert r.status_code == 200, r.text
|
||||
job = r.json()["job"]
|
||||
assert job["status"] == "cancelled" and job["finished_at"]
|
||||
assert {t["status"] for t in job["tasks"]} == {"cancelled"}
|
||||
again = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/cancel")
|
||||
assert again.status_code == 409 and again.json()["code"] == "catalog_draft_finished"
|
||||
|
||||
|
||||
async def test_cancelling_a_running_job_stops_the_call_and_keeps_what_is_written(client, monkeypatch):
|
||||
sid = await _source()
|
||||
model = Held(REPLIES, hold={"visits"})
|
||||
use(monkeypatch, model)
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs",
|
||||
json={"tables": ["channels", "visits"], "use_model": True})).json()["job"]
|
||||
|
||||
async def channels_written() -> bool:
|
||||
tasks = {t["table_name"]: t for t in (await get_job(client, sid, job["id"]))["tasks"]}
|
||||
return tasks["channels"]["status"] == "ok" and tasks["visits"]["status"] == "running"
|
||||
|
||||
await until(channels_written)
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/cancel")
|
||||
assert r.status_code == 200, r.text
|
||||
after = r.json()["job"]
|
||||
by = {t["table_name"]: t for t in after["tasks"]}
|
||||
assert after["status"] == "cancelled" and after["finished_at"]
|
||||
assert by["channels"]["status"] == "ok" and by["visits"]["status"] == "cancelled"
|
||||
assert model.cancelled == 1, "正在进行的模型调用被停掉"
|
||||
detail = (await client.get(f"{_base(sid)}/channels")).json()
|
||||
assert detail["notes"]["kind"]["value"] == "dimension", "已经写进目录的不回滚"
|
||||
assert (await client.get(f"{_base(sid)}/draft-jobs/active")).json()["job"] is None
|
||||
|
||||
|
||||
async def test_startup_marks_unfinished_jobs_interrupted_and_does_not_resume(client):
|
||||
sid = await _source()
|
||||
async with SessionLocal() as session:
|
||||
job = CatalogDraftJob(source_id=sid, status="running", use_model=True)
|
||||
session.add(job)
|
||||
await session.flush()
|
||||
for seq, (name, status) in enumerate([("visits", "ok"), ("channels", "running"), ("gates", "pending")]):
|
||||
session.add(CatalogDraftTask(job_id=job.id, seq=seq, table_name=name, status=status))
|
||||
await session.commit()
|
||||
job_id = job.id
|
||||
assert await catalog_jobs.startup() >= 1
|
||||
got = await get_job(client, sid, job_id)
|
||||
by = {t["table_name"]: t for t in got["tasks"]}
|
||||
assert got["status"] == "interrupted" and got["error"] == "服务重启，起草中断" and got["finished_at"]
|
||||
assert by["visits"]["status"] == "ok" and by["visits"]["error"] is None
|
||||
assert by["channels"]["status"] == by["gates"]["status"] == "cancelled"
|
||||
assert by["channels"]["error"] == by["gates"]["error"] == "服务重启，起草中断"
|
||||
assert not catalog_jobs.manager.running(job_id), "不自动续跑"
|
||||
|
||||
|
||||
async def test_shutdown_stops_running_jobs_as_interrupted(client, monkeypatch):
|
||||
sid = await _source()
|
||||
model = Held(REPLIES, hold={"visits"})
|
||||
use(monkeypatch, model)
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["visits"], "use_model": True})).json()["job"]
|
||||
|
||||
async def running() -> bool:
|
||||
return (await get_job(client, sid, job["id"]))["tasks"][0]["status"] == "running"
|
||||
|
||||
await until(running)
|
||||
await catalog_jobs.manager.shutdown(grace=5)
|
||||
got = await get_job(client, sid, job["id"])
|
||||
assert got["status"] == "interrupted" and got["tasks"][0]["status"] == "cancelled"
|
||||
assert got["tasks"][0]["error"] == "服务重启，起草中断"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 重试
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_retry_takes_only_failed_partial_and_cancelled_tables(client, monkeypatch):
|
||||
sid = await _source()
|
||||
use(monkeypatch, FakeModel(REPLIES, fail={"visits"}))
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "旧的要求"})
|
||||
first = (await client.post(f"{_base(sid)}/draft-jobs", json={
|
||||
"tables": ["visits", "channels", "nope"], "use_model": True, "note": "第一次"})).json()["job"]
|
||||
await finished(first["id"])
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "新的要求"})
|
||||
assert (await get_job(client, sid, first["id"]))["guidance"] == "旧的要求", "任务记下的是当时的快照"
|
||||
|
||||
use(monkeypatch, FakeModel(REPLIES))
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs/{first['id']}/retry", headers=ACTOR)
|
||||
assert r.status_code == 202, r.text
|
||||
retry = r.json()["job"]
|
||||
assert [t["table_name"] for t in retry["tasks"]] == ["visits"]
|
||||
assert retry["use_model"] is True and retry["note"] == "第一次" and retry["guidance"] == "新的要求"
|
||||
assert retry["retry_of"] == first["id"] and retry["created_by"] == "王敏"
|
||||
await finished(retry["id"])
|
||||
assert (await get_job(client, sid, retry["id"]))["status"] == "succeeded"
|
||||
|
||||
nothing = await client.post(f"{_base(sid)}/draft-jobs/{retry['id']}/retry")
|
||||
assert nothing.status_code == 422 and "没有" in nothing.json()["detail"]
|
||||
|
||||
|
||||
async def test_retry_while_busy_is_409(client, monkeypatch):
|
||||
sid = await _source()
|
||||
monkeypatch.setattr(catalog_jobs.manager, "start", lambda job_id: None)
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["gates"]})).json()["job"]
|
||||
await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/cancel")
|
||||
other = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["channels"]})).json()["job"]
|
||||
assert other["status"] == "queued"
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/retry")
|
||||
assert r.status_code == 409 and r.json()["code"] == "catalog_draft_busy"
|
||||
|
||||
|
||||
async def test_retry_with_an_unavailable_model_is_409(client, monkeypatch):
|
||||
sid = await _source()
|
||||
monkeypatch.setattr(catalog_jobs.manager, "start", lambda job_id: None)
|
||||
use(monkeypatch, FakeModel(REPLIES))
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["gates"], "use_model": True})).json()["job"]
|
||||
await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/cancel")
|
||||
|
||||
async def unavailable(session):
|
||||
raise catalog.CatalogModelUnavailable("当前只有演示用的模型接入，无法用于起草数据目录")
|
||||
|
||||
monkeypatch.setattr(catalog, "resolve_draft_model", unavailable)
|
||||
r = await client.post(f"{_base(sid)}/draft-jobs/{job['id']}/retry")
|
||||
assert r.status_code == 409 and r.json()["code"] == "catalog_draft_model_unavailable"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 目录清单上的起草状态
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_the_list_carries_last_draft_per_table(client, monkeypatch):
|
||||
sid = await _source()
|
||||
use(monkeypatch, FakeModel(REPLIES, fail={"visits"}))
|
||||
monkeypatch.setattr(catalog, "MODEL_BATCH_TABLES", 1)
|
||||
first = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["visits", "channels"],
|
||||
"use_model": True})).json()["job"]
|
||||
await finished(first["id"])
|
||||
second = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["channels"]})).json()["job"]
|
||||
await finished(second["id"])
|
||||
body = (await client.get(_base(sid))).json()
|
||||
rows = {r["table_name"]: r for r in body["tables"]}
|
||||
assert rows["gates"]["last_draft"] is None
|
||||
visits = rows["visits"]["last_draft"]
|
||||
assert visits["status"] == "failed" and visits["job_id"] == first["id"] and visits["finished_at"]
|
||||
assert "网关超时" in visits["model_error"] and visits["error"] is None
|
||||
assert rows["channels"]["last_draft"]["job_id"] == second["id"], "取最近一次任务里的那一项"
|
||||
assert body["guidance"] is None and body["active_job"] is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 起草要求
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_guidance_is_saved_trimmed_cleared_and_limited(client):
|
||||
sid = await _source()
|
||||
r = await client.put(f"{_base(sid)}/guidance", json={"guidance": " 票务库，金额按元。\n状态 9 是作废 "})
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json() == {"guidance": "票务库，金额按元。\n状态 9 是作废"}
|
||||
assert (await client.get(_base(sid))).json()["guidance"] == "票务库，金额按元。\n状态 9 是作废"
|
||||
too_long = await client.put(f"{_base(sid)}/guidance", json={"guidance": "字" * 2001})
|
||||
assert too_long.status_code == 422 and "2000" in too_long.json()["detail"]
|
||||
assert (await client.get(_base(sid))).json()["guidance"] == "票务库，金额按元。\n状态 9 是作废"
|
||||
assert (await client.put(f"{_base(sid)}/guidance", json={"guidance": "字" * 2000})).status_code == 200
|
||||
cleared = await client.put(f"{_base(sid)}/guidance", json={"guidance": " "})
|
||||
assert cleared.json() == {"guidance": None}
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(DataSource, sid)
|
||||
assert CATALOG_DRAFT_GUIDANCE_OPTION not in (row.options or {})
|
||||
assert (await client.put(f"{_base('missing-source')}/guidance", json={"guidance": "x"})).status_code == 404
|
||||
|
||||
|
||||
async def test_a_table_named_guidance_can_still_be_saved(client):
|
||||
"""起草要求的路径和「整份提交一张表」撞名：表名恰好叫 guidance 时，{notes, if_version} 照旧是那张表的提交，
|
||||
不会被当成清除起草要求；两样都没写的 422。"""
|
||||
tables = _tables()
|
||||
tables["guidance"] = {"qualified": "guidance", "columns": [_col("id"), _col("body", "TEXT")],
|
||||
"primary_key": ["id"], "comment": None, "is_view": False, "foreign_keys": []}
|
||||
sid = await _source(tables)
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "金额按元"})
|
||||
r = await client.put(f"{_base(sid)}/guidance", json={"notes": {"label": {"value": "操作指引"}}, "if_version": 0},
|
||||
headers=ACTOR)
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["table_name"] == "guidance" and r.json()["notes"]["label"]["value"] == "操作指引"
|
||||
assert (await client.get(_base(sid))).json()["guidance"] == "金额按元"
|
||||
for body in ({}, {"guidance": 3}, {"guidance": ["金额按元"]}):
|
||||
r = await client.put(f"{_base(sid)}/guidance", json=body)
|
||||
assert r.status_code == 422, body
|
||||
assert isinstance(r.json()["detail"], str) and "起草要求" in r.json()["detail"], "说明是中文的一句话"
|
||||
assert (await client.get(_base(sid))).json()["guidance"] == "金额按元"
|
||||
|
||||
|
||||
async def test_guidance_of_an_uploaded_source_lives_on_its_own_row(client):
|
||||
sid = await _source(origin="upload")
|
||||
r = await client.put(f"{_base(sid)}/guidance", json={"guidance": "按期上传的客流日报"})
|
||||
assert r.status_code == 200, r.text
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(DataSource, sid)
|
||||
assert row.options[CATALOG_DRAFT_GUIDANCE_OPTION] == "按期上传的客流日报"
|
||||
|
||||
|
||||
async def test_the_datasource_form_neither_shows_nor_drops_the_guidance(client):
|
||||
sid = await _source(options={"schema": "main"})
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "金额按元"})
|
||||
shown = (await client.get("/api/datasources")).json()
|
||||
mine = next(s for s in shown if s["id"] == sid)
|
||||
assert CATALOG_DRAFT_GUIDANCE_OPTION not in mine["options"] and mine["options"]["schema"] == "main"
|
||||
# 表单整份回写 options（不带起草要求，或者带着一份别的）：起草要求留着，只认 PUT /catalog/guidance
|
||||
r = await client.patch(f"/api/datasources/{sid}", json={"options": {"schema": "other",
|
||||
CATALOG_DRAFT_GUIDANCE_OPTION: "篡改"}})
|
||||
assert r.status_code == 200, r.text
|
||||
assert CATALOG_DRAFT_GUIDANCE_OPTION not in r.json()["options"]
|
||||
assert (await client.get(_base(sid))).json()["guidance"] == "金额按元"
|
||||
async with SessionLocal() as session:
|
||||
assert (await session.get(DataSource, sid)).options == {"schema": "other",
|
||||
CATALOG_DRAFT_GUIDANCE_OPTION: "金额按元"}
|
||||
|
||||
|
||||
def test_the_guidance_never_reaches_the_driver():
|
||||
src = SimpleNamespace(kind="mysql", host="db.example", port=3306, database="shop", username="reader",
|
||||
password=None, options={"charset": "utf8mb4", CATALOG_DRAFT_GUIDANCE_OPTION: "金额按元"})
|
||||
url = build_url(src)
|
||||
assert "charset=utf8mb4" in url and CATALOG_DRAFT_GUIDANCE_OPTION not in url
|
||||
|
||||
|
||||
async def test_the_sync_draft_uses_the_source_guidance(client, monkeypatch):
|
||||
sid = await _source(description="虚构的景区票务库")
|
||||
model = FakeModel(REPLIES)
|
||||
use(monkeypatch, model)
|
||||
await client.put(f"{_base(sid)}/guidance", json={"guidance": "渠道表的名称是对外展示名"})
|
||||
r = await client.post(f"{_base(sid)}/draft", json={"tables": ["channels"], "use_model": True})
|
||||
assert r.status_code == 200, r.text
|
||||
assert "数据源说明：虚构的景区票务库\n起草要求：渠道表的名称是对外展示名" in model.calls[0]
|
||||
assert set(r.json()) == {"tables", "total", "model_used", "model", "model_error"}, "旧的同步接口返回形状不变"
|
||||
|
||||
|
||||
async def test_jobs_go_away_with_their_source(client):
|
||||
sid = await _source()
|
||||
job = (await client.post(f"{_base(sid)}/draft-jobs", json={"tables": ["gates"]})).json()["job"]
|
||||
await finished(job["id"])
|
||||
await drop_source(sid)
|
||||
_CREATED.remove(sid)
|
||||
async with SessionLocal() as session:
|
||||
assert (await session.execute(select(CatalogDraftJob).where(CatalogDraftJob.id == job["id"]))).first() is None
|
||||
assert (await session.execute(select(CatalogDraftTask).where(CatalogDraftTask.job_id == job["id"]))).first() \
|
||||
is None
|
||||
@@ -0,0 +1,270 @@
|
||||
"""起草模型（设置里的 catalog_draft 一组）与起草取模型的顺序。
|
||||
|
||||
- 设置组 catalog_draft：{provider, model}，经 PUT /api/settings 写、GET 读，没存过补 {provider: null, model: null}；
|
||||
provider 不是字符串、点名的接入不存在或已停用 → 422，说法写设置页上的叫法「起草模型 · 接入」「起草模型 · 模型」
|
||||
- 取模型的顺序（数据目录起草和按配方导入的 AI 起草共用 draft_model.draft_model_choice）：起草模型填了接入或模型就整组
|
||||
用它；否则沿用助手的模型设置；再没有就是默认接入。不跨级拼接入和模型
|
||||
- 数据目录起草不关思考（按接入默认）、额度 32768，服务方嫌大退回 8192；配方 AI 起草只改取模型的顺序，额度、思考不动
|
||||
|
||||
接入名、模型名都是虚构的。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from sqlalchemy import delete
|
||||
|
||||
from app.data import catalog, draft_model
|
||||
from app.data import recipe_ai as A
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Provider, Setting
|
||||
from app.main import app
|
||||
|
||||
LIVE, OFF, GONE = "起草用接入", "起草用停用接入", "起草用不存在的接入"
|
||||
|
||||
|
||||
async def _clean() -> None:
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key.in_([draft_model.CATALOG_DRAFT_SETTING_KEY])))
|
||||
await session.execute(delete(Provider).where(Provider.name.in_([LIVE, OFF])))
|
||||
await session.commit()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def providers():
|
||||
"""一个启用、一个停用的演示类接入（不联网）；用完删掉。"""
|
||||
await _clean()
|
||||
async with SessionLocal() as session:
|
||||
session.add_all([
|
||||
Provider(name=LIVE, kind="mock", models=[{"id": "起草模型甲"}], default_model="起草模型甲", enabled=True,
|
||||
extra={}),
|
||||
Provider(name=OFF, kind="mock", models=[{"id": "起草模型乙"}], default_model="起草模型乙", enabled=False,
|
||||
extra={}),
|
||||
])
|
||||
await session.commit()
|
||||
yield
|
||||
await _clean()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client():
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
|
||||
yield c
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 设置组的读写与校验
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_the_group_is_exposed_with_empty_defaults(client, providers):
|
||||
got = (await client.get("/api/settings")).json()
|
||||
assert got["catalog_draft"] == {"provider": None, "model": None}
|
||||
|
||||
|
||||
async def test_the_group_round_trips(client, providers):
|
||||
r = await client.put("/api/settings", json={"values": {"catalog_draft": {"provider": LIVE, "model": "起草模型甲"}}})
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["catalog_draft"] == {"provider": LIVE, "model": "起草模型甲"}
|
||||
got = (await client.get("/api/settings")).json()["catalog_draft"]
|
||||
assert got == {"provider": LIVE, "model": "起草模型甲"}
|
||||
# 只写模型、接入留空：照存，接入是 null
|
||||
r = await client.put("/api/settings", json={"values": {"catalog_draft": {"provider": " ", "model": "x"}}})
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["catalog_draft"] == {"provider": None, "model": "x"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value,label", [
|
||||
({"provider": 3}, "起草模型 · 接入"),
|
||||
({"model": ["a"]}, "起草模型 · 模型"),
|
||||
({"temperature": 0}, "temperature"),
|
||||
("起草模型甲", "起草模型"),
|
||||
])
|
||||
async def test_bad_values_are_refused_with_the_settings_page_wording(client, providers, value, label):
|
||||
r = await client.put("/api/settings", json={"values": {"catalog_draft": value, "ui": {"theme": "dark"}}})
|
||||
assert r.status_code == 422, r.text
|
||||
assert label in r.json()["detail"]
|
||||
got = (await client.get("/api/settings")).json()
|
||||
assert got["catalog_draft"] == {"provider": None, "model": None} and got["ui"]["theme"] != "dark"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name,reason,fix", [
|
||||
(GONE, f"模型接入「{GONE}」不存在", "请在「设置 → 模型接入」中核对名称"),
|
||||
(OFF, f"模型接入「{OFF}」已停用", "请在「设置 → 模型接入」中启用它"),
|
||||
])
|
||||
async def test_a_missing_or_disabled_provider_is_refused_with_the_field(client, providers, name, reason, fix):
|
||||
r = await client.put("/api/settings", json={"values": {"catalog_draft": {"provider": name}}})
|
||||
assert r.status_code == 422, r.text
|
||||
body = r.json()
|
||||
assert body["code"] == "settings_provider_invalid" and body["field"] == "catalog_draft.provider"
|
||||
assert body["detail"] == f"「起草模型 · 接入」指定的{reason}。{fix}，或换一个接入"
|
||||
async with SessionLocal() as session:
|
||||
assert await session.get(Setting, "catalog_draft") is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 取模型的顺序
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _MemDB:
|
||||
"""独立的内存库：取模型只看这里的接入和设置，不受别的测试建的接入影响。"""
|
||||
|
||||
async def __aenter__(self):
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
import app.db.models # noqa: F401
|
||||
from app.db.base import Base
|
||||
|
||||
self.engine = create_async_engine("sqlite+aiosqlite:///:memory:", poolclass=StaticPool)
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
self.session = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)()
|
||||
return self.session
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
await self.session.close()
|
||||
await self.engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def mem():
|
||||
async with _MemDB() as session:
|
||||
yield session
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_env_keys(monkeypatch):
|
||||
for k in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "OPENAI_API_KEY", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.delenv(k, raising=False)
|
||||
|
||||
|
||||
async def _setup(session, *, draft: dict | None = None, copilot: dict | None = None) -> None:
|
||||
from app.core.crypto import encrypt
|
||||
|
||||
for name, model in (("甲接入", "模型甲"), ("乙接入", "模型乙"), ("丙接入", "模型丙")):
|
||||
session.add(Provider(name=name, kind="anthropic", api_key=encrypt("sk-test"), models=[{"id": model}],
|
||||
default_model=model, enabled=True, extra={}))
|
||||
if draft is not None:
|
||||
session.add(Setting(key=draft_model.CATALOG_DRAFT_SETTING_KEY, value=draft))
|
||||
if copilot is not None:
|
||||
session.add(Setting(key=draft_model.COPILOT_SETTING_KEY, value=copilot))
|
||||
await session.commit()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("draft,copilot,expected", [
|
||||
(None, None, (None, None, "default")),
|
||||
(None, {"provider": "乙接入", "model": ""}, ("乙接入", None, "assistant")),
|
||||
({"provider": None, "model": None}, {"provider": "", "model": "模型乙"}, (None, "模型乙", "assistant")),
|
||||
({"provider": "", "model": "模型丙"}, {"provider": "乙接入", "model": "模型乙"}, (None, "模型丙", "draft")),
|
||||
({"provider": "甲接入", "model": None}, {"provider": "乙接入", "model": "模型乙"}, ("甲接入", None, "draft")),
|
||||
({"provider": 7}, {"provider": "乙接入"}, ("乙接入", None, "assistant")), # 存坏了的按没填
|
||||
])
|
||||
async def test_choice_order(mem, draft, copilot, expected):
|
||||
await _setup(mem, draft=draft, copilot=copilot)
|
||||
got = await draft_model.draft_model_choice(mem)
|
||||
assert (got.provider, got.model, got.origin) == expected
|
||||
|
||||
|
||||
def _spy(monkeypatch) -> list[Any]:
|
||||
from app.providers import factory
|
||||
|
||||
specs: list[Any] = []
|
||||
|
||||
class Built:
|
||||
def __init__(self, provider, spec):
|
||||
self.provider, self.spec = provider.name, spec
|
||||
|
||||
def fake_build(provider, spec):
|
||||
specs.append((provider.name, spec))
|
||||
return Built(provider, spec)
|
||||
|
||||
monkeypatch.setattr(factory, "build_chat_model", fake_build)
|
||||
return specs
|
||||
|
||||
|
||||
async def test_catalog_draft_uses_the_draft_setting_keeps_thinking_and_a_big_budget(mem, no_env_keys, monkeypatch):
|
||||
await _setup(mem, draft={"provider": "甲接入", "model": None}, copilot={"provider": "乙接入", "model": "模型乙"})
|
||||
specs = _spy(monkeypatch)
|
||||
model, model_id = await catalog.resolve_draft_model(mem)
|
||||
assert model_id == "模型甲"
|
||||
assert [p for p, _ in specs] == ["甲接入", "甲接入"]
|
||||
assert [s.max_tokens for _, s in specs] == [catalog.MODEL_MAX_TOKENS, catalog.MODEL_FALLBACK_MAX_TOKENS]
|
||||
assert catalog.MODEL_MAX_TOKENS == 32768 and catalog.MODEL_FALLBACK_MAX_TOKENS == 8192
|
||||
assert all(s.thinking is None and s.temperature == 0 for _, s in specs), "不关思考，按接入默认"
|
||||
assert all(s.timeout == catalog.MODEL_TIMEOUT_S == 420 for _, s in specs)
|
||||
assert all(s.max_retries == 0 for _, s in specs), "重试只由起草引擎按失败原因来做，SDK 自己不再重试"
|
||||
assert model.budget == 32768 and model.model_name == "模型甲"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,extra", [
|
||||
("openai_compatible", {}),
|
||||
("anthropic", {}),
|
||||
("anthropic", {"auth_style": "bearer"}),
|
||||
], ids=["openai-compatible", "anthropic", "anthropic-bearer"])
|
||||
def test_the_built_model_does_not_retry_on_its_own_when_asked(kind, extra):
|
||||
"""ModelSpec.max_retries 一直传到 SDK 客户端（Bearer 鉴权换上的客户端也是）；不写就照 SDK 默认。"""
|
||||
from app.core.crypto import encrypt
|
||||
from app.providers.factory import ModelSpec, build_chat_model
|
||||
|
||||
provider = Provider(name="虚构接入", kind=kind, api_key=encrypt("sk-test"), base_url="https://gateway.example/v1",
|
||||
models=[], default_model="model-x", extra=extra)
|
||||
|
||||
def client_retries(chat: Any) -> int:
|
||||
client = chat.root_async_client if kind == "openai_compatible" else chat._async_client
|
||||
return client.max_retries
|
||||
|
||||
assert client_retries(build_chat_model(provider, ModelSpec(model="model-x", max_retries=0))) == 0
|
||||
assert client_retries(build_chat_model(provider, ModelSpec(model="model-x"))) == 2
|
||||
|
||||
|
||||
async def test_catalog_draft_falls_back_to_the_assistant_model(mem, no_env_keys, monkeypatch):
|
||||
await _setup(mem, copilot={"provider": "", "model": "模型乙"})
|
||||
specs = _spy(monkeypatch)
|
||||
_, model_id = await catalog.resolve_draft_model(mem)
|
||||
assert model_id == "模型乙" and specs[0][0] == "乙接入"
|
||||
|
||||
|
||||
async def test_recipe_ai_draft_follows_the_same_order_but_keeps_its_own_spec(mem, no_env_keys, monkeypatch):
|
||||
await _setup(mem, draft={"provider": "", "model": "模型丙"}, copilot={"provider": "乙接入", "model": "模型乙"})
|
||||
specs = _spy(monkeypatch)
|
||||
model, model_id, avail = await A.resolve_draft_model(mem)
|
||||
assert avail.available and model_id == "模型丙" and avail.provider == "丙接入"
|
||||
assert [s.max_tokens for _, s in specs] == [A.AI_MAX_TOKENS, A.AI_FALLBACK_MAX_TOKENS]
|
||||
assert all(s.thinking == "off" and s.temperature == 0 for _, s in specs), "配方 AI 起草的思考设置不动"
|
||||
assert all(s.max_retries is None for _, s in specs), "配方 AI 起草的重试照旧（SDK 默认）"
|
||||
|
||||
|
||||
async def test_a_broken_draft_setting_says_where_to_fix_it(mem, no_env_keys):
|
||||
"""起草模型点名的接入不存在：报错说的是起草模型，不是助手的模型。沿用助手模型时照旧说助手。"""
|
||||
from app.core.crypto import encrypt
|
||||
|
||||
mem.add(Provider(name="停用的接入", kind="anthropic", api_key=encrypt("k"), models=[{"id": "模型丁"}],
|
||||
default_model="模型丁", enabled=False, extra={}))
|
||||
mem.add(Setting(key=draft_model.CATALOG_DRAFT_SETTING_KEY, value={"provider": "停用的接入"}))
|
||||
await mem.commit()
|
||||
with pytest.raises(catalog.CatalogModelUnavailable) as e:
|
||||
await catalog.resolve_draft_model(mem)
|
||||
assert "起草模型" in str(e.value) and "停用的接入" in str(e.value) and "已停用" in str(e.value)
|
||||
assert "助手" not in str(e.value)
|
||||
_, _, avail = await A.resolve_draft_model(mem)
|
||||
assert not avail.available and "起草模型" in avail.reason and "助手" not in avail.reason
|
||||
|
||||
await mem.execute(delete(Setting))
|
||||
mem.add(Setting(key=draft_model.COPILOT_SETTING_KEY, value={"provider": "停用的接入"}))
|
||||
await mem.commit()
|
||||
with pytest.raises(catalog.CatalogModelUnavailable) as e:
|
||||
await catalog.resolve_draft_model(mem)
|
||||
assert "助手使用的" in str(e.value)
|
||||
|
||||
|
||||
async def test_only_a_demo_provider_is_refused(mem, no_env_keys):
|
||||
mem.add(Provider(name="演示", kind="mock", models=[{"id": "mock-fast"}], default_model="mock-fast", enabled=True,
|
||||
extra={}))
|
||||
await mem.commit()
|
||||
with pytest.raises(catalog.CatalogModelUnavailable) as e:
|
||||
await catalog.resolve_draft_model(mem)
|
||||
assert "演示" in str(e.value)
|
||||
Reference in New Issue
Block a user