File size: 3,721 Bytes
20e9e63 | 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 113 114 115 116 117 118 119 | """
session_store.py
----------------
Persistence for `DialogueState`, keyed by session_id.
Two backends:
- SQLite (default for local / Render / Docker with persistent volume)
- In-memory dict (for Hugging Face Spaces where the filesystem is ephemeral)
The backend is selected via the `PAYPARSE_SESSION_BACKEND` env var:
"sqlite" (default) → SQLite file
"memory" → in-memory dict (lost on restart, fine for a demo)
"""
from __future__ import annotations
import os
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
from schema import DialogueState
DEFAULT_DB_PATH = Path(__file__).parent / "sessions.db"
_SCHEMA = """
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
state_json TEXT NOT NULL,
updated_at TEXT NOT NULL
);
"""
class SessionStore:
"""Thread-safe session store. Delegates to SQLite or in-memory dict."""
def __init__(self, db_path: Path | str = DEFAULT_DB_PATH):
backend = os.getenv("PAYPARSE_SESSION_BACKEND", "sqlite").lower()
self._lock = threading.Lock()
if backend == "memory":
self._backend = _MemoryBackend()
else:
self._backend = _SQLiteBackend(db_path)
def get(self, session_id: str) -> Optional[DialogueState]:
return self._backend.get(session_id)
def get_or_create(self, session_id: str) -> DialogueState:
return self.get(session_id) or DialogueState(session_id=session_id)
def save(self, state: DialogueState) -> None:
self._backend.save(state)
def delete(self, session_id: str) -> bool:
return self._backend.delete(session_id)
def close(self) -> None:
self._backend.close()
class _SQLiteBackend:
"""SQLite-backed persistence (local / Render / Docker with volume)."""
def __init__(self, db_path: Path | str):
self._conn = sqlite3.connect(str(db_path), check_same_thread=False)
self._conn.execute(_SCHEMA)
self._conn.commit()
def get(self, session_id: str) -> Optional[DialogueState]:
row = self._conn.execute(
"SELECT state_json FROM sessions WHERE session_id = ?",
(session_id,),
).fetchone()
return DialogueState.model_validate_json(row[0]) if row else None
def save(self, state: DialogueState) -> None:
now = datetime.now(timezone.utc).isoformat()
self._conn.execute(
"INSERT INTO sessions (session_id, state_json, updated_at) "
"VALUES (?, ?, ?) "
"ON CONFLICT(session_id) DO UPDATE SET "
"state_json = excluded.state_json, updated_at = excluded.updated_at",
(state.session_id, state.model_dump_json(), now),
)
self._conn.commit()
def delete(self, session_id: str) -> bool:
cursor = self._conn.execute(
"DELETE FROM sessions WHERE session_id = ?", (session_id,)
)
self._conn.commit()
return cursor.rowcount > 0
def close(self) -> None:
self._conn.close()
class _MemoryBackend:
"""In-memory dict store (Hugging Face Spaces — ephemeral filesystem)."""
def __init__(self):
self._data: dict[str, str] = {}
def get(self, session_id: str) -> Optional[DialogueState]:
json_str = self._data.get(session_id)
return DialogueState.model_validate_json(json_str) if json_str else None
def save(self, state: DialogueState) -> None:
self._data[state.session_id] = state.model_dump_json()
def delete(self, session_id: str) -> bool:
return self._data.pop(session_id, None) is not None
def close(self) -> None:
pass
|