Spaces:
Sleeping
Sleeping
File size: 3,591 Bytes
c27eaf1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """SQLite-backed async session storage."""
from __future__ import annotations
import json
import os
from datetime import datetime, timezone
import aiosqlite
from models import Message, Phase, SessionState
DB_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data")
DB_PATH = os.path.join(DB_DIR, "sessions.db")
# ---------------------------------------------------------------------------
# DB initialisation
# ---------------------------------------------------------------------------
async def init_db() -> None:
"""Create the sessions table if it doesn't exist."""
os.makedirs(DB_DIR, exist_ok=True)
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
state_json TEXT NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
await db.commit()
# ---------------------------------------------------------------------------
# CRUD helpers
# ---------------------------------------------------------------------------
def _now() -> str:
return datetime.now(timezone.utc).isoformat()
async def load_session(session_id: str) -> SessionState | None:
"""Load a session from the database. Returns None if not found."""
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute(
"SELECT state_json FROM sessions WHERE session_id = ?",
(session_id,),
)
row = await cursor.fetchone()
if row is None:
return None
data = json.loads(row[0])
return SessionState(**data)
async def save_session(session_id: str, state: SessionState) -> None:
"""Insert or update a session."""
now = _now()
state.updated_at = now
if not state.created_at:
state.created_at = now
state_json = state.model_dump_json()
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"""
INSERT INTO sessions (session_id, state_json, created_at, updated_at)
VALUES (?, ?, ?, ?)
ON CONFLICT(session_id)
DO UPDATE SET state_json = excluded.state_json,
updated_at = excluded.updated_at
""",
(session_id, state_json, state.created_at, now),
)
await db.commit()
async def list_sessions() -> list[dict]:
"""Return lightweight info about all sessions."""
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute(
"SELECT session_id, state_json, created_at, updated_at FROM sessions ORDER BY updated_at DESC"
)
rows = await cursor.fetchall()
results = []
for sid, sj, ca, ua in rows:
data = json.loads(sj)
results.append(
{
"session_id": sid,
"phase": data.get("phase", "discovery"),
"confidence": data.get("confidence", 0.0),
"message_count": len(data.get("messages", [])),
"created_at": ca,
"updated_at": ua,
}
)
return results
async def delete_session(session_id: str) -> bool:
"""Delete a session. Returns True if it existed."""
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute(
"DELETE FROM sessions WHERE session_id = ?", (session_id,)
)
await db.commit()
return cursor.rowcount > 0
|