Spaces:
Paused
Paused
| """The shared chat store, and the ceiling it stops at. | |
| Ephemeral is the default and needs nothing from here: those chats live in the | |
| viewer's browser. This is the other option -- a Bit Trading Company database a | |
| signed-in viewer's chats can follow them across devices. | |
| Two rules shape the whole module: | |
| * **A chat is scoped to whoever wrote it.** No owner, no store: an unauthenticated | |
| viewer would otherwise be writing into a bucket everyone else can read. | |
| * **Full is a normal state, not an error.** Past the ceiling the store simply | |
| stops accepting writes and the client keeps everything on the device. There is | |
| a real budget behind this, and silently growing past it is how a demo becomes | |
| an invoice. | |
| Enabled by ``ASE_CHAT_DB`` (a directory). Unset -- which is the case on the | |
| Space today -- and every endpoint reports the store as unavailable. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import re | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| # A few gigabytes. Raise it if the data turns out to be worth the space; the | |
| # client needs no change, it reads the ceiling out of /api/config. | |
| DEFAULT_MAX_BYTES = 3 * 1024**3 | |
| # One conversation with a hundred turns and a dozen canvases is well under this. | |
| MAX_CHAT_BYTES = 512 * 1024 | |
| _SAFE_ID = re.compile(r"^[A-Za-z0-9_-]{1,64}$") | |
| class StoreState: | |
| remote: bool | |
| max_bytes: int | |
| used_bytes: int | |
| full: bool | |
| detail: str = "" | |
| def to_json(self) -> dict: | |
| return { | |
| "remote": self.remote, | |
| "maxBytes": self.max_bytes, | |
| "usedBytes": self.used_bytes, | |
| "full": self.full, | |
| "detail": self.detail, | |
| } | |
| def root() -> Path | None: | |
| """Where chats are kept, or None when no store is configured.""" | |
| configured = os.environ.get("ASE_CHAT_DB") | |
| return Path(configured) if configured else None | |
| def max_bytes() -> int: | |
| try: | |
| return max(0, int(os.environ.get("ASE_CHAT_DB_MAX_BYTES", DEFAULT_MAX_BYTES))) | |
| except ValueError: | |
| return DEFAULT_MAX_BYTES | |
| def used_bytes(base: Path) -> int: | |
| total = 0 | |
| for path in base.rglob("*.json"): | |
| try: | |
| total += path.stat().st_size | |
| except OSError: | |
| continue | |
| return total | |
| def state(owner: str | None) -> StoreState: | |
| """What the client should be told about the store, for this viewer.""" | |
| base = root() | |
| if base is None: | |
| return StoreState(False, 0, 0, False, "No shared store is configured on this server.") | |
| if not owner: | |
| return StoreState(False, max_bytes(), 0, False, | |
| "Sign in with Hugging Face to use the shared store.") | |
| try: | |
| base.mkdir(parents=True, exist_ok=True) | |
| used = used_bytes(base) | |
| except OSError as e: | |
| return StoreState(False, max_bytes(), 0, False, f"The store is unreachable: {e}") | |
| cap = max_bytes() | |
| full = used >= cap | |
| return StoreState(True, cap, used, full, | |
| "Full — new chats stay on the device." if full else "") | |
| def _dir(base: Path, owner: str) -> Path: | |
| # The owner is a Hugging Face username, but it arrives from a token rather | |
| # than from our own validation, so it is hashed into a safe segment instead | |
| # of being trusted as a path component. | |
| safe = re.sub(r"[^A-Za-z0-9_-]", "_", owner)[:48] | |
| return base / safe | |
| def load_all(owner: str) -> list[dict]: | |
| base = root() | |
| if base is None or not owner: | |
| return [] | |
| folder = _dir(base, owner) | |
| if not folder.is_dir(): | |
| return [] | |
| out = [] | |
| for path in sorted(folder.glob("*.json")): | |
| try: | |
| out.append(json.loads(path.read_text("utf-8"))) | |
| except (OSError, json.JSONDecodeError): | |
| # One unreadable chat must not take the whole history down with it. | |
| continue | |
| return out | |
| def save(owner: str, chat_id: str, chat: dict) -> StoreState: | |
| """Write one chat. A refusal is reported as state, never as an exception.""" | |
| st = state(owner) | |
| if not st.remote or st.full or not _SAFE_ID.match(chat_id): | |
| return st | |
| if chat.get("ephemeral"): | |
| # An ephemeral chat is never written, whatever the request says. The | |
| # client enforces this too; doing it in both places is deliberate. | |
| return st | |
| body = json.dumps(chat, ensure_ascii=False) | |
| if len(body.encode("utf-8")) > MAX_CHAT_BYTES: | |
| st.detail = "That chat is too large for the shared store." | |
| return st | |
| base = root() | |
| assert base is not None | |
| folder = _dir(base, owner) | |
| try: | |
| folder.mkdir(parents=True, exist_ok=True) | |
| # Written beside the target and moved into place, so a crash mid-write | |
| # leaves the previous version rather than half a file. | |
| tmp = folder / f".{chat_id}.tmp" | |
| tmp.write_text(body, "utf-8") | |
| tmp.replace(folder / f"{chat_id}.json") | |
| except OSError as e: | |
| st.detail = f"Could not write: {e}" | |
| return st | |
| return state(owner) | |
| def delete(owner: str, chat_id: str) -> StoreState: | |
| base = root() | |
| if base is not None and owner and _SAFE_ID.match(chat_id): | |
| try: | |
| (_dir(base, owner) / f"{chat_id}.json").unlink(missing_ok=True) | |
| except OSError: | |
| pass | |
| return state(owner) | |