diff --git a/src/plyngent/__init__.py b/src/plyngent/__init__.py index 4715794..65a5492 100644 --- a/src/plyngent/__init__.py +++ b/src/plyngent/__init__.py @@ -1 +1,3 @@ from . import config as config +from . import memory as memory +from . import runtime as runtime diff --git a/src/plyngent/memory/__init__.py b/src/plyngent/memory/__init__.py index e69de29..1e0ff83 100644 --- a/src/plyngent/memory/__init__.py +++ b/src/plyngent/memory/__init__.py @@ -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 diff --git a/src/plyngent/memory/database/__init__.py b/src/plyngent/memory/database/__init__.py index e69de29..641b1e2 100644 --- a/src/plyngent/memory/database/__init__.py +++ b/src/plyngent/memory/database/__init__.py @@ -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 diff --git a/src/plyngent/memory/database/engine.py b/src/plyngent/memory/database/engine.py new file mode 100644 index 0000000..47039ca --- /dev/null +++ b/src/plyngent/memory/database/engine.py @@ -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) diff --git a/src/plyngent/memory/database/store.py b/src/plyngent/memory/database/store.py new file mode 100644 index 0000000..2603601 --- /dev/null +++ b/src/plyngent/memory/database/store.py @@ -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()