Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/stores/database.py: 30%
82 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""SQL-backed session store via the DatabaseProviderProtocol."""
3from __future__ import annotations
5from datetime import datetime
6from typing import Any
8from lexigram.contracts.ai.session import (
9 SessionCheckpoint,
10 SessionState,
11 SessionStatus,
12 SessionTurn,
13)
14from lexigram.contracts.data.sql.database import DatabaseProviderProtocol
15from lexigram.logging import (
16 get_logger,
17)
18from lexigram.serialization.backends.json import dumps_str, loads
20logger = get_logger(__name__)
22# SQL DDL
23_CREATE_SESSIONS_TABLE = """
24CREATE TABLE IF NOT EXISTS ai_sessions (
25 session_id TEXT PRIMARY KEY,
26 user_id TEXT NOT NULL,
27 status TEXT NOT NULL,
28 turns TEXT NOT NULL DEFAULT '[]',
29 metadata TEXT NOT NULL DEFAULT '{}',
30 active_tools TEXT NOT NULL DEFAULT '[]',
31 active_skills TEXT NOT NULL DEFAULT '[]',
32 system_prompt TEXT,
33 variables TEXT NOT NULL DEFAULT '{}',
34 created_at TEXT NOT NULL,
35 updated_at TEXT NOT NULL,
36 checkpoint_id TEXT,
37 total_tokens INTEGER NOT NULL DEFAULT 0,
38 total_cost REAL NOT NULL DEFAULT 0.0,
39 turn_count INTEGER NOT NULL DEFAULT 0,
40 parent_session_id TEXT,
41 branch_name TEXT
42);
43"""
45_CREATE_CHECKPOINTS_TABLE = """
46CREATE TABLE IF NOT EXISTS ai_checkpoints (
47 checkpoint_id TEXT PRIMARY KEY,
48 session_id TEXT NOT NULL,
49 state TEXT NOT NULL,
50 created_at TEXT NOT NULL,
51 parent_checkpoint_id TEXT,
52 metadata TEXT NOT NULL DEFAULT '{}'
53);
54"""
57def _row_to_state(row: Any) -> SessionState:
58 """Convert a DB row (sequence or mapping) to a ``SessionState``."""
59 turns_raw = loads(row[3] if isinstance(row, (list, tuple)) else row["turns"])
60 turns = [
61 SessionTurn(
62 turn_id=t["turn_id"],
63 role=t["role"],
64 content=t["content"],
65 timestamp=datetime.fromisoformat(t["timestamp"]),
66 tool_calls=t.get("tool_calls", []),
67 skill_results=t.get("skill_results", []),
68 metadata=t.get("metadata", {}),
69 tokens_used=t.get("tokens_used", 0),
70 cost=t.get("cost", 0.0),
71 model=t.get("model"),
72 provider=t.get("provider"),
73 )
74 for t in turns_raw
75 ]
77 def _col(row: Any, idx: int, name: str) -> Any:
78 if isinstance(row, (list, tuple)):
79 return row[idx]
80 return row[name]
82 return SessionState(
83 session_id=_col(row, 0, "session_id"),
84 user_id=_col(row, 1, "user_id"),
85 status=SessionStatus(_col(row, 2, "status")),
86 turns=turns,
87 metadata=loads(_col(row, 4, "metadata")),
88 active_tools=loads(_col(row, 5, "active_tools")),
89 active_skills=loads(_col(row, 6, "active_skills")),
90 system_prompt=_col(row, 7, "system_prompt"),
91 variables=loads(_col(row, 8, "variables")),
92 created_at=datetime.fromisoformat(_col(row, 9, "created_at")),
93 updated_at=datetime.fromisoformat(_col(row, 10, "updated_at")),
94 checkpoint_id=_col(row, 11, "checkpoint_id"),
95 total_tokens=_col(row, 12, "total_tokens"),
96 total_cost=_col(row, 13, "total_cost"),
97 turn_count=_col(row, 14, "turn_count"),
98 parent_session_id=_col(row, 15, "parent_session_id"),
99 branch_name=_col(row, 16, "branch_name"),
100 )
103class DatabaseSessionStore:
104 """SQL-backed session store for durable production storage.
106 Uses ``DatabaseProviderProtocol`` from ``lexigram-contracts`` so the
107 concrete driver (asyncpg, aiosqlite, …) is never imported directly.
108 All state is JSON-serialised for the SQL TEXT columns.
110 Tables used:
111 - ``ai_sessions`` — one row per session
112 - ``ai_checkpoints`` — one row per checkpoint
114 Args:
115 db: Any object implementing ``DatabaseProviderProtocol``
116 (provides ``scoped_context()`` and ``get_scoped_connection()``).
117 """
119 def __init__(self, db: DatabaseProviderProtocol) -> None:
120 self._db = db
122 async def _ensure_tables(self) -> None:
123 """Create the ai_sessions and ai_checkpoints tables if they do not exist."""
124 await self._db.execute(_CREATE_SESSIONS_TABLE)
125 await self._db.execute(_CREATE_CHECKPOINTS_TABLE)
127 # ------------------------------------------------------------------
128 # Session CRUD
129 # ------------------------------------------------------------------
131 async def save(self, state: SessionState) -> None:
132 """Upsert *state* in the database.
134 Args:
135 state: The session state to persist.
136 """
137 turns_json = dumps_str(
138 [
139 {
140 "turn_id": t.turn_id,
141 "role": t.role,
142 "content": t.content,
143 "timestamp": t.timestamp.isoformat(),
144 "tool_calls": t.tool_calls,
145 "skill_results": t.skill_results,
146 "metadata": t.metadata,
147 "tokens_used": t.tokens_used,
148 "cost": t.cost,
149 "model": t.model,
150 "provider": t.provider,
151 }
152 for t in state.turns
153 ]
154 )
155 sql = """
156 INSERT INTO ai_sessions
157 (session_id, user_id, status, turns, metadata, active_tools,
158 active_skills, system_prompt, variables, created_at, updated_at,
159 checkpoint_id, total_tokens, total_cost, turn_count,
160 parent_session_id, branch_name)
161 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)
162 ON CONFLICT (session_id) DO UPDATE SET
163 status=EXCLUDED.status, turns=EXCLUDED.turns,
164 metadata=EXCLUDED.metadata, active_tools=EXCLUDED.active_tools,
165 active_skills=EXCLUDED.active_skills,
166 system_prompt=EXCLUDED.system_prompt,
167 variables=EXCLUDED.variables,
168 updated_at=EXCLUDED.updated_at,
169 checkpoint_id=EXCLUDED.checkpoint_id,
170 total_tokens=EXCLUDED.total_tokens,
171 total_cost=EXCLUDED.total_cost,
172 turn_count=EXCLUDED.turn_count
173 """
174 async with self._db.scoped_context():
175 conn = await self._db.get_scoped_connection()
176 await conn.execute(
177 sql,
178 state.session_id,
179 state.user_id,
180 state.status.value,
181 turns_json,
182 dumps_str(state.metadata),
183 dumps_str(state.active_tools),
184 dumps_str(state.active_skills),
185 state.system_prompt,
186 dumps_str(state.variables),
187 state.created_at.isoformat(),
188 state.updated_at.isoformat(),
189 state.checkpoint_id,
190 state.total_tokens,
191 state.total_cost,
192 state.turn_count,
193 state.parent_session_id,
194 state.branch_name,
195 )
197 async def load(self, session_id: str) -> SessionState | None:
198 """Return the session state for *session_id*, or ``None``.
200 Args:
201 session_id: The session to load.
203 Returns:
204 Deserialised ``SessionState`` or ``None``.
205 """
206 sql = "SELECT * FROM ai_sessions WHERE session_id = $1"
207 async with self._db.scoped_context():
208 conn = await self._db.get_scoped_connection()
209 row = await conn.fetchrow(sql, session_id)
210 if row is None:
211 return None
212 return _row_to_state(row)
214 async def delete(self, session_id: str) -> None:
215 """Remove a session from the database.
217 Args:
218 session_id: The session to delete.
219 """
220 async with self._db.scoped_context():
221 conn = await self._db.get_scoped_connection()
222 await conn.execute(
223 "DELETE FROM ai_sessions WHERE session_id = $1", session_id
224 )
226 async def list_sessions(self, user_id: str) -> list[SessionState]:
227 """List all sessions belonging to *user_id*.
229 Args:
230 user_id: User to filter by.
232 Returns:
233 All sessions for that user.
234 """
235 sql = "SELECT * FROM ai_sessions WHERE user_id = $1 ORDER BY created_at"
236 async with self._db.scoped_context():
237 conn = await self._db.get_scoped_connection()
238 rows = await conn.fetch(sql, user_id)
239 return [_row_to_state(r) for r in rows]
241 # ------------------------------------------------------------------
242 # Checkpoint CRUD
243 # ------------------------------------------------------------------
245 async def save_checkpoint(self, checkpoint: SessionCheckpoint) -> None:
246 """Persist an immutable checkpoint.
248 Args:
249 checkpoint: The checkpoint to save.
250 """
251 from lexigram.ai.session.stores.cache import _state_to_dict # local helper
253 sql = """
254 INSERT INTO ai_checkpoints
255 (checkpoint_id, session_id, state, created_at,
256 parent_checkpoint_id, metadata)
257 VALUES ($1,$2,$3,$4,$5,$6)
258 ON CONFLICT (checkpoint_id) DO NOTHING
259 """
260 async with self._db.scoped_context():
261 conn = await self._db.get_scoped_connection()
262 await conn.execute(
263 sql,
264 checkpoint.checkpoint_id,
265 checkpoint.session_id,
266 dumps_str(_state_to_dict(checkpoint.state)),
267 checkpoint.created_at.isoformat(),
268 checkpoint.parent_checkpoint_id,
269 dumps_str(checkpoint.metadata),
270 )
272 async def load_checkpoint(self, checkpoint_id: str) -> SessionCheckpoint | None:
273 """Return the checkpoint for *checkpoint_id*, or ``None``.
275 Args:
276 checkpoint_id: The checkpoint to load.
278 Returns:
279 Deserialised ``SessionCheckpoint`` or ``None``.
280 """
281 sql = "SELECT * FROM ai_checkpoints WHERE checkpoint_id = $1"
282 async with self._db.scoped_context():
283 conn = await self._db.get_scoped_connection()
284 row = await conn.fetchrow(sql, checkpoint_id)
285 if row is None:
286 return None
287 return self._row_to_checkpoint(row)
289 async def list_checkpoints(self, session_id: str) -> list[SessionCheckpoint]:
290 """List all checkpoints for *session_id*, oldest-first.
292 Args:
293 session_id: The session to query.
295 Returns:
296 Checkpoints in chronological order.
297 """
298 sql = "SELECT * FROM ai_checkpoints WHERE session_id = $1 ORDER BY created_at"
299 async with self._db.scoped_context():
300 conn = await self._db.get_scoped_connection()
301 rows = await conn.fetch(sql, session_id)
302 return [self._row_to_checkpoint(r) for r in rows]
304 async def delete_checkpoint(self, checkpoint_id: str) -> None:
305 """Remove a checkpoint from the database.
307 Args:
308 checkpoint_id: The checkpoint to delete.
309 """
310 async with self._db.scoped_context():
311 conn = await self._db.get_scoped_connection()
312 await conn.execute(
313 "DELETE FROM ai_checkpoints WHERE checkpoint_id = $1", checkpoint_id
314 )
316 # ------------------------------------------------------------------
317 # Internal helpers
318 # ------------------------------------------------------------------
320 @staticmethod
321 def _row_to_checkpoint(row: Any) -> SessionCheckpoint:
322 """Deserialise a DB row to a ``SessionCheckpoint``."""
323 from lexigram.ai.session.stores.cache import _dict_to_state # local helper
325 def _col(row: Any, idx: int, name: str) -> Any:
326 if isinstance(row, (list, tuple)):
327 return row[idx]
328 return row[name]
330 state_dict = loads(_col(row, 2, "state"))
331 return SessionCheckpoint(
332 checkpoint_id=_col(row, 0, "checkpoint_id"),
333 session_id=_col(row, 1, "session_id"),
334 state=_dict_to_state(state_dict),
335 created_at=datetime.fromisoformat(_col(row, 3, "created_at")),
336 parent_checkpoint_id=_col(row, 4, "parent_checkpoint_id"),
337 metadata=loads(_col(row, 5, "metadata")),
338 )
341__all__ = ["DatabaseSessionStore"]