Coverage for agentos/checkpoint/postgres.py: 31%
68 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 21:19 +0800
1"""
2Postgres Checkpointer — 生产级持久化后端。
4需安装: pip install asyncpg
6参考 LangGraph PostgresSaver 的 schema 设计。
7"""
9from __future__ import annotations
11import json
12from typing import Any
14from agentos.checkpoint.base import (
15 Checkpoint,
16 CheckpointBackend,
17 CheckpointMetadata,
18)
20__all__ = ["PostgresCheckpointer"]
22_SCHEMA = """
23CREATE TABLE IF NOT EXISTS checkpoints (
24 id BIGSERIAL PRIMARY KEY,
25 thread_id TEXT NOT NULL,
26 checkpoint_id TEXT NOT NULL UNIQUE,
27 parent_id TEXT,
28 step INTEGER NOT NULL,
29 created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
30 tags JSONB NOT NULL DEFAULT '[]',
31 summary TEXT NOT NULL DEFAULT '',
32 messages_blob JSONB NOT NULL DEFAULT '[]',
33 state_blob JSONB NOT NULL DEFAULT '{}',
34 tools_blob JSONB NOT NULL DEFAULT '{}',
35 next_node TEXT NOT NULL DEFAULT ''
36);
38CREATE INDEX IF NOT EXISTS idx_thread_step ON checkpoints(thread_id, step DESC);
39CREATE INDEX IF NOT EXISTS idx_checkpoint_id ON checkpoints(checkpoint_id);
40CREATE INDEX IF NOT EXISTS idx_parent ON checkpoints(parent_id);
41CREATE INDEX IF NOT EXISTS idx_created_at ON checkpoints(created_at DESC);
42"""
45class PostgresCheckpointer(CheckpointBackend):
46 """Postgres 后端 Checkpointer — 生产环境推荐。
48 用法:
49 cp = PostgresCheckpointer(dsn="postgresql://user:pass@localhost:5432/agentos")
50 await cp.put(checkpoint)
51 latest = await cp.get_latest("thread_abc")
52 """
54 def __init__(self, dsn: str = "", **kwargs: Any):
55 self._dsn = dsn or "postgresql://localhost:5432/agentos"
56 self._kwargs = kwargs
57 self._pool: Any = None
58 self._initialized = False
60 async def _ensure_pool(self):
61 if self._pool is not None:
62 return
63 import asyncpg
65 self._pool = await asyncpg.create_pool(dsn=self._dsn, **self._kwargs)
66 async with self._pool.acquire() as conn:
67 await conn.execute(_SCHEMA)
69 async def put(self, checkpoint: Checkpoint) -> str:
70 await self._ensure_pool()
71 meta = checkpoint.metadata
72 async with self._pool.acquire() as conn:
73 await conn.execute(
74 """INSERT INTO checkpoints
75 (thread_id, checkpoint_id, parent_id, step, created_at, tags, summary,
76 messages_blob, state_blob, tools_blob, next_node)
77 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
78 ON CONFLICT (checkpoint_id) DO UPDATE SET
79 step=$4, messages_blob=$8, state_blob=$9, tools_blob=$10, next_node=$11""",
80 meta.thread_id,
81 meta.checkpoint_id,
82 meta.parent_checkpoint_id,
83 meta.step,
84 meta.created_at,
85 json.dumps(meta.tags),
86 meta.summary,
87 json.dumps(checkpoint.messages, ensure_ascii=False),
88 json.dumps(checkpoint.state, ensure_ascii=False),
89 json.dumps(checkpoint.tools_result, ensure_ascii=False),
90 checkpoint.next_node,
91 )
92 return meta.checkpoint_id
94 async def get(self, checkpoint_id: str) -> Checkpoint | None:
95 await self._ensure_pool()
96 async with self._pool.acquire() as conn:
97 row = await conn.fetchrow(
98 "SELECT * FROM checkpoints WHERE checkpoint_id = $1", checkpoint_id
99 )
100 return self._row_to_checkpoint(row) if row else None
102 async def get_latest(self, thread_id: str) -> Checkpoint | None:
103 await self._ensure_pool()
104 async with self._pool.acquire() as conn:
105 row = await conn.fetchrow(
106 "SELECT * FROM checkpoints WHERE thread_id = $1 ORDER BY step DESC LIMIT 1",
107 thread_id,
108 )
109 return self._row_to_checkpoint(row) if row else None
111 async def list_threads(self, limit: int = 50, offset: int = 0) -> list[CheckpointMetadata]:
112 await self._ensure_pool()
113 async with self._pool.acquire() as conn:
114 rows = await conn.fetch(
115 """SELECT DISTINCT ON (thread_id) *
116 FROM checkpoints
117 ORDER BY thread_id, step DESC
118 LIMIT $1 OFFSET $2""",
119 limit,
120 offset,
121 )
122 return [self._row_to_metadata(r) for r in rows]
124 async def list_checkpoints(
125 self, thread_id: str, limit: int = 100, offset: int = 0
126 ) -> list[CheckpointMetadata]:
127 await self._ensure_pool()
128 async with self._pool.acquire() as conn:
129 rows = await conn.fetch(
130 "SELECT * FROM checkpoints WHERE thread_id = $1 ORDER BY step DESC LIMIT $2 OFFSET $3",
131 thread_id,
132 limit,
133 offset,
134 )
135 return [self._row_to_metadata(r) for r in rows]
137 async def delete_thread(self, thread_id: str) -> int:
138 await self._ensure_pool()
139 async with self._pool.acquire() as conn:
140 result = await conn.execute("DELETE FROM checkpoints WHERE thread_id = $1", thread_id)
141 return int(result.split()[-1]) if result else 0
143 async def delete_before(self, thread_id: str, before_step: int) -> int:
144 await self._ensure_pool()
145 async with self._pool.acquire() as conn:
146 result = await conn.execute(
147 "DELETE FROM checkpoints WHERE thread_id = $1 AND step < $2",
148 thread_id,
149 before_step,
150 )
151 return int(result.split()[-1]) if result else 0
153 async def close(self) -> None:
154 if self._pool:
155 await self._pool.close()
156 self._pool = None
158 @staticmethod
159 def _row_to_metadata(row: Any) -> CheckpointMetadata:
160 return CheckpointMetadata(
161 thread_id=row["thread_id"],
162 checkpoint_id=row["checkpoint_id"],
163 parent_checkpoint_id=row["parent_id"],
164 step=row["step"],
165 created_at=str(row["created_at"]),
166 tags=row["tags"] if isinstance(row["tags"], list) else json.loads(row["tags"]),
167 summary=row["summary"],
168 )
170 @staticmethod
171 def _row_to_checkpoint(row: Any) -> Checkpoint:
172 messages = (
173 row["messages_blob"]
174 if isinstance(row["messages_blob"], list)
175 else json.loads(row["messages_blob"])
176 )
177 state = (
178 row["state_blob"]
179 if isinstance(row["state_blob"], dict)
180 else json.loads(row["state_blob"])
181 )
182 tools = (
183 row["tools_blob"]
184 if isinstance(row["tools_blob"], dict)
185 else json.loads(row["tools_blob"])
186 )
187 return Checkpoint(
188 metadata=PostgresCheckpointer._row_to_metadata(row),
189 messages=messages,
190 state=state,
191 tools_result=tools,
192 next_node=row["next_node"],
193 )