mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 02:11:58 +00:00
ThearchyHelios
975dde1146
- 均分：超过 40 列的表，列分段数是列数除以 40 向上取整，各段列数平均（相差不超过 1），不再留下 2 到 8 列的尾段（实测 会思考的模型常对这种小段只给空壳）；表级那段、窄表拼段、parts_total 不变。 - 拆不开的段（不到 10 列）只给空壳：再问一次，结构化输出换纯文本，已经走纯文本的原样再问，仍是空壳才记没写出来。 - 超时当作这一段太大：不再原样重试同样大小的段，按额度用完的办法拆小再问；网络错误、429、5xx 仍原样退避重试。 - 段大小随模型调整：一次起草里额度用完或超时累计 2 次，之后新开始的段列数上限减半（40 → 20 → 10，最低 10），新开始 的段超过上限就先按上限均分开（多表的段按表重新拼）；任务记当前的上限 part_columns。 - 封顶：一段（连同它拆出来的、补问的、重试的）最多调用 8 次，到了次数还没写完的记没写出来、原因写清。止损照旧。 - 提醒：最近 6 次调用有 4 次以上没写出来（额度用完、超时、空壳、出错），任务的 notice 写一句提醒（次数、原因，建议在 「设置 → 偏好设置」的「起草模型」换一个模型、减少表或停止）；之后最近 6 次里少于 2 次就清掉。 - 用量：进度回调 usage 换成 call_done（CallReport：用量、没写出来的原因种类、当前的提醒和列数上限），任务的用量加 timeouts、failed_calls；超时的调用 token 照实记 0。 - 留空的列：请模型写了、模型也列出了，中文名和含义却都没给（或者给的都没通过校验）的列记在 ModelOutcome.blank，逐表 记 blank_columns 和最多 5 个列名 blank_sample，任务上是合计，目录清单的 last_draft 也带列数；表的状态不因此变。已确认 的、库注释已经说清的列不算。补问时把漏的列列出来、字段却留空的，算留空，不再记成「漏了这些列」。 - 公平排队：宽表的表级那段写完，它的列分段排到已经排着的段后面（先来先服务），不再插到队首让窄表一直等。 - 原因：几段因为同一个原因没写出来时说「这几段」，同一段拆出来的几块仍是「这一段」；只有真拆小过才补「已拆小重试仍 没写出」（截断后只补剩下的不算）；每段写上列的范围，例如「第 5 段（列名 到 列名，共 33 列）」，窄表写出是哪几列。 - 库：catalog_draft_jobs 加 timeouts、failed_calls、part_columns、notice，catalog_draft_tasks 加 blank_columns、 blank_sample，旧库启动时补列（存量回填 0 和空）。假模型加 hang（头几次调用一直不回，用来模拟超时）。
244 lines
12 KiB
Python
244 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from typing import Any, AsyncIterator
|
||
|
||
from sqlalchemy import JSON, DateTime, MetaData, event
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||
from sqlalchemy.types import TypeDecorator
|
||
|
||
from app.core.config import settings
|
||
|
||
NAMING_CONVENTION = {
|
||
"ix": "ix_%(column_0_label)s",
|
||
"uq": "uq_%(table_name)s_%(column_0_name)s",
|
||
"fk": "fk_%(table_name)s_%(column_0_name)s",
|
||
"pk": "pk_%(table_name)s",
|
||
}
|
||
|
||
|
||
def utcnow() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def new_id() -> str:
|
||
return uuid.uuid4().hex
|
||
|
||
|
||
class UTCDateTime(TypeDecorator):
|
||
"""存 UTC、读出来也是带时区的 UTC。
|
||
|
||
SQLite 没有时区类型：DateTime(timezone=True) 写进去的只是一串不带偏移的
|
||
字面时间，读回来是 naive 的。API 于是吐出 '2026-09-26T01:11:19'，浏览器把
|
||
不带偏移的 ISO 串当本地时间解析，运行列表整体慢了 8 小时。在这一层补上 UTC，
|
||
pydantic 序列化自然带 +00:00，所有接口一次修好。
|
||
|
||
写入前先换算成 UTC 再去掉时区：SQLite 方言格式化时直接丢弃 tzinfo，
|
||
一个 +08:00 的时间原样写进去，读回来就差出 8 小时。
|
||
"""
|
||
|
||
impl = DateTime
|
||
cache_ok = True
|
||
|
||
def process_bind_param(self, value: datetime | None, dialect: Any) -> datetime | None:
|
||
if value is None:
|
||
return None
|
||
if value.tzinfo is not None:
|
||
value = value.astimezone(timezone.utc).replace(tzinfo=None)
|
||
return value
|
||
|
||
def process_result_value(self, value: datetime | None, dialect: Any) -> datetime | None:
|
||
if value is None:
|
||
return None
|
||
if value.tzinfo is None:
|
||
return value.replace(tzinfo=timezone.utc)
|
||
return value.astimezone(timezone.utc)
|
||
|
||
|
||
class Base(DeclarativeBase):
|
||
metadata = MetaData(naming_convention=NAMING_CONVENTION)
|
||
|
||
# 图定义、节点配置、事件负载都是任意形状的 JSON，统一映射到 SQLite 的 JSON 列
|
||
type_annotation_map = {dict[str, Any]: JSON, list[Any]: JSON, datetime: UTCDateTime}
|
||
|
||
|
||
class TimestampMixin:
|
||
created_at: Mapped[datetime] = mapped_column(UTCDateTime, default=utcnow)
|
||
updated_at: Mapped[datetime] = mapped_column(UTCDateTime, default=utcnow, onupdate=utcnow)
|
||
|
||
|
||
def _stop_aiosqlite_once() -> None:
|
||
"""aiosqlite 连接的工作线程只停一次；之后再要求停的，拿到的是同一个结果。
|
||
|
||
aiosqlite 0.22 的 stop() 每调一次都往工作线程队列里放一个"关闭并退出线程"
|
||
的请求，而线程处理完第一个就退出了——之后放进去的请求没人处理，等它的一方
|
||
永远挂着。close() 的 finally 里会调 stop()，SQLAlchemy 的强制 terminate 也会
|
||
直接调；两个 close 并发、或者 close 撞上强制 terminate，就会卡死。实测同一
|
||
连接并发 close 两次，10 次卡 10 次。
|
||
|
||
在这里触发它的路径：并行的两条支路一条失败，LangGraph 取消另一条，而那一条
|
||
正在写工件。取消打进进行中的数据库操作，SQLAlchemy 在两条路径上各 terminate
|
||
一次同一条连接——节点任务取消不完，整个运行卡在 running，占着的并发名额也不
|
||
还。大约五次里有一次。连接池和 checkpointer 都走 aiosqlite，补在这一层都管得到。
|
||
"""
|
||
import weakref
|
||
|
||
import aiosqlite
|
||
|
||
original = aiosqlite.Connection.stop
|
||
if getattr(original, "_agentlab_once", False):
|
||
return
|
||
stopping: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
|
||
|
||
def stop(self): # type: ignore[no-untyped-def]
|
||
if self not in stopping:
|
||
stopping[self] = original(self)
|
||
return stopping[self]
|
||
|
||
stop._agentlab_once = True # type: ignore[attr-defined]
|
||
aiosqlite.Connection.stop = stop # type: ignore[method-assign]
|
||
|
||
|
||
_stop_aiosqlite_once()
|
||
# pool_pre_ping：取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时
|
||
# （_close_connection 记完日志会把 CancelledError 原样抛出，跳过了把连接记录
|
||
# 清空的那一行），已经关掉的连接会留在池里，下一个取到它的人拿到的是
|
||
# "no active connection"。探活失败会被当成断开，换一条新的——SQLite 本地一次
|
||
# SELECT 1 的代价，换池子不被污染。
|
||
engine = create_async_engine(settings.db_url, echo=False, future=True, pool_pre_ping=True)
|
||
|
||
|
||
@event.listens_for(engine.sync_engine, "connect")
|
||
def _set_sqlite_pragma(dbapi_conn, _record): # pragma: no cover - 驱动回调
|
||
"""WAL + 外键约束。WAL 让读写不互相阻塞，事件写入和前端轮询能并行。"""
|
||
cur = dbapi_conn.cursor()
|
||
cur.execute("PRAGMA journal_mode=WAL")
|
||
cur.execute("PRAGMA foreign_keys=ON")
|
||
cur.execute("PRAGMA busy_timeout=5000")
|
||
cur.close()
|
||
|
||
|
||
SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
||
|
||
|
||
async def get_session() -> AsyncIterator[AsyncSession]:
|
||
"""FastAPI 依赖注入用的会话。"""
|
||
async with SessionLocal() as session:
|
||
yield session
|
||
|
||
|
||
# 轻量迁移：create_all 只建新表，不给已有表加列。SQLite 的 ADD COLUMN
|
||
# 便宜且带默认值回填，够用；等真需要改列类型再考虑 alembic。
|
||
_COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [
|
||
("runs", "run_class", "TEXT DEFAULT 'exploratory'"),
|
||
("runs", "version", "INTEGER"),
|
||
("runs", "version_hash", "VARCHAR(64)"),
|
||
("runs", "manifest_hash", "VARCHAR(64)"),
|
||
("runs", "manifest_seq", "INTEGER"),
|
||
("runs", "started_by", "VARCHAR(100)"),
|
||
# 续跑、恢复要沿用第一次发起时的记忆域、知识库和审批默认值；以前只在第一段
|
||
# _drive 里传过一次，恢复后回落成 "default"，同一次运行前后两段查的不是同一个库
|
||
("runs", "memory_scope", "VARCHAR(100)"),
|
||
("runs", "collection", "VARCHAR(100)"),
|
||
("runs", "approval_default", "VARCHAR(20)"),
|
||
("runs", "tool_trust", "JSON"),
|
||
("runs", "agent_limits", "JSON"),
|
||
("runs", "error_node_id", "VARCHAR(64)"),
|
||
("workflows", "status", "TEXT DEFAULT 'draft'"),
|
||
("workflows", "published_version", "INTEGER"),
|
||
("workflows", "published_by", "VARCHAR(100)"),
|
||
("workflow_versions", "graph_hash", "VARCHAR(64)"),
|
||
# 存量版本回填成 NULL：不知道当时按哪一档发布的，正式运行照旧看工作流的 status
|
||
("workflow_versions", "level", "VARCHAR(20)"),
|
||
# 发布时的数据目录版本（阶段 4A）。存量版本回填成 NULL：不知道当时按哪一版目录发布的，正式运行不提醒
|
||
("workflow_versions", "catalog_versions", "JSON"),
|
||
("approvals", "resolved_by", "VARCHAR(100)"),
|
||
# 存量向量没有这两列，回填成空串 / 0 —— 正好表示"不知道是谁建的"，
|
||
# 检索侧会把它们当成对不上当前 embedder 处理
|
||
("chunks", "embed_model", "VARCHAR(100) DEFAULT ''"),
|
||
("chunks", "embed_dim", "INTEGER DEFAULT 0"),
|
||
("memories", "embed_model", "VARCHAR(100) DEFAULT ''"),
|
||
("memories", "embed_dim", "INTEGER DEFAULT 0"),
|
||
("chunks", "token_len", "INTEGER DEFAULT 0"),
|
||
("documents", "status", "TEXT DEFAULT 'ready'"),
|
||
("documents", "error", "TEXT DEFAULT ''"),
|
||
# 存量轮次没有复核结论。NULL 正好表示"这一轮没复核过"，和"复核过、没问题"
|
||
# （存 verdict=ok）是两回事，前端要分得开
|
||
("conversation_turns", "review", "JSON"),
|
||
("conversation_turns", "meta", "JSON"),
|
||
("data_sources", "last_check", "JSON"),
|
||
("mcp_servers", "last_check", "JSON"),
|
||
# 存量的源一律先记成 manual；老的上传源由 table_versions.startup 按路径认出来，
|
||
# 补上 v0 版本后再改成 upload
|
||
("data_sources", "origin", "TEXT DEFAULT 'manual'"),
|
||
("data_sources", "current_snapshot_id", "VARCHAR(64)"),
|
||
# 存量运行没有钉过版本，NULL 正好表示「升级前发起的」
|
||
("runs", "data_versions", "JSON"),
|
||
# 按配方导入（期 2）。新表 table_recipes、import_stagings 由 create_all 建；这些是已有表的新列，
|
||
# 存量行一律 NULL：存量的源都是简单导入或手工源，存量的导入记录都不是按配方导入的
|
||
("data_sources", "current_recipe_id", "VARCHAR(32)"),
|
||
("table_imports", "recipe_id", "VARCHAR(32)"),
|
||
("table_imports", "period_start", "VARCHAR(10)"),
|
||
("table_imports", "period_end", "VARCHAR(10)"),
|
||
("table_imports", "context", "JSON"),
|
||
("table_imports", "checks", "JSON"),
|
||
("table_imports", "overrides", "JSON"),
|
||
("table_imports", "waivers", "JSON"),
|
||
("table_imports", "confirmations", "JSON"),
|
||
("table_imports", "manifest_artifact", "VARCHAR(64)"),
|
||
("table_imports", "signed_by", "VARCHAR(100)"),
|
||
("table_imports", "staging_id", "VARCHAR(32)"),
|
||
# 期 3：按期累积、版本页（P3-SPEC 9.6）。新表 snapshot_activations 由 create_all 建；存量行一律 NULL，
|
||
# 含义是「期 3 之前的」：快照的 mode 按 replace、recipe_id 取最后一期导入的、activated_at 按 created_at
|
||
("source_snapshots", "mode", "VARCHAR(20)"),
|
||
("source_snapshots", "recipe_id", "VARCHAR(32)"),
|
||
("source_snapshots", "activated_at", "DATETIME"),
|
||
("source_snapshots", "manifest_artifact", "VARCHAR(64)"),
|
||
("table_imports", "revoked", "JSON"),
|
||
("import_stagings", "edits", "JSON"),
|
||
("import_stagings", "answers_dropped", "JSON"),
|
||
("import_stagings", "redraft_sha256", "VARCHAR(64)"),
|
||
# 运行选项（发起时的快照，整次运行沿用）和发起它的定时任务。存量运行回填 {} 和 NULL：
|
||
# 没有选项，也不是定时任务发起的
|
||
("runs", "options", "JSON DEFAULT '{}'"),
|
||
("runs", "schedule_id", "VARCHAR(32)"),
|
||
# 数据目录起草任务（验收后修正）：超时、没写出来的调用次数，当前的列数上限，提醒；逐表的模型留空的列。
|
||
# 存量任务回填 0 和 NULL：升级前没记这些
|
||
("catalog_draft_jobs", "timeouts", "INTEGER DEFAULT 0"),
|
||
("catalog_draft_jobs", "failed_calls", "INTEGER DEFAULT 0"),
|
||
("catalog_draft_jobs", "part_columns", "INTEGER"),
|
||
("catalog_draft_jobs", "notice", "TEXT"),
|
||
("catalog_draft_tasks", "blank_columns", "INTEGER DEFAULT 0"),
|
||
("catalog_draft_tasks", "blank_sample", "JSON"),
|
||
]
|
||
|
||
#: 迁移补上的列要建的索引：(索引名, 表, 列)。create_all 不碰已有的表，表上新列的索引也就不会建。
|
||
#: 索引名照 NAMING_CONVENTION 取，和新库 create_all 建出来的同名，IF NOT EXISTS 才不会重复建
|
||
_INDEX_MIGRATIONS: list[tuple[str, str, str]] = [
|
||
# 定时任务的触发历史按它找这个任务发起的运行
|
||
("ix_runs_schedule_id", "runs", "schedule_id"),
|
||
]
|
||
|
||
|
||
def _migrate(conn) -> None: # pragma: no cover - 同步回调
|
||
from sqlalchemy import text
|
||
|
||
for table, column, ddl in _COLUMN_MIGRATIONS:
|
||
cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))}
|
||
if cols and column not in cols:
|
||
conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}"))
|
||
for name, table, column in _INDEX_MIGRATIONS:
|
||
cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))}
|
||
if column in cols:
|
||
conn.execute(text(f"CREATE INDEX IF NOT EXISTS {name} ON {table} ({column})"))
|
||
|
||
|
||
async def init_db() -> None:
|
||
from app.db import models # noqa: F401 确保模型已注册到 metadata
|
||
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
await conn.run_sync(_migrate)
|