132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
import asyncio
|
|
from datetime import UTC, datetime
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
from uuid import UUID, uuid4
|
|
|
|
from pydantic import ValidationError
|
|
|
|
from domain import Session
|
|
|
|
|
|
class SessionNotFoundError(Exception):
|
|
pass
|
|
|
|
|
|
class SessionStorageError(Exception):
|
|
pass
|
|
|
|
|
|
class JsonSessionStorage:
|
|
def __init__(self, data_dir: Path) -> None:
|
|
self.data_dir = data_dir
|
|
|
|
async def create(self, system_prompt: str, user_id: str | None = None) -> Session:
|
|
return await asyncio.to_thread(self._create, system_prompt, user_id)
|
|
|
|
async def read(self, session_id: str) -> Session:
|
|
return await asyncio.to_thread(self._read, session_id)
|
|
|
|
async def write(self, session: Session) -> None:
|
|
await asyncio.to_thread(self._write, session)
|
|
|
|
async def list_by_user(self, user_id: str) -> list[Session]:
|
|
return await asyncio.to_thread(self._list_by_user, user_id)
|
|
|
|
def _path(self, session_id: str) -> Path:
|
|
try:
|
|
normalized = str(UUID(session_id))
|
|
except ValueError as exc:
|
|
raise SessionNotFoundError(session_id) from exc
|
|
return self.data_dir / f"{normalized}.json"
|
|
|
|
def _create(self, system_prompt: str, user_id: str | None) -> Session:
|
|
try:
|
|
self.data_dir.mkdir(parents=True, exist_ok=True)
|
|
while True:
|
|
now = datetime.now(UTC)
|
|
session = Session(
|
|
user_id=user_id,
|
|
session_id=str(uuid4()),
|
|
system_prompt=system_prompt,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
path = self._path(session.session_id)
|
|
try:
|
|
with path.open("x", encoding="utf-8") as file:
|
|
json.dump(
|
|
session.model_dump(mode="json"),
|
|
file,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
return session
|
|
except FileExistsError:
|
|
continue
|
|
except OSError as exc:
|
|
raise SessionStorageError("could not create session") from exc
|
|
|
|
def _read(self, session_id: str) -> Session:
|
|
path = self._path(session_id)
|
|
try:
|
|
with path.open(encoding="utf-8") as file:
|
|
return Session.model_validate(json.load(file))
|
|
except FileNotFoundError as exc:
|
|
raise SessionNotFoundError(session_id) from exc
|
|
except (OSError, json.JSONDecodeError, ValidationError) as exc:
|
|
raise SessionStorageError("could not read session") from exc
|
|
|
|
def _list_by_user(self, user_id: str) -> list[Session]:
|
|
sessions: list[Session] = []
|
|
try:
|
|
for path in self.data_dir.glob("*.json"):
|
|
try:
|
|
UUID(path.stem)
|
|
except ValueError:
|
|
continue
|
|
with path.open(encoding="utf-8") as file:
|
|
session = Session.model_validate(json.load(file))
|
|
if session.user_id == user_id:
|
|
sessions.append(session)
|
|
except (OSError, json.JSONDecodeError, ValidationError) as exc:
|
|
raise SessionStorageError("could not list sessions") from exc
|
|
return sessions
|
|
|
|
def _write(self, session: Session) -> None:
|
|
path = self._path(session.session_id)
|
|
if not path.exists():
|
|
raise SessionNotFoundError(session.session_id)
|
|
|
|
temporary_path: str | None = None
|
|
try:
|
|
self.data_dir.mkdir(parents=True, exist_ok=True)
|
|
with NamedTemporaryFile(
|
|
mode="w",
|
|
encoding="utf-8",
|
|
dir=self.data_dir,
|
|
prefix=f".{session.session_id}.",
|
|
suffix=".tmp",
|
|
delete=False,
|
|
) as file:
|
|
temporary_path = file.name
|
|
json.dump(
|
|
session.model_dump(mode="json"),
|
|
file,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
file.flush()
|
|
os.fsync(file.fileno())
|
|
os.replace(temporary_path, path)
|
|
except OSError as exc:
|
|
raise SessionStorageError("could not write session") from exc
|
|
finally:
|
|
if temporary_path:
|
|
try:
|
|
Path(temporary_path).unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|