core/memory: add async MemoryStore with default local user

Open SQLite via aiosqlite, create schema, ensure default user, and
CRUD sessions/messages with msgspec AnyChatMessage JSON.
This commit is contained in:
2026-07-14 17:25:14 +08:00
parent 8a1994a785
commit 0c5cc20d0b
5 changed files with 193 additions and 0 deletions
+2
View File
@@ -1 +1,3 @@
from . import config as config from . import config as config
from . import memory as memory
from . import runtime as runtime
+9
View File
@@ -0,0 +1,9 @@
from .database import DEFAULT_USER_EMAIL as DEFAULT_USER_EMAIL
from .database import DEFAULT_USER_NAME as DEFAULT_USER_NAME
from .database import MemoryStore as MemoryStore
from .database import Message as Message
from .database import Session as Session
from .database import UnsupportedDatabaseError as UnsupportedDatabaseError
from .database import User as User
from .database import build_async_url as build_async_url
from .database import create_engine as create_engine
+10
View File
@@ -0,0 +1,10 @@
from .engine import UnsupportedDatabaseError as UnsupportedDatabaseError
from .engine import build_async_url as build_async_url
from .engine import create_engine as create_engine
from .schema import Message as Message
from .schema import PlyngentBase as PlyngentBase
from .schema import Session as Session
from .schema import User as User
from .store import DEFAULT_USER_EMAIL as DEFAULT_USER_EMAIL
from .store import DEFAULT_USER_NAME as DEFAULT_USER_NAME
from .store import MemoryStore as MemoryStore
+35
View File
@@ -0,0 +1,35 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
if TYPE_CHECKING:
from plyngent.config.models import DatabaseConfig
class UnsupportedDatabaseError(NotImplementedError):
"""Raised when the configured database implementation is not supported."""
def build_async_url(config: DatabaseConfig) -> str:
"""Build a SQLAlchemy async database URL from :class:`DatabaseConfig`."""
implementation = config.implementation.lower()
if implementation != "sqlite":
msg = f"database implementation {config.implementation!r} is not supported"
raise UnsupportedDatabaseError(msg)
url = config.url
if url in {":memory:", ""}:
return "sqlite+aiosqlite:///:memory:"
if url.startswith("sqlite+aiosqlite://"):
return url
if url.startswith("sqlite://"):
return "sqlite+aiosqlite://" + url.removeprefix("sqlite://")
# File path (relative or absolute)
return f"sqlite+aiosqlite:///{url}"
def create_engine(config: DatabaseConfig, *, echo: bool = False) -> AsyncEngine:
"""Create an async SQLAlchemy engine for the given database config."""
return create_async_engine(build_async_url(config), echo=echo)
+137
View File
@@ -0,0 +1,137 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Self
import msgspec
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
from plyngent.config.models import DatabaseConfig
from plyngent.lmproto.openai_compatible.model import AnyChatMessage
from .engine import create_engine
from .schema import Message, PlyngentBase, Session, User
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
DEFAULT_USER_NAME = "local"
DEFAULT_USER_EMAIL = "local@localhost"
DEFAULT_USER_PASSWORD_HASH = ""
class MemoryStore:
"""Async persistence for users, chat sessions, and messages."""
_engine: AsyncEngine
_session_factory: async_sessionmaker[AsyncSession]
def __init__(self, engine: AsyncEngine) -> None:
self._engine = engine
self._session_factory = async_sessionmaker(engine, expire_on_commit=False)
@classmethod
async def open(
cls,
database: DatabaseConfig | Mapping[str, object],
*,
echo: bool = False,
init_schema: bool = True,
ensure_default_user: bool = True,
) -> Self:
"""Open a store from database config, optionally creating tables and default user."""
config = database if isinstance(database, DatabaseConfig) else msgspec.convert(dict(database), DatabaseConfig)
store = cls(create_engine(config, echo=echo))
if init_schema:
await store.create_schema()
if ensure_default_user:
_ = await store.ensure_default_user()
return store
async def create_schema(self) -> None:
"""Create all tables if they do not exist."""
async with self._engine.begin() as conn:
_ = await conn.run_sync(PlyngentBase.metadata.create_all)
async def close(self) -> None:
"""Dispose the underlying engine."""
await self._engine.dispose()
async def ensure_default_user(self) -> User:
"""Return the default local user, creating it if missing."""
async with self._session_factory() as session:
result = await session.execute(select(User).where(User.name == DEFAULT_USER_NAME))
user = result.scalar_one_or_none()
if user is not None:
return user
user = User(
name=DEFAULT_USER_NAME,
email=DEFAULT_USER_EMAIL,
password_hash=DEFAULT_USER_PASSWORD_HASH,
)
session.add(user)
await session.commit()
await session.refresh(user)
return user
async def get_user_by_name(self, name: str) -> User | None:
async with self._session_factory() as session:
result = await session.execute(select(User).where(User.name == name))
return result.scalar_one_or_none()
async def create_session(self, *, uid: int | None = None, name: str = "default") -> Session:
"""Create a chat session for ``uid`` (default local user when omitted)."""
if uid is None:
user = await self.ensure_default_user()
uid = user.uid
async with self._session_factory() as session:
row = Session(uid=uid, name=name)
session.add(row)
await session.commit()
await session.refresh(row)
return row
async def get_session(self, sid: int) -> Session | None:
async with self._session_factory() as session:
return await session.get(Session, sid)
async def list_sessions(self, *, uid: int | None = None) -> Sequence[Session]:
async with self._session_factory() as session:
stmt = select(Session).order_by(Session.sid)
if uid is not None:
stmt = stmt.where(Session.uid == uid)
result = await session.execute(stmt)
return result.scalars().all()
async def append_message(self, sid: int, message: AnyChatMessage) -> Message:
"""Append a chat message to a session with the next sequence number."""
data = msgspec.to_builtins(message)
if not isinstance(data, dict):
msg = "chat message must serialize to a JSON object"
raise TypeError(msg)
async with self._session_factory() as session:
result = await session.execute(
select(Message.seq).where(Message.sid == sid).order_by(Message.seq.desc()).limit(1)
)
last_seq = result.scalar_one_or_none()
seq = 0 if last_seq is None else last_seq + 1
row = Message(sid=sid, seq=seq, data=data)
session.add(row)
await session.commit()
await session.refresh(row)
return row
async def list_messages(self, sid: int) -> list[AnyChatMessage]:
"""Load chat messages for a session in sequence order."""
async with self._session_factory() as session:
result = await session.execute(select(Message).where(Message.sid == sid).order_by(Message.seq))
rows = result.scalars().all()
return [msgspec.convert(row.data, type=AnyChatMessage) for row in rows]
async def list_message_rows(self, sid: int) -> Sequence[Message]:
"""Load raw message rows for a session in sequence order."""
async with self._session_factory() as session:
result = await session.execute(select(Message).where(Message.sid == sid).order_by(Message.seq))
return result.scalars().all()