asa / scripts /lib /chat /session_state.py
defenestrate's picture
ASA Space: squashed history (instructions and settings live in the private pack)
526950e
Raw History Blame Contribute Delete
4.59 kB
"""Per-conversation dedupe state, persisted to
derived/student-chat/sessions/{conv_id}.json.
The dedupe set lets the orchestrator skip re-retrieving a doc that was
already pulled earlier in the same conversation. Persisted so a Ctrl-C
on the Gradio app doesn't lose continuity for in-flight chats.
"""
from __future__ import annotations
import json
import logging
import os
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from scripts.lib.runtime_config import SESSIONS_DIR
if sys.platform == "win32":
import msvcrt
def _try_lock(fileobj) -> bool:
try:
msvcrt.locking(fileobj.fileno(), msvcrt.LK_NBLCK, 1)
return True
except OSError:
return False
def _unlock(fileobj) -> None:
try:
fileobj.seek(0)
msvcrt.locking(fileobj.fileno(), msvcrt.LK_UNLCK, 1)
except OSError:
pass
else:
import fcntl
def _try_lock(fileobj) -> bool:
try:
fcntl.flock(fileobj.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
return True
except OSError:
return False
def _unlock(fileobj) -> None:
try:
fcntl.flock(fileobj.fileno(), fcntl.LOCK_UN)
except OSError:
pass
logger = logging.getLogger(__name__)
_SESSIONS_DIR = SESSIONS_DIR
@dataclass
class SessionState:
conversation_id: str
seen_doc_ids: set[str] = field(default_factory=set)
created_at: str = ""
updated_at: str = ""
def _now_iso() -> str:
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
def _path_for(conv_id: str, sessions_dir: Path) -> Path:
return sessions_dir / f"{conv_id}.json"
def load_or_create(conv_id: str, *, sessions_dir: Path = _SESSIONS_DIR) -> SessionState:
"""Load an existing session or return a fresh one. Never raises."""
path = _path_for(conv_id, sessions_dir)
if path.is_file():
try:
data = json.loads(path.read_text(encoding="utf-8"))
return SessionState(
conversation_id=data.get("conversation_id", conv_id),
seen_doc_ids=set(data.get("seen_doc_ids", [])),
created_at=data.get("created_at", _now_iso()),
updated_at=data.get("updated_at", _now_iso()),
)
except (json.JSONDecodeError, OSError) as e:
logger.warning(
"session_state.load_or_create: failed to read %s (%s) — starting fresh",
path, e,
)
now = _now_iso()
return SessionState(conversation_id=conv_id, created_at=now, updated_at=now)
def record_retrieved(state: SessionState, doc_ids: list[str]) -> None:
"""Add doc_ids to the dedupe set. In-place."""
state.seen_doc_ids.update(doc_ids)
state.updated_at = _now_iso()
def persist(state: SessionState, *, sessions_dir: Path = _SESSIONS_DIR) -> None:
"""Atomically write the state to disk via stage-and-rename, with a
short non-blocking lock to guard against concurrent writes for the
same conv_id. If lock acquisition fails, log a warning and skip the
write — the next turn will re-persist."""
sessions_dir.mkdir(parents=True, exist_ok=True)
path = _path_for(state.conversation_id, sessions_dir)
payload = {
"conversation_id": state.conversation_id,
"seen_doc_ids": sorted(state.seen_doc_ids),
"created_at": state.created_at,
"updated_at": state.updated_at,
}
staging = path.with_name(path.name + ".tmp")
try:
with open(staging, "w", encoding="utf-8") as f:
if not _try_lock(f):
logger.warning(
"session_state.persist: lock contention on %s — skipping write",
staging,
)
# Drop the staging file the with-block would otherwise leave.
# Use missing_ok in case f.close() races with another process.
try:
f.close()
staging.unlink(missing_ok=True)
except OSError:
pass
return
try:
f.write(json.dumps(payload, ensure_ascii=False))
f.flush()
os.fsync(f.fileno())
finally:
_unlock(f)
os.replace(staging, path)
except BaseException:
if staging.exists():
try:
os.unlink(staging)
except OSError:
pass
raise