Coverage for agentos/storage/base.py: 54%

1""" 

2AgentOS v0.20 持久化存储层。 

3Base + SQLite实现，支持Checkpoint持久化。 

4""" 

5 

6from __future__ import annotations 

7 

8import json 

9import sqlite3 

10import time 

11from abc import ABC, abstractmethod 

12from dataclasses import dataclass 

13 

14 

15# ── 抽象基类 ──────────────────────────────────── 

16 

17class CheckpointStore(ABC): 

18 

19 """检查点存储基类。""" 

20 

21 @abstractmethod 

22 async def save(self, session_id: str, snapshot: dict): ... 

23 @abstractmethod 

24 async def load(self, session_id: str) -> dict | None: ... 

25 @abstractmethod 

26 async def delete(self, session_id: str): ... 

27 @abstractmethod 

28 async def list_sessions(self, limit: int = 50) -> list[str]: ... 

29 

30 

31@dataclass 

32class SqliteStore(CheckpointStore): 

33 """SQLite 持久化存储。""" 

34 

35 path: str = ":memory:" 

36 

37 def __post_init__(self): 

38 self._conn = sqlite3.connect(self.path, check_same_thread=False) 

39 self._conn.execute( 

40 """CREATE TABLE IF NOT EXISTS checkpoints ( 

41 session_id TEXT PRIMARY KEY, 

42 snapshot TEXT NOT NULL, 

43 created_at REAL NOT NULL, 

44 updated_at REAL NOT NULL 

45 )""" 

46 ) 

47 self._conn.execute("CREATE INDEX IF NOT EXISTS idx_updated ON checkpoints(updated_at DESC)") 

48 self._conn.commit() 

49 

50 async def save(self, session_id: str, snapshot: dict): 

51 now = time.time() 

52 self._conn.execute( 

53 """INSERT INTO checkpoints(session_id, snapshot, created_at, updated_at) 

54 VALUES(?, ?, ?, ?) 

55 ON CONFLICT(session_id) DO UPDATE SET 

56 snapshot=excluded.snapshot, updated_at=excluded.updated_at""", 

57 (session_id, json.dumps(snapshot, default=str), now, now), 

58 ) 

59 self._conn.commit() 

60 

61 async def load(self, session_id: str) -> dict | None: 

62 row = self._conn.execute( 

63 "SELECT snapshot FROM checkpoints WHERE session_id=?", (session_id,) 

64 ).fetchone() 

65 return json.loads(row[0]) if row else None 

66 

67 async def delete(self, session_id: str): 

68 self._conn.execute("DELETE FROM checkpoints WHERE session_id=?", (session_id,)) 

69 self._conn.commit() 

70 

71 async def list_sessions(self, limit: int = 50) -> list[str]: 

72 rows = self._conn.execute( 

73 "SELECT session_id FROM checkpoints ORDER BY updated_at DESC LIMIT ?", (limit,) 

74 ).fetchall() 

75 return [r[0] for r in rows]