494 lines
19 KiB
Python
494 lines
19 KiB
Python
"""
|
||
LangGraph Checkpointer 配置模块
|
||
|
||
支持多种持久化后端：
|
||
1. MemorySaver - 内存存储（仅用于测试）
|
||
2. PostgresSaver - PostgreSQL 数据库（生产环境）
|
||
|
||
使用 message_id 和 chat_id 作为 thread_id 来恢复工作流状态
|
||
"""
|
||
|
||
import os
|
||
import asyncio
|
||
from typing import Optional
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
from config import POSTGRES_CONNECTION_STRING
|
||
|
||
CHECKPOINTER_TYPE = "postgres"
|
||
|
||
|
||
def get_thread_id(chat_id: str, message_id: str) -> str:
|
||
"""
|
||
生成唯一的 thread_id
|
||
|
||
Args:
|
||
chat_id: 聊天会话 ID
|
||
message_id: 消息 ID
|
||
|
||
Returns:
|
||
组合后的唯一 thread_id
|
||
"""
|
||
return f"{chat_id}:{message_id}"
|
||
|
||
|
||
def parse_thread_id(thread_id: str) -> tuple[str, str]:
|
||
"""
|
||
解析 thread_id 获取 chat_id 和 message_id
|
||
|
||
Args:
|
||
thread_id: 组合的 thread_id
|
||
|
||
Returns:
|
||
(chat_id, message_id) 元组
|
||
"""
|
||
parts = thread_id.split(":", 1)
|
||
if len(parts) == 2:
|
||
return parts[0], parts[1]
|
||
return parts[0], ""
|
||
|
||
|
||
async def verify_checkpointer_tables_exist() -> bool:
|
||
"""
|
||
验证 checkpointer 所需的数据库表是否存在且结构正确
|
||
|
||
Returns:
|
||
True 如果所有必需的表都存在且结构正确，否则 False
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError:
|
||
return False
|
||
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
# 检查 checkpoints 表
|
||
await cur.execute("""
|
||
SELECT EXISTS (
|
||
SELECT FROM information_schema.tables
|
||
WHERE table_schema = 'public'
|
||
AND table_name = 'checkpoints'
|
||
)
|
||
""")
|
||
checkpoints_exists = (await cur.fetchone())[0]
|
||
|
||
# 检查 checkpoint_writes 表
|
||
await cur.execute("""
|
||
SELECT EXISTS (
|
||
SELECT FROM information_schema.tables
|
||
WHERE table_schema = 'public'
|
||
AND table_name = 'checkpoint_writes'
|
||
)
|
||
""")
|
||
writes_exists = (await cur.fetchone())[0]
|
||
|
||
# 检查 checkpoint_blobs 表
|
||
await cur.execute("""
|
||
SELECT EXISTS (
|
||
SELECT FROM information_schema.tables
|
||
WHERE table_schema = 'public'
|
||
AND table_name = 'checkpoint_blobs'
|
||
)
|
||
""")
|
||
blobs_exists = (await cur.fetchone())[0]
|
||
|
||
if not (checkpoints_exists and writes_exists and blobs_exists):
|
||
return False
|
||
|
||
# 检查 checkpoints 表的 checkpoint 字段是否为 JSONB 类型
|
||
await cur.execute("""
|
||
SELECT data_type
|
||
FROM information_schema.columns
|
||
WHERE table_name = 'checkpoints'
|
||
AND column_name = 'checkpoint'
|
||
""")
|
||
result = await cur.fetchone()
|
||
if not result or result[0] != 'jsonb':
|
||
print(f"[Checkpointer] checkpoints.checkpoint column type is {result[0] if result else 'NULL'}, expected 'jsonb'")
|
||
return False
|
||
|
||
# 检查 checkpoint_writes 表是否有 task_path 列
|
||
await cur.execute("""
|
||
SELECT EXISTS (
|
||
SELECT FROM information_schema.columns
|
||
WHERE table_name = 'checkpoint_writes'
|
||
AND column_name = 'task_path'
|
||
)
|
||
""")
|
||
task_path_exists = (await cur.fetchone())[0]
|
||
if not task_path_exists:
|
||
print(f"[Checkpointer] checkpoint_writes.task_path column does not exist")
|
||
return False
|
||
|
||
# 检查 checkpoint_writes 表的主键是否正确
|
||
await cur.execute("""
|
||
SELECT COUNT(*)
|
||
FROM information_schema.table_constraints tc
|
||
JOIN information_schema.key_column_usage kcu
|
||
ON tc.constraint_name = kcu.constraint_name
|
||
WHERE tc.table_name = 'checkpoint_writes'
|
||
AND tc.constraint_type = 'PRIMARY KEY'
|
||
""")
|
||
pk_count = (await cur.fetchone())[0]
|
||
if pk_count != 5: # 应该有 5 个主键列
|
||
print(f"[Checkpointer] checkpoint_writes has {pk_count} primary key columns, expected 5")
|
||
return False
|
||
|
||
return True
|
||
except Exception as e:
|
||
print(f"[Checkpointer] Error verifying tables: {e}")
|
||
return False
|
||
|
||
|
||
async def init_checkpointer_db():
|
||
"""
|
||
初始化 PostgreSQL checkpointer 数据库表
|
||
使用 AsyncPostgresSaver 的 setup 方法创建正确的表结构
|
||
"""
|
||
if CHECKPOINTER_TYPE != "postgres":
|
||
return
|
||
|
||
try:
|
||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||
from psycopg_pool import AsyncConnectionPool
|
||
import psycopg
|
||
except ImportError as e:
|
||
print(f"[Checkpointer] ERROR: Missing required package: {e}")
|
||
print(f"[Checkpointer] Please install: pip install langgraph-checkpoint-postgres psycopg-pool")
|
||
raise ImportError(
|
||
"langgraph-checkpoint-postgres or psycopg-pool not installed. "
|
||
"Run: pip install langgraph-checkpoint-postgres psycopg-pool"
|
||
) from e
|
||
|
||
# 首先检查表是否已存在
|
||
tables_exist = await verify_checkpointer_tables_exist()
|
||
if tables_exist:
|
||
print(f"[Checkpointer] Database tables already exist, skipping initialization")
|
||
return
|
||
|
||
print(f"[Checkpointer] Initializing PostgreSQL database tables...")
|
||
|
||
# 尝试使用 AsyncPostgresSaver 的 setup 方法
|
||
try:
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
checkpointer = AsyncPostgresSaver(pool)
|
||
await checkpointer.setup()
|
||
print(f"[Checkpointer] AsyncPostgresSaver.setup() executed")
|
||
except Exception as e:
|
||
print(f"[Checkpointer] AsyncPostgresSaver.setup() failed: {e}")
|
||
|
||
# 验证表是否真的被创建
|
||
tables_exist_after = await verify_checkpointer_tables_exist()
|
||
|
||
if not tables_exist_after:
|
||
print(f"[Checkpointer] Tables not created by setup(), using direct SQL...")
|
||
try:
|
||
await create_checkpointer_tables_directly()
|
||
except Exception as e2:
|
||
print(f"[Checkpointer] Direct SQL initialization also failed: {e2}")
|
||
raise
|
||
else:
|
||
print(f"[Checkpointer] PostgreSQL database tables verified successfully")
|
||
|
||
|
||
async def create_checkpointer_tables_directly():
|
||
"""
|
||
直接通过 SQL 语句创建 checkpointer 所需的数据库表
|
||
这是备用方案，当 AsyncPostgresSaver.setup() 失败时使用
|
||
|
||
注意：LangGraph AsyncPostgresSaver 使用 JSONB 类型存储数据
|
||
表名必须是: checkpoints, checkpoint_writes, checkpoint_blobs
|
||
"""
|
||
try:
|
||
from psycopg_pool import AsyncConnectionPool
|
||
except ImportError as e:
|
||
print(f"[Checkpointer] ERROR: Missing psycopg-pool: {e}")
|
||
raise
|
||
|
||
print(f"[Checkpointer] Creating checkpointer tables using direct SQL...")
|
||
|
||
async with AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
kwargs={"autocommit": True}
|
||
) as pool:
|
||
async with pool.connection() as conn:
|
||
async with conn.cursor() as cur:
|
||
# 先删除旧表（如果存在错误的表结构）
|
||
await cur.execute("DROP TABLE IF EXISTS checkpoints CASCADE")
|
||
await cur.execute("DROP TABLE IF EXISTS writes CASCADE")
|
||
await cur.execute("DROP TABLE IF EXISTS checkpoint_blobs CASCADE")
|
||
await cur.execute("DROP TABLE IF EXISTS checkpoint_writes CASCADE")
|
||
await cur.execute("DROP TABLE IF EXISTS blobs CASCADE")
|
||
await cur.execute("DROP TABLE IF EXISTS checkpoint_migrations CASCADE")
|
||
print(f"[Checkpointer] Dropped existing tables")
|
||
|
||
# 创建 checkpoint_migrations 表（用于跟踪迁移版本）
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS checkpoint_migrations (
|
||
v INTEGER PRIMARY KEY
|
||
)
|
||
""")
|
||
|
||
# 创建 checkpoints 表 - 使用 JSONB 类型
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS checkpoints (
|
||
thread_id TEXT NOT NULL,
|
||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||
checkpoint_id TEXT NOT NULL,
|
||
parent_checkpoint_id TEXT,
|
||
type TEXT,
|
||
checkpoint JSONB NOT NULL DEFAULT '{}',
|
||
metadata JSONB NOT NULL DEFAULT '{}',
|
||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
|
||
)
|
||
""")
|
||
|
||
# 创建 checkpoint_writes 表 - 使用 BYTEA 类型
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS checkpoint_writes (
|
||
thread_id TEXT NOT NULL,
|
||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||
checkpoint_id TEXT NOT NULL,
|
||
task_id TEXT NOT NULL,
|
||
task_path TEXT NOT NULL DEFAULT '',
|
||
idx INTEGER NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
type TEXT,
|
||
blob BYTEA,
|
||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||
)
|
||
""")
|
||
|
||
# 创建 checkpoint_blobs 表 - 使用 BYTEA 类型存储二进制数据
|
||
await cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS checkpoint_blobs (
|
||
thread_id TEXT NOT NULL,
|
||
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
||
channel TEXT NOT NULL,
|
||
version TEXT NOT NULL,
|
||
type TEXT NOT NULL,
|
||
blob BYTEA,
|
||
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
|
||
)
|
||
""")
|
||
|
||
# 创建索引以提高查询性能
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id)
|
||
""")
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id)
|
||
""")
|
||
await cur.execute("""
|
||
CREATE INDEX IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id)
|
||
""")
|
||
|
||
# 插入迁移版本记录
|
||
await cur.execute("""
|
||
INSERT INTO checkpoint_migrations (v) VALUES (1) ON CONFLICT DO NOTHING
|
||
""")
|
||
|
||
print(f"[Checkpointer] Checkpointer tables created successfully via direct SQL")
|
||
|
||
|
||
def create_checkpointer(checkpointer_type: str = None):
|
||
"""
|
||
创建同步 checkpointer 实例
|
||
|
||
Args:
|
||
checkpointer_type: checkpointer 类型，可选 "memory", "postgres"
|
||
|
||
Returns:
|
||
checkpointer 实例
|
||
"""
|
||
checkpointer_type = checkpointer_type or CHECKPOINTER_TYPE
|
||
|
||
if checkpointer_type == "memory":
|
||
return MemorySaver()
|
||
|
||
elif checkpointer_type == "postgres":
|
||
try:
|
||
from langgraph.checkpoint.postgres import PostgresSaver
|
||
from psycopg import Connection
|
||
if not POSTGRES_CONNECTION_STRING:
|
||
raise ValueError("POSTGRES_CONNECTION_STRING 环境变量未设置")
|
||
conn = Connection.connect(POSTGRES_CONNECTION_STRING)
|
||
return PostgresSaver(conn)
|
||
except ImportError:
|
||
print("[WARNING] psycopg 或 langgraph.checkpoint.postgres 未安装，回退到 memory")
|
||
return MemorySaver()
|
||
|
||
else:
|
||
print(f"[WARNING] 未知的 checkpointer 类型: {checkpointer_type}，使用 memory")
|
||
return MemorySaver()
|
||
|
||
|
||
def create_async_checkpointer(checkpointer_type: str = None):
|
||
"""
|
||
创建异步 checkpointer 实例
|
||
|
||
Args:
|
||
checkpointer_type: checkpointer 类型
|
||
|
||
Returns:
|
||
异步 checkpointer 实例
|
||
"""
|
||
checkpointer_type = checkpointer_type or CHECKPOINTER_TYPE
|
||
print(f"[Checkpointer] 创建异步 checkpointer, 类型配置: {checkpointer_type}")
|
||
|
||
if checkpointer_type == "memory":
|
||
print("[Checkpointer] 使用 MemorySaver")
|
||
return MemorySaver()
|
||
|
||
elif checkpointer_type == "postgres":
|
||
try:
|
||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||
from psycopg_pool import AsyncConnectionPool
|
||
if not POSTGRES_CONNECTION_STRING:
|
||
raise ValueError("POSTGRES_CONNECTION_STRING 环境变量未设置")
|
||
pool = AsyncConnectionPool(POSTGRES_CONNECTION_STRING)
|
||
return AsyncPostgresSaver(pool)
|
||
except ImportError:
|
||
print("[WARNING] psycopg_pool 或 langgraph.checkpoint.postgres 未安装，回退到 memory")
|
||
return MemorySaver()
|
||
|
||
else:
|
||
print(f"[WARNING] 未知的 checkpointer 类型: {checkpointer_type}，使用 memory")
|
||
return MemorySaver()
|
||
|
||
|
||
async def get_or_create_async_checkpointer(checkpointer_type: str = None):
|
||
"""
|
||
获取或创建异步 checkpointer 实例（正确处理上下文管理器）
|
||
|
||
Args:
|
||
checkpointer_type: checkpointer 类型
|
||
|
||
Returns:
|
||
已初始化的异步 checkpointer 实例
|
||
"""
|
||
checkpointer = create_async_checkpointer(checkpointer_type)
|
||
|
||
if hasattr(checkpointer, '__aenter__'):
|
||
checkpointer = await checkpointer.__aenter__()
|
||
|
||
return checkpointer
|
||
|
||
|
||
class CheckpointerManager:
|
||
"""
|
||
Checkpointer 管理器 - 单例模式
|
||
|
||
用于管理多个工作流的 checkpointer 实例
|
||
"""
|
||
|
||
_instance = None
|
||
_sync_checkpointer = None
|
||
_async_conn_pool = None
|
||
_async_checkpointer = None
|
||
_initialized = False
|
||
|
||
def __new__(cls):
|
||
if cls._instance is None:
|
||
cls._instance = super().__new__(cls)
|
||
return cls._instance
|
||
|
||
def get_checkpointer(self):
|
||
"""获取同步 checkpointer 实例"""
|
||
if self._sync_checkpointer is None:
|
||
self._sync_checkpointer = create_checkpointer()
|
||
return self._sync_checkpointer
|
||
|
||
async def get_async_checkpointer(self):
|
||
"""
|
||
获取异步 checkpointer 实例
|
||
"""
|
||
if self._async_checkpointer is None:
|
||
if CHECKPOINTER_TYPE == "postgres":
|
||
try:
|
||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||
from psycopg_pool import AsyncConnectionPool
|
||
if not POSTGRES_CONNECTION_STRING:
|
||
raise ValueError("POSTGRES_CONNECTION_STRING not set")
|
||
|
||
self._async_conn_pool = AsyncConnectionPool(
|
||
POSTGRES_CONNECTION_STRING,
|
||
open=False
|
||
)
|
||
await self._async_conn_pool.open()
|
||
self._async_checkpointer = AsyncPostgresSaver(self._async_conn_pool)
|
||
|
||
# 检查表是否存在
|
||
tables_exist = await verify_checkpointer_tables_exist()
|
||
|
||
if not tables_exist:
|
||
# 尝试使用 setup() 方法
|
||
try:
|
||
await self._async_checkpointer.setup()
|
||
print(f"[Checkpointer] AsyncPostgresSaver.setup() executed")
|
||
except Exception as setup_e:
|
||
print(f"[Checkpointer] AsyncPostgresSaver.setup() failed: {setup_e}")
|
||
|
||
# 再次验证表是否被创建
|
||
tables_exist_after = await verify_checkpointer_tables_exist()
|
||
|
||
if not tables_exist_after:
|
||
print(f"[Checkpointer] Tables not created by setup(), using direct SQL...")
|
||
await create_checkpointer_tables_directly()
|
||
else:
|
||
print(f"[Checkpointer] Database tables already exist")
|
||
|
||
print(f"[Checkpointer] AsyncPostgresSaver created, connection: {POSTGRES_CONNECTION_STRING.split('@')[1] if '@' in POSTGRES_CONNECTION_STRING else POSTGRES_CONNECTION_STRING}")
|
||
except ImportError as e:
|
||
print(f"[WARNING] psycopg_pool or langgraph.checkpoint.postgres not installed: {e}, fallback to memory")
|
||
self._async_checkpointer = MemorySaver()
|
||
except Exception as e:
|
||
print(f"[Checkpointer] ERROR: Failed to create AsyncPostgresSaver: {e}")
|
||
raise
|
||
elif CHECKPOINTER_TYPE == "memory":
|
||
self._async_checkpointer = MemorySaver()
|
||
print("[Checkpointer] Using MemorySaver")
|
||
else:
|
||
self._async_checkpointer = MemorySaver()
|
||
print(f"[Checkpointer] Unknown type {CHECKPOINTER_TYPE}, using MemorySaver")
|
||
|
||
print(f"[Checkpointer] Async checkpointer type: {type(self._async_checkpointer).__name__}")
|
||
return self._async_checkpointer
|
||
|
||
async def setup(self):
|
||
"""
|
||
初始化 checkpointer（创建数据库表等）
|
||
"""
|
||
if not self._initialized:
|
||
await init_checkpointer_db()
|
||
try:
|
||
from tools.fault_record_db import init_fault_records_table
|
||
await init_fault_records_table()
|
||
except Exception as e:
|
||
print(f"[Checkpointer] 初始化 fault_records 表失败: {e}")
|
||
self._initialized = True
|
||
|
||
def reset(self):
|
||
"""重置 checkpointer（主要用于测试）"""
|
||
self._sync_checkpointer = None
|
||
self._async_checkpointer = None
|
||
self._async_conn_pool = None
|
||
self._initialized = False
|
||
|
||
|
||
checkpointer_manager = CheckpointerManager()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("Checkpointer 配置测试:")
|
||
print(f" 类型: {CHECKPOINTER_TYPE}")
|
||
print(f" PostgreSQL 连接: {'已配置' if POSTGRES_CONNECTION_STRING else '未配置'}")
|
||
|