Spaces:
Sleeping
Sleeping
Download scripts/lib/chat/session_state.py from defenestrate/asa: direct link, hf CLI and curl.
- Browser
- Download file 4.59 kB
-
https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/chat/session_state.py
- Command line
-
hf download hf://spaces/defenestrate/asa/scripts/lib/chat/session_state.py
-
curl -L -o session_state.py https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/chat/session_state.py
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 | |
| 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 | |