杜宇轩
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
import json
|
|
import os
|
|
|
|
from .base import Session, SessionStore, sanitize_id
|
|
|
|
|
|
class JSONSessionStore(SessionStore):
|
|
"""A simple JSON file-based session store."""
|
|
|
|
def __init__(self, base_dir: str):
|
|
self.base_dir = base_dir
|
|
os.makedirs(base_dir, exist_ok=True)
|
|
|
|
def _path(self, user_id: str, session_id: str) -> str:
|
|
return os.path.join(
|
|
self.base_dir,
|
|
f"{sanitize_id(user_id, 'anon')}__{sanitize_id(session_id, 'default')}.json")
|
|
|
|
def load(self, user_id: str, session_id: str) -> Session:
|
|
path = self._path(user_id, session_id)
|
|
messages = []
|
|
if os.path.exists(path):
|
|
try:
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
messages = json.load(f).get('messages', [])
|
|
|
|
except (json.JSONDecodeError, OSError):
|
|
messages = []
|
|
|
|
return Session(session_id=session_id, user_id=user_id, messages=messages)
|
|
|
|
|
|
def save(self, session: Session) -> None:
|
|
path = self._path(session.user_id, session.session_id)
|
|
tmp = path + ".tmp"
|
|
|
|
with open(path, 'w', encoding='utf-8') as f:
|
|
json.dump(
|
|
{
|
|
"user_id": session.user_id,
|
|
"session_id": session.session_id,
|
|
"messages": session.messages,
|
|
},
|
|
f,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
os.replace(tmp, path)
|
|
|
|
def clear(self, user_id: str, session_id: str) -> None:
|
|
path = self._path(user_id, session_id)
|
|
if os.path.exists(path):
|
|
os.remove(path) |