Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/stores/cache.py: 31%

83 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Redis-backed session store via the CacheBackendProtocol protocol.""" 

2 

3from __future__ import annotations 

4 

5from datetime import datetime 

6from typing import Any 

7 

8from lexigram.contracts.ai.session import ( 

9 SessionCheckpoint, 

10 SessionState, 

11 SessionStatus, 

12 SessionTurn, 

13) 

14from lexigram.contracts.infra.cache import CacheBackendProtocol 

15from lexigram.logging import ( 

16 get_logger, 

17) 

18from lexigram.serialization.backends.json import dumps_str, loads 

19 

20logger = get_logger(__name__) 

21 

22_SESSION_PREFIX = "ai:session:" 

23_CHECKPOINT_PREFIX = "ai:checkpoint:" 

24_USER_INDEX_PREFIX = "ai:session:user:" 

25 

26 

27def _state_to_dict(state: SessionState) -> dict[str, Any]: 

28 """Serialise a ``SessionState`` to a JSON-compatible dict.""" 

29 return { 

30 "session_id": state.session_id, 

31 "user_id": state.user_id, 

32 "status": state.status.value, 

33 "turns": [ 

34 { 

35 "turn_id": t.turn_id, 

36 "role": t.role, 

37 "content": t.content, 

38 "timestamp": t.timestamp.isoformat(), 

39 "tool_calls": t.tool_calls, 

40 "skill_results": t.skill_results, 

41 "metadata": t.metadata, 

42 "tokens_used": t.tokens_used, 

43 "cost": t.cost, 

44 "model": t.model, 

45 "provider": t.provider, 

46 } 

47 for t in state.turns 

48 ], 

49 "metadata": state.metadata, 

50 "active_tools": state.active_tools, 

51 "active_skills": state.active_skills, 

52 "system_prompt": state.system_prompt, 

53 "variables": state.variables, 

54 "created_at": state.created_at.isoformat(), 

55 "updated_at": state.updated_at.isoformat(), 

56 "checkpoint_id": state.checkpoint_id, 

57 "total_tokens": state.total_tokens, 

58 "total_cost": state.total_cost, 

59 "turn_count": state.turn_count, 

60 "parent_session_id": state.parent_session_id, 

61 "branch_name": state.branch_name, 

62 } 

63 

64 

65def _dict_to_state(d: dict[str, Any]) -> SessionState: 

66 """Deserialise a dict back to a ``SessionState``.""" 

67 turns = [ 

68 SessionTurn( 

69 turn_id=t["turn_id"], 

70 role=t["role"], 

71 content=t["content"], 

72 timestamp=datetime.fromisoformat(t["timestamp"]), 

73 tool_calls=t.get("tool_calls", []), 

74 skill_results=t.get("skill_results", []), 

75 metadata=t.get("metadata", {}), 

76 tokens_used=t.get("tokens_used", 0), 

77 cost=t.get("cost", 0.0), 

78 model=t.get("model"), 

79 provider=t.get("provider"), 

80 ) 

81 for t in d.get("turns", []) 

82 ] 

83 return SessionState( 

84 session_id=d["session_id"], 

85 user_id=d["user_id"], 

86 status=SessionStatus(d["status"]), 

87 turns=turns, 

88 metadata=d.get("metadata", {}), 

89 active_tools=d.get("active_tools", []), 

90 active_skills=d.get("active_skills", []), 

91 system_prompt=d.get("system_prompt"), 

92 variables=d.get("variables", {}), 

93 created_at=datetime.fromisoformat(d["created_at"]), 

94 updated_at=datetime.fromisoformat(d["updated_at"]), 

95 checkpoint_id=d.get("checkpoint_id"), 

96 total_tokens=d.get("total_tokens", 0), 

97 total_cost=d.get("total_cost", 0.0), 

98 turn_count=d.get("turn_count", 0), 

99 parent_session_id=d.get("parent_session_id"), 

100 branch_name=d.get("branch_name"), 

101 ) 

102 

103 

104def _checkpoint_to_dict(cp: SessionCheckpoint) -> dict[str, Any]: 

105 """Serialise a ``SessionCheckpoint`` to a JSON-compatible dict.""" 

106 return { 

107 "checkpoint_id": cp.checkpoint_id, 

108 "session_id": cp.session_id, 

109 "state": _state_to_dict(cp.state), 

110 "created_at": cp.created_at.isoformat(), 

111 "parent_checkpoint_id": cp.parent_checkpoint_id, 

112 "metadata": cp.metadata, 

113 } 

114 

115 

116def _dict_to_checkpoint(d: dict[str, Any]) -> SessionCheckpoint: 

117 """Deserialise a dict back to a ``SessionCheckpoint``.""" 

118 return SessionCheckpoint( 

119 checkpoint_id=d["checkpoint_id"], 

120 session_id=d["session_id"], 

121 state=_dict_to_state(d["state"]), 

122 created_at=datetime.fromisoformat(d["created_at"]), 

123 parent_checkpoint_id=d.get("parent_checkpoint_id"), 

124 metadata=d.get("metadata", {}), 

125 ) 

126 

127 

128class CacheSessionStore: 

129 """Redis-backed session store for multi-process deployments. 

130 

131 Uses the ``CacheBackendProtocol`` protocol from ``lexigram-contracts`` so the 

132 concrete Redis client is never imported directly. Sessions and 

133 checkpoints are JSON-serialised. A per-user set index tracks which 

134 sessions belong to a given user. 

135 

136 Args: 

137 cache: Any object implementing ``CacheBackendProtocol`` (get/set/delete). 

138 ttl: Session TTL in seconds (default 24 h). 

139 """ 

140 

141 def __init__(self, cache: CacheBackendProtocol, ttl: int = 86400) -> None: 

142 self._cache = cache 

143 self._ttl = ttl 

144 

145 # ------------------------------------------------------------------ 

146 # Session CRUD 

147 # ------------------------------------------------------------------ 

148 

149 async def save(self, state: SessionState) -> None: 

150 """Persist *state* in the cache. 

151 

152 Args: 

153 state: Session state to save. 

154 """ 

155 key = _SESSION_PREFIX + state.session_id 

156 payload = dumps_str(_state_to_dict(state)) 

157 await self._cache.set(key, payload, ttl=self._ttl) 

158 # Maintain user → session_id index 

159 index_key = _USER_INDEX_PREFIX + state.user_id 

160 existing_raw = await self._cache.get(index_key) 

161 ids: list[str] = loads(existing_raw) if existing_raw else [] # type: ignore[arg-type] 

162 if state.session_id not in ids: 

163 ids.append(state.session_id) 

164 await self._cache.set(index_key, dumps_str(ids), ttl=self._ttl) 

165 

166 async def load(self, session_id: str) -> SessionState | None: 

167 """Return the session state for *session_id*, or ``None``. 

168 

169 Args: 

170 session_id: The session to load. 

171 

172 Returns: 

173 Deserialised ``SessionState`` or ``None``. 

174 """ 

175 raw = await self._cache.get(_SESSION_PREFIX + session_id) 

176 if raw is None: 

177 return None 

178 return _dict_to_state(loads(raw)) # type: ignore[arg-type] 

179 

180 async def delete(self, session_id: str) -> None: 

181 """Remove a session from the cache. 

182 

183 Args: 

184 session_id: The session to delete. 

185 """ 

186 await self._cache.delete(_SESSION_PREFIX + session_id) 

187 

188 async def list_sessions(self, user_id: str) -> list[SessionState]: 

189 """List all sessions for *user_id*. 

190 

191 Args: 

192 user_id: The user to query. 

193 

194 Returns: 

195 All sessions found for that user. 

196 """ 

197 index_key = _USER_INDEX_PREFIX + user_id 

198 raw = await self._cache.get(index_key) 

199 if not raw: 

200 return [] 

201 ids = loads(raw) # type: ignore[arg-type] 

202 results: list[SessionState] = [] 

203 for sid in ids: 

204 state = await self.load(sid) 

205 if state is not None: 

206 results.append(state) 

207 return results 

208 

209 # ------------------------------------------------------------------ 

210 # Checkpoint CRUD 

211 # ------------------------------------------------------------------ 

212 

213 async def save_checkpoint(self, checkpoint: SessionCheckpoint) -> None: 

214 """Persist a checkpoint in the cache. 

215 

216 Args: 

217 checkpoint: The checkpoint to save. 

218 """ 

219 key = _CHECKPOINT_PREFIX + checkpoint.checkpoint_id 

220 payload = dumps_str(_checkpoint_to_dict(checkpoint)) 

221 await self._cache.set(key, payload, ttl=self._ttl) 

222 # Maintain session → checkpoint_id index 

223 idx_key = _CHECKPOINT_PREFIX + "idx:" + checkpoint.session_id 

224 raw = await self._cache.get(idx_key) 

225 ids: list[str] = loads(raw) if raw else [] # type: ignore[arg-type] 

226 if checkpoint.checkpoint_id not in ids: 

227 ids.append(checkpoint.checkpoint_id) 

228 await self._cache.set(idx_key, dumps_str(ids), ttl=self._ttl) 

229 

230 async def load_checkpoint(self, checkpoint_id: str) -> SessionCheckpoint | None: 

231 """Return the checkpoint for *checkpoint_id*, or ``None``. 

232 

233 Args: 

234 checkpoint_id: The checkpoint to load. 

235 

236 Returns: 

237 Deserialised ``SessionCheckpoint`` or ``None``. 

238 """ 

239 raw = await self._cache.get(_CHECKPOINT_PREFIX + checkpoint_id) 

240 if raw is None: 

241 return None 

242 return _dict_to_checkpoint(loads(raw)) # type: ignore[arg-type] 

243 

244 async def list_checkpoints(self, session_id: str) -> list[SessionCheckpoint]: 

245 """List all checkpoints for *session_id*, oldest-first. 

246 

247 Args: 

248 session_id: The session to query. 

249 

250 Returns: 

251 Checkpoints in chronological order. 

252 """ 

253 idx_key = _CHECKPOINT_PREFIX + "idx:" + session_id 

254 raw = await self._cache.get(idx_key) 

255 if not raw: 

256 return [] 

257 ids: list[str] = loads(raw) # type: ignore[arg-type] 

258 checkpoints: list[SessionCheckpoint] = [] 

259 for cid in ids: 

260 cp = await self.load_checkpoint(cid) 

261 if cp is not None: 

262 checkpoints.append(cp) 

263 return sorted(checkpoints, key=lambda c: c.created_at) 

264 

265 async def delete_checkpoint(self, checkpoint_id: str) -> None: 

266 """Remove a checkpoint from the cache. 

267 

268 Args: 

269 checkpoint_id: The checkpoint to delete. 

270 """ 

271 await self._cache.delete(_CHECKPOINT_PREFIX + checkpoint_id) 

272 

273 

274__all__ = ["CacheSessionStore"]