223 lines
7.9 KiB
Python
223 lines
7.9 KiB
Python
"""SQLite-backed external-conversation to Agent-session mapping."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Protocol, Union
|
|
|
|
|
|
class SessionStoreError(RuntimeError):
|
|
"""Raised when persisted session state is invalid or unavailable."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SessionRecord:
|
|
conversation_id: str
|
|
session_id: str
|
|
external_subject_id: Optional[str]
|
|
metadata: Dict[str, Any]
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"conversation_id": self.conversation_id,
|
|
"session_id": self.session_id,
|
|
"external_subject_id": self.external_subject_id,
|
|
"metadata": self.metadata,
|
|
"created_at": self.created_at,
|
|
"updated_at": self.updated_at,
|
|
}
|
|
|
|
|
|
class SessionStore(Protocol):
|
|
"""Storage contract that can also be implemented by a shared database adapter."""
|
|
|
|
def get(self, conversation_id: str) -> Optional[SessionRecord]:
|
|
...
|
|
|
|
def put(
|
|
self,
|
|
*,
|
|
conversation_id: str,
|
|
session_id: str,
|
|
external_subject_id: Optional[str],
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> SessionRecord:
|
|
...
|
|
|
|
def delete(self, conversation_id: str) -> bool:
|
|
...
|
|
|
|
|
|
class SQLiteSessionStore:
|
|
"""Small thread-safe SQLite store suitable for a single application instance."""
|
|
|
|
def __init__(self, path: Union[str, Path]) -> None:
|
|
self.path = Path(path)
|
|
self._lock = threading.RLock()
|
|
try:
|
|
if str(self.path) != ":memory:":
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._connection = sqlite3.connect(
|
|
str(self.path),
|
|
timeout=5.0,
|
|
check_same_thread=False,
|
|
)
|
|
self._connection.row_factory = sqlite3.Row
|
|
self._connection.execute("PRAGMA busy_timeout = 5000")
|
|
self._initialize()
|
|
except SessionStoreError:
|
|
connection = getattr(self, "_connection", None)
|
|
if connection is not None:
|
|
connection.close()
|
|
raise
|
|
except (OSError, sqlite3.Error) as exc:
|
|
connection = getattr(self, "_connection", None)
|
|
if connection is not None:
|
|
connection.close()
|
|
raise SessionStoreError(f"cannot initialize session database: {exc}") from exc
|
|
|
|
def __enter__(self) -> "SQLiteSessionStore":
|
|
return self
|
|
|
|
def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
try:
|
|
with self._lock:
|
|
self._connection.close()
|
|
except sqlite3.Error as exc:
|
|
raise SessionStoreError(f"cannot close session database: {exc}") from exc
|
|
|
|
def _initialize(self) -> None:
|
|
try:
|
|
with self._lock, self._connection:
|
|
self._connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS agent_sessions (
|
|
conversation_id TEXT PRIMARY KEY NOT NULL,
|
|
session_id TEXT NOT NULL,
|
|
external_subject_id TEXT,
|
|
metadata_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
except sqlite3.Error as exc:
|
|
raise SessionStoreError(f"cannot create session database schema: {exc}") from exc
|
|
|
|
def get(self, conversation_id: str) -> Optional[SessionRecord]:
|
|
conversation_id = self._required("conversation_id", conversation_id)
|
|
try:
|
|
with self._lock:
|
|
row = self._connection.execute(
|
|
"SELECT * FROM agent_sessions WHERE conversation_id = ?",
|
|
(conversation_id,),
|
|
).fetchone()
|
|
except sqlite3.Error as exc:
|
|
raise SessionStoreError(f"cannot read session mapping: {exc}") from exc
|
|
return self._record_from_row(row) if row is not None else None
|
|
|
|
def put(
|
|
self,
|
|
*,
|
|
conversation_id: str,
|
|
session_id: str,
|
|
external_subject_id: Optional[str],
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> SessionRecord:
|
|
conversation_id = self._required("conversation_id", conversation_id)
|
|
session_id = self._required("session_id", session_id)
|
|
metadata_value = dict(metadata or {})
|
|
try:
|
|
metadata_json = json.dumps(
|
|
metadata_value,
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
)
|
|
except (TypeError, ValueError) as exc:
|
|
raise SessionStoreError("session metadata must be JSON serializable") from exc
|
|
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
with self._lock, self._connection:
|
|
self._connection.execute(
|
|
"""
|
|
INSERT INTO agent_sessions (
|
|
conversation_id,
|
|
session_id,
|
|
external_subject_id,
|
|
metadata_json,
|
|
created_at,
|
|
updated_at
|
|
) VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(conversation_id) DO UPDATE SET
|
|
session_id = excluded.session_id,
|
|
external_subject_id = excluded.external_subject_id,
|
|
metadata_json = excluded.metadata_json,
|
|
updated_at = excluded.updated_at
|
|
""",
|
|
(
|
|
conversation_id,
|
|
session_id,
|
|
external_subject_id,
|
|
metadata_json,
|
|
now,
|
|
now,
|
|
),
|
|
)
|
|
row = self._connection.execute(
|
|
"SELECT * FROM agent_sessions WHERE conversation_id = ?",
|
|
(conversation_id,),
|
|
).fetchone()
|
|
except sqlite3.Error as exc:
|
|
raise SessionStoreError(f"cannot persist session mapping: {exc}") from exc
|
|
if row is None:
|
|
raise SessionStoreError("session mapping was not persisted")
|
|
return self._record_from_row(row)
|
|
|
|
def delete(self, conversation_id: str) -> bool:
|
|
conversation_id = self._required("conversation_id", conversation_id)
|
|
try:
|
|
with self._lock, self._connection:
|
|
cursor = self._connection.execute(
|
|
"DELETE FROM agent_sessions WHERE conversation_id = ?",
|
|
(conversation_id,),
|
|
)
|
|
except sqlite3.Error as exc:
|
|
raise SessionStoreError(f"cannot delete session mapping: {exc}") from exc
|
|
return cursor.rowcount > 0
|
|
|
|
@staticmethod
|
|
def _required(name: str, value: str) -> str:
|
|
normalized = str(value).strip()
|
|
if not normalized:
|
|
raise ValueError(f"{name} must not be empty")
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _record_from_row(row: sqlite3.Row) -> SessionRecord:
|
|
try:
|
|
metadata = json.loads(row["metadata_json"])
|
|
except (TypeError, ValueError) as exc:
|
|
raise SessionStoreError("persisted session metadata is invalid JSON") from exc
|
|
if not isinstance(metadata, dict):
|
|
raise SessionStoreError("persisted session metadata must be an object")
|
|
return SessionRecord(
|
|
conversation_id=row["conversation_id"],
|
|
session_id=row["session_id"],
|
|
external_subject_id=row["external_subject_id"],
|
|
metadata=metadata,
|
|
created_at=row["created_at"],
|
|
updated_at=row["updated_at"],
|
|
)
|