92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
"""
|
||
Checkpoint 工厂：PostgreSQL 优先，开发可降级内存。
|
||
|
||
本文件职责：按 CHECKPOINT_BACKEND / PG_DSN 创建 LangGraph checkpointer。
|
||
副作用：postgres 模式会 setup() 建 checkpoint 表（仅当前 PG_DATABASE）。
|
||
禁止：把 checkpoint 写入 Redis；禁止指向现网 ytd_runtime。
|
||
线程：Worker 进程启动时创建一次，供多槽共享（同一 thread_id 写冲突由 runtime 锁）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from dataclasses import dataclass
|
||
from typing import Any, Optional
|
||
|
||
from agent.config import Settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class CheckpointBundle:
|
||
"""
|
||
checkpointer 与可选连接资源。
|
||
|
||
postgres 模式下保留 context/cm 以便进程退出时关闭连接池。
|
||
"""
|
||
|
||
checkpointer: Any
|
||
backend: str
|
||
_closer: Optional[Any] = None
|
||
|
||
def close(self) -> None:
|
||
"""释放 postgres 连接（若有）。无副作用于业务数据。"""
|
||
if self._closer is None:
|
||
return
|
||
try:
|
||
# from_conn_string 返回上下文管理器
|
||
self._closer.__exit__(None, None, None)
|
||
except Exception: # noqa: BLE001 — 关闭失败只记日志
|
||
logger.exception("关闭 Postgres checkpointer 失败")
|
||
self._closer = None
|
||
|
||
def create_checkpointer(settings: Settings) -> CheckpointBundle:
|
||
"""
|
||
创建 checkpointer。
|
||
|
||
- CHECKPOINT_BACKEND=memory：仅本地骨架冒烟，进程重启丢书签。
|
||
- postgres：需要可连的 resolve_pg_dsn()；失败时若 ytd_env=test 可降级 memory 并打警告。
|
||
"""
|
||
backend = (settings.checkpoint_backend or "postgres").strip().lower()
|
||
if backend == "memory":
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
|
||
logger.warning("Graph checkpoint 使用 MemorySaver（仅骨架/本机；正式必须 postgres）")
|
||
return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")
|
||
|
||
dsn = settings.resolve_pg_dsn()
|
||
if not dsn:
|
||
if (settings.ytd_env or "").lower() == "prod":
|
||
raise RuntimeError("正式环境必须配置 PG_DSN 或 PG_HOST/PG_DATABASE/PG_USER")
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
|
||
logger.warning(
|
||
"未配置 PG_DSN，ytd_env=%s 降级 MemorySaver；正式切流前必须接 inquiry_robot_runtime",
|
||
settings.ytd_env,
|
||
)
|
||
return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")
|
||
|
||
try:
|
||
from langgraph.checkpoint.postgres import PostgresSaver
|
||
|
||
# from_conn_string 返回上下文管理器；进入后得到 saver，并需保持打开
|
||
cm = PostgresSaver.from_conn_string(dsn)
|
||
saver = cm.__enter__()
|
||
saver.setup()
|
||
logger.info(
|
||
"Graph Postgres checkpointer 就绪 database=%s",
|
||
settings.pg_database or "(from DSN)",
|
||
)
|
||
return CheckpointBundle(checkpointer=saver, backend="postgres", _closer=cm)
|
||
except Exception as exc: # noqa: BLE001
|
||
if (settings.ytd_env or "").lower() == "prod":
|
||
raise
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
|
||
logger.warning(
|
||
"Postgres checkpointer 初始化失败，降级 MemorySaver：%s",
|
||
exc,
|
||
)
|
||
return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")
|