Coverage for agentos/checkpoint/sqlite.py: 36%

1""" 

2SQLite Checkpointer — 零依赖本地持久化。 

3 

4适用场景: 单机部署、开发调试、POC。 

5生产多机部署请使用 PostgresCheckpointer。 

6""" 

7 

8from __future__ import annotations 

9 

10import json 

11import os 

12import sqlite3 

13 

14from agentos.checkpoint.base import ( 

15 Checkpoint, 

16 CheckpointBackend, 

17 CheckpointMetadata, 

18) 

19 

20__all__ = ["SQLiteCheckpointer"] 

21 

22_SCHEMA = """ 

23CREATE TABLE IF NOT EXISTS checkpoints ( 

24 id INTEGER PRIMARY KEY AUTOINCREMENT, 

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 TEXT NOT NULL, 

30 tags TEXT NOT NULL DEFAULT '[]', 

31 summary TEXT NOT NULL DEFAULT '', 

32 messages_blob TEXT NOT NULL DEFAULT '[]', 

33 state_blob TEXT NOT NULL DEFAULT '{}', 

34 tools_blob TEXT NOT NULL DEFAULT '{}', 

35 next_node TEXT NOT NULL DEFAULT '' 

36); 

37 

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); 

41""" 

42 

43 

44class SQLiteCheckpointer(CheckpointBackend): 

45 """SQLite 后端 Checkpointer。 

46 

47 用法: 

48 cp = SQLiteCheckpointer(db_path="data/checkpoints.db") 

49 await cp.put(checkpoint) 

50 latest = await cp.get_latest("thread_abc") 

51 """ 

52 

53 def __init__(self, db_path: str = "checkpoints.db"): 

54 self._db_path = db_path 

55 os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) 

56 self._init_db() 

57 

58 def _init_db(self) -> None: 

59 with sqlite3.connect(self._db_path) as conn: 

60 conn.executescript(_SCHEMA) 

61 conn.commit() 

62 

63 def _get_conn(self) -> sqlite3.Connection: 

64 conn = sqlite3.connect(self._db_path) 

65 conn.row_factory = sqlite3.Row 

66 return conn 

67 

68 def _row_to_metadata(self, row: sqlite3.Row) -> CheckpointMetadata: 

69 return CheckpointMetadata( 

70 thread_id=row["thread_id"], 

71 checkpoint_id=row["checkpoint_id"], 

72 parent_checkpoint_id=row["parent_id"], 

73 step=row["step"], 

74 created_at=row["created_at"], 

75 tags=json.loads(row["tags"]), 

76 summary=row["summary"], 

77 ) 

78 

79 def _row_to_checkpoint(self, row: sqlite3.Row) -> Checkpoint: 

80 return Checkpoint( 

81 metadata=self._row_to_metadata(row), 

82 messages=json.loads(row["messages_blob"]), 

83 state=json.loads(row["state_blob"]), 

84 tools_result=json.loads(row["tools_blob"]), 

85 next_node=row["next_node"], 

86 ) 

87 

88 async def put(self, checkpoint: Checkpoint) -> str: 

89 meta = checkpoint.metadata 

90 with self._get_conn() as conn: 

91 conn.execute( 

92 """INSERT OR REPLACE INTO checkpoints 

93 (thread_id, checkpoint_id, parent_id, step, created_at, tags, summary, 

94 messages_blob, state_blob, tools_blob, next_node) 

95 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", 

96 ( 

97 meta.thread_id, 

98 meta.checkpoint_id, 

99 meta.parent_checkpoint_id, 

100 meta.step, 

101 meta.created_at, 

102 json.dumps(meta.tags), 

103 meta.summary, 

104 json.dumps(checkpoint.messages, ensure_ascii=False), 

105 json.dumps(checkpoint.state, ensure_ascii=False), 

106 json.dumps(checkpoint.tools_result, ensure_ascii=False), 

107 checkpoint.next_node, 

108 ), 

109 ) 

110 conn.commit() 

111 return meta.checkpoint_id 

112 

113 async def get(self, checkpoint_id: str) -> Checkpoint | None: 

114 with self._get_conn() as conn: 

115 row = conn.execute( 

116 "SELECT * FROM checkpoints WHERE checkpoint_id = ?", (checkpoint_id,) 

117 ).fetchone() 

118 return self._row_to_checkpoint(row) if row else None 

119 

120 async def get_latest(self, thread_id: str) -> Checkpoint | None: 

121 with self._get_conn() as conn: 

122 row = conn.execute( 

123 "SELECT * FROM checkpoints WHERE thread_id = ? ORDER BY step DESC LIMIT 1", 

124 (thread_id,), 

125 ).fetchone() 

126 return self._row_to_checkpoint(row) if row else None 

127 

128 async def list_threads( 

129 self, limit: int = 50, offset: int = 0 

130 ) -> list[CheckpointMetadata]: 

131 with self._get_conn() as conn: 

132 rows = conn.execute( 

133 """SELECT * FROM checkpoints 

134 WHERE checkpoint_id IN ( 

135 SELECT checkpoint_id FROM checkpoints 

136 GROUP BY thread_id HAVING step = MAX(step) 

137 ) 

138 ORDER BY created_at DESC LIMIT ? OFFSET ?""", 

139 (limit, offset), 

140 ).fetchall() 

141 return [self._row_to_metadata(r) for r in rows] 

142 

143 async def list_checkpoints( 

144 self, thread_id: str, limit: int = 100, offset: int = 0 

145 ) -> list[CheckpointMetadata]: 

146 with self._get_conn() as conn: 

147 rows = conn.execute( 

148 "SELECT * FROM checkpoints WHERE thread_id = ? ORDER BY step DESC LIMIT ? OFFSET ?", 

149 (thread_id, limit, offset), 

150 ).fetchall() 

151 return [self._row_to_metadata(r) for r in rows] 

152 

153 async def delete_thread(self, thread_id: str) -> int: 

154 with self._get_conn() as conn: 

155 cur = conn.execute("DELETE FROM checkpoints WHERE thread_id = ?", (thread_id,)) 

156 conn.commit() 

157 return cur.rowcount 

158 

159 async def delete_before(self, thread_id: str, before_step: int) -> int: 

160 with self._get_conn() as conn: 

161 cur = conn.execute( 

162 "DELETE FROM checkpoints WHERE thread_id = ? AND step < ?", 

163 (thread_id, before_step), 

164 ) 

165 conn.commit() 

166 return cur.rowcount