杜宇轩
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import json
|
|
import sqlite3
|
|
from typing import Optional
|
|
|
|
from .base import Session, SessionStore
|
|
|
|
_SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
user_id TEXT NOT NULL,
|
|
session_id TEXT NOT NULL,
|
|
messages TEXT NOT NULL, -- JSON-encoded message list
|
|
updated_at REAL NOT NULL,
|
|
PRIMARY KEY (user_id, session_id)
|
|
);
|
|
"""
|
|
|
|
|
|
class SQLiteSessionStore(SessionStore):
|
|
"""A simple SQLite-based session store."""
|
|
|
|
def __init__(self, db_path: str):
|
|
self.db_path = db_path
|
|
self._init_schema()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def _init_schema(self):
|
|
with self._connect() as conn:
|
|
conn.execute(_SCHEMA)
|
|
|
|
def load(self, user_id: str, session_id: str) -> Session:
|
|
with self._connect() as conn:
|
|
row: Optional[sqlite3.Row] = conn.execute(
|
|
"SELECT messages FROM agent_session WHERE user_id = ? AND session_id = ?",
|
|
(user_id, session_id),
|
|
).fetchone()
|
|
messages = []
|
|
if row is not None:
|
|
try:
|
|
messages = json.loads(row["messages"])
|
|
except (json.JSONDecodeError, TypeError):
|
|
messages = []
|
|
return Session(session_id=session_id, user_id=user_id, messages=messages)
|
|
|
|
def save(self, session: Session) -> None:
|
|
import time
|
|
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO agent_session (user_id, session_id, messages, updated_at)
|
|
VALUES (?, ?, ?, ?)
|
|
ON CONFLICT(user_id, session_id)
|
|
DO UPDATE SET messages=excluded.messages, updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
session.user_id,
|
|
session.session_id,
|
|
json.dumps(session.messages, ensure_ascii=False),
|
|
time.time()
|
|
),
|
|
)
|
|
conn.commit()
|
|
|
|
def clear(self, user_id: str, session_id: str) -> None:
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"DELETE FROM agent_session WHERE user_id = ? AND session_id = ?",
|
|
(user_id, session_id),
|
|
)
|
|
conn.commit()
|