feat: prepare ARR for controlled public deployment
This commit is contained in:
Wyndham ARR
222
agent_integration/sessions.py
Normal file
222
agent_integration/sessions.py
Normal file
@@ -0,0 +1,222 @@
|
||||
"""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"],
|
||||
)
|
||||
Reference in New Issue
Block a user