mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 01:11:59 +00:00
ThearchyHelios
2e7d6103ea
- 配置：database_url、checkpoint_url（启动时校验写法，只认 PostgreSQL）、连接池四项；启动日志写用的是哪种库（主机、库名、用户，不写密码） - 引擎（db/base.make_engine）：SQLite 照旧 WAL、外键、忙等；PostgreSQL 用连接池、探活、会话时区 UTC，装上方言调整（db/postgres.py）：建表时 String(n) 是不限长的 VARCHAR（SQLite 从不查长度，导过来的存量也可能超长）、ORDER BY 里可空列的 NULL 位置照 SQLite 补写（升序在前、降序在后，窗口函数里同样）；Integer 列在 PostgreSQL 上是 BIGINT（列挂到表上时换成 with_variant，绑定参数跟着列按 BIGINT 转换）；JSON 列两边都存原文（PostgreSQL 的 json，不用 JSONB：JSONB 重排对象里的键，快照的列序、表的先后读回来就变了）；JSON 里的 NaN 写成 null、NUL 去掉（PostgreSQL 不收） - 迁移按方言：已有的列和索引用 inspector 认；PostgreSQL 上补列的类型取模型里这一列的、缺省值取迁移表里的（布尔 0/1 换成 false/true），ADD COLUMN IF NOT EXISTS；去重语句里的布尔按方言写；部分唯一索引补 postgresql_where；init_db 在 PostgreSQL 上拿咨询锁（几个进程同时启动不撞） - 应用库上的 SQLite 专用写法改成按方言编译的表达式（db/portable.py）：json_extract 的十几处（口径召回、出具、对账、个人应用、幂等键、访问控制、数据源变化通知、剖析任务、待办超时）换成 json_eq / json_flag / json_text，json_type 换成 json_is_string（两种库同一套带类型的比较）；能力登记的 sqlite insert 换成按会话方言的 insert_ignore；能力使用次数的 json_each 按方言分支；附件「没绑记录」的条件不靠 json 的等号；按报错文字认唯一约束的地方换成 unique_violation（按表和列认，两种库的报错都认得） - 断点（engine/checkpoints.py）：PostgreSQL 时开 psycopg 连接池、在咨询锁里 setup；runner 只认 BaseCheckpointSaver，读断点仍只走 saver 的公开接口 - 依赖：新 extra postgres（asyncpg、psycopg[binary] 3、psycopg-pool、langgraph-checkpoint-postgres 3.1.2）；langgraph-checkpoint 卡在 <4.3（断点格式就在这个包里，postgres 包新版要 4.3，不动已装的 4.2.0）；锁文件补上新包 - 运行收尾：记 action.unknown、安排核实做完才把运行摘掉（以前先摘掉，等「不在跑了」再读事件的一方在库慢时读不到这一条） - 测试：AGENTLAB_TEST_DATABASE_URL 时跑在 PostgreSQL 上（每个 xdist 进程建自己的库、按事件循环分的连接池 tests/pgpool.py），不设时盖掉环境里的数据库地址照旧用 SQLite；scripts/test-postgres.sh 起临时的 postgres:16 容器跑全量；sqlite_only 标记；另开连接改库的几个竞态测试、触发器测试改成两种库都能跑（tests/appdb.py）；新增 test_db_portable 守着方言调整、JSON 键序、超长字符串、大整数、NULL 排序、时区、PostgreSQL 老库迁移
88 lines
4.0 KiB
Python
88 lines
4.0 KiB
Python
"""LangGraph 的断点库：应用库是 SQLite 时用 AsyncSqliteSaver（数据目录里的 checkpoints.db，和业务表分开免得抢写锁），
|
||
是 PostgreSQL 时用官方的 AsyncPostgresSaver（psycopg 3 + 连接池），放在同一个库里，AGENTLAB_CHECKPOINT_URL 另给就放那里。
|
||
|
||
它建的四张表（checkpoints、checkpoint_blobs、checkpoint_writes、checkpoint_migrations）和应用的表不重名
|
||
（test_db_portable 守着），所以同一个库、同一个 schema 就行。
|
||
|
||
应用里读断点一律走 saver 的公开接口（aget_state / aget_state_history / aget_tuple / alist），不直接查这几张表：
|
||
两种库的表结构完全不一样（SQLite 整份断点是一个 msgpack，PostgreSQL 拆成 JSONB 加 blob）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
from collections.abc import AsyncIterator
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||
|
||
from app.core.config import settings
|
||
|
||
#: 建断点表时拿的会话级咨询锁的号：几个进程同时启动，各自的 setup（建表、记迁移版本）一个一个来
|
||
_SETUP_LOCK = 0x4167656E74436B70 # "AgentCkp"
|
||
|
||
|
||
def _pool_size() -> int:
|
||
# 运行里每个超步写一次断点、每个任务写一次 pending writes，都是一来一回的短操作。并发运行多时多给几条
|
||
return max(4, min(settings.max_concurrent_runs, 20))
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def open_postgres_saver(conninfo: str, *, max_size: int | None = None,
|
||
application_name: str = "agentlab") -> AsyncIterator[BaseCheckpointSaver]:
|
||
"""PostgreSQL 的断点库：开连接池、建表（幂等），退出时关池子。application_name 同应用库（db/base.APPLICATION_NAME）。"""
|
||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||
from psycopg.rows import dict_row
|
||
from psycopg_pool import AsyncConnectionPool
|
||
|
||
# 连接参数照 AsyncPostgresSaver.from_conn_string：自动提交（它不自己管事务）、语句第一次执行就预备、行按 dict 取
|
||
pool = AsyncConnectionPool(
|
||
conninfo, min_size=1, max_size=max_size or _pool_size(), open=False,
|
||
kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row,
|
||
"application_name": application_name},
|
||
check=AsyncConnectionPool.check_connection, name="agentlab-checkpoints")
|
||
await pool.open(wait=True)
|
||
try:
|
||
saver = AsyncPostgresSaver(pool) # type: ignore[arg-type]
|
||
async with pool.connection() as conn:
|
||
await conn.execute("SELECT pg_advisory_lock(%s)", (_SETUP_LOCK,))
|
||
try:
|
||
await saver.setup()
|
||
finally:
|
||
await conn.execute("SELECT pg_advisory_unlock(%s)", (_SETUP_LOCK,))
|
||
yield saver
|
||
finally:
|
||
await pool.close()
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def open_sqlite_saver(path: Path | str) -> AsyncIterator[BaseCheckpointSaver]:
|
||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||
|
||
async with AsyncSqliteSaver.from_conn_string(str(path)) as saver:
|
||
yield saver
|
||
|
||
|
||
def open_saver() -> contextlib.AbstractAsyncContextManager[BaseCheckpointSaver]:
|
||
"""按配置开这个进程的断点库。"""
|
||
conninfo = settings.checkpoint_conninfo
|
||
if conninfo:
|
||
return open_postgres_saver(conninfo)
|
||
return open_sqlite_saver(settings.checkpoint_path)
|
||
|
||
|
||
def kind() -> str:
|
||
""""sqlite" 或 "postgresql"：断点库放在哪种库里。"""
|
||
return "postgresql" if settings.checkpoint_conninfo else "sqlite"
|
||
|
||
|
||
async def count(saver: Any, thread_id: str | None = None) -> tuple[int, int]:
|
||
"""(断点条数, pending writes 条数)。thread_id 为空数全部。copy-to-postgres 核对用，走公开接口。"""
|
||
config = {"configurable": {"thread_id": thread_id}} if thread_id else None
|
||
checkpoints = writes = 0
|
||
async for tup in saver.alist(config):
|
||
checkpoints += 1
|
||
writes += len(tup.pending_writes or ())
|
||
return checkpoints, writes
|