Merge pull request #1 from Open-LLM-VTuber/dev

Dev

Author: Stewitch

alembic.ini

@@ -0,0 +1,35 @@
+[alembic]
+script_location = bot_sdk/db/migrations
+prepend_sys_path = .
+
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+
+[logger_sqlalchemy]
+level = WARN
+handlers = console
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers = console
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stdout,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s

bot_sdk/__init__.py

@@ -14,6 +14,7 @@
 from .runner import BotRunner
 from .logging import setup_logging
 from .storage import BotStorage, CachedStorage
+from .config import StorageConfig
 from .models import (
 	Event,
 	EventsResponse,
@@ -29,6 +30,18 @@
     UpdatePresenceRequest,
     GetUserGroupsRequest,
 	GetUserGroupsResponse,
+    DataModel,
+    Timestamped,
+    IDModel,
+    SoftDelete,
+)
+from .db import (
+    Base,
+    make_sqlite_url,
+    create_engine,
+    create_sessionmaker,
+    session_scope,
+    AsyncRepository,
 )
 
 __all__ = [
@@ -45,6 +58,7 @@
 	"setup_logging",
 	"BotStorage",
 	"CachedStorage",
+    "StorageConfig",
 	"Event",
 	"EventsResponse",
 	"Message",
@@ -61,4 +75,14 @@
 	"GetUserGroupsResponse",
     "Validator",
 	"OptionValidator",
+	"DataModel",
+	"Timestamped",
+	"IDModel",
+	"SoftDelete",
+	"Base",
+	"make_sqlite_url",
+	"create_engine",
+	"create_sessionmaker",
+	"session_scope",
+	"AsyncRepository",
 ]

bot_sdk/bot.py

@@ -11,7 +11,7 @@
 from .commands import CommandInvocation, CommandParser, CommandSpec, CommandArgument
 from .storage import BotStorage
 from .permissions import PermissionPolicy
-from .config import BotLocalConfig, load_bot_local_config, save_bot_local_config
+from .config import BotLocalConfig, StorageConfig, load_bot_local_config, save_bot_local_config
 from .models import (
     Event,
     Message,
@@ -37,6 +37,7 @@ class BaseBot(abc.ABC):
     # Storage configuration
     enable_storage = True
     storage_path: Optional[str] = None  # Defaults to "bot_data/{bot_name}.db"
+    storage_config: Optional[StorageConfig] = None
 
     def __init__(self, client: AsyncClient) -> None:
         self.client = client
@@ -54,6 +55,10 @@ def __init__(self, client: AsyncClient) -> None:
     def set_runner(self, runner: "BotRunner") -> None:
         """Called by BotRunner to allow commands to signal runner actions."""
         self._runner = runner
+
+    def set_storage_config(self, storage_config: Optional[StorageConfig]) -> None:
+        """Inject per-bot storage configuration from runner/config layer."""
+        self.storage_config = storage_config
 
     async def post_init(self) -> None:
         """Hook for post-initialization logic. Override if needed.
@@ -76,9 +81,14 @@ async def _init_storage(self) -> None:
             return
         if self.storage_path is None:
             self.storage_path = f"bot_data/{bot_name}.db"
+        auto_cfg = self.storage_config or StorageConfig()
         self.storage = BotStorage(
             db_path=self.storage_path,
-            namespace=f"bot_{bot_name}"
+            namespace=f"bot_{bot_name}",
+            auto_cache=auto_cfg.auto_cache,
+            auto_flush_interval=auto_cfg.auto_flush_interval,
+            auto_flush_retry=auto_cfg.auto_flush_retry,
+            auto_flush_max_retries=auto_cfg.auto_flush_max_retries,
         )
         logger.info(f"Initialized storage at {self.storage_path}")
         # Initialize permissions helper

bot_sdk/config.py

@@ -7,6 +7,21 @@
 from pathlib import Path
 
 
+class StorageConfig(BaseModel):
+    """Per-bot storage options for KV backend.
+
+    auto_cache: enable always-on cache with periodic flush
+    auto_flush_interval: seconds between flush attempts
+    auto_flush_retry: retry delay after a flush failure (e.g., DB locked)
+    auto_flush_max_retries: max retries per key per flush cycle
+    """
+
+    auto_cache: bool = False
+    auto_flush_interval: float = 5.0
+    auto_flush_retry: float = 1.0
+    auto_flush_max_retries: int = 3
+
+
 class BotConfig(BaseModel):
     name: str
     module: Optional[str] = None
@@ -15,6 +30,7 @@ class BotConfig(BaseModel):
     zuliprc: Optional[str] = None
     event_types: list[str] = Field(default_factory=lambda: ["message"])
     config: dict[str, Any] = Field(default_factory=dict)
+    storage: Optional[StorageConfig] = None
 
 
 class AppConfig(BaseModel):
@@ -66,6 +82,7 @@ def save_bot_local_config(path: str | Path, config: BotLocalConfig) -> None:
 __all__ = [
     "AppConfig",
     "BotConfig",
+    "StorageConfig",
     "load_config",
     "BotLocalConfig",
     "load_bot_local_config",

bot_sdk/db/__init__.py

@@ -0,0 +1,17 @@
+from .database import (
+    Base,
+    make_sqlite_url,
+    create_engine,
+    create_sessionmaker,
+    session_scope,
+)
+from .repository import AsyncRepository
+
+__all__ = [
+    "Base",
+    "make_sqlite_url",
+    "create_engine",
+    "create_sessionmaker",
+    "session_scope",
+    "AsyncRepository",
+]

bot_sdk/db/database.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncIterator
+
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.orm import DeclarativeBase
+
+__all__ = [
+    "Base",
+    "make_sqlite_url",
+    "create_engine",
+    "create_sessionmaker",
+    "session_scope",
+]
+
+
+class Base(DeclarativeBase):
+    """Declarative base for ORM models."""
+
+
+def make_sqlite_url(db_path: str | Path) -> str:
+    """Build an async SQLite URL using aiosqlite driver."""
+    path = Path(db_path).expanduser().resolve()
+    return f"sqlite+aiosqlite:///{path.as_posix()}"
+
+
+def create_engine(url: str, *, echo: bool = False) -> AsyncEngine:
+    """Create an async SQLAlchemy engine.
+
+    Defaults are tuned for SQLite; adjust pool settings if swapping drivers.
+    """
+    return create_async_engine(url, echo=echo, future=True)
+
+
+def create_sessionmaker(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
+    """Return an async sessionmaker bound to the given engine."""
+    return async_sessionmaker(engine, expire_on_commit=False)
+
+
+@asynccontextmanager
+async def session_scope(session_factory: async_sessionmaker[AsyncSession]) -> AsyncIterator[AsyncSession]:
+    """Async context manager that yields a session and commits/rolls back safely."""
+    session = session_factory()
+    try:
+        yield session
+        await session.commit()
+    except Exception:
+        await session.rollback()
+        raise
+    finally:
+        await session.close()

bot_sdk/db/migrations/README.md

@@ -0,0 +1,5 @@
+# Migrations layout
+
+- SDK-owned migrations live here and should only cover shared/SDK tables.
+- Bot-specific schemas should keep their own Alembic env/migrations under each bot directory (e.g., `bots/my_bot/migrations`).
+- Point `alembic.ini` in each bot to its own `script_location`; the SDK env in this folder remains for shared models only.

bot_sdk/db/migrations/env.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from logging.config import fileConfig
+from pathlib import Path
+
+from alembic import context
+from sqlalchemy import pool
+from sqlalchemy.engine import Connection
+from sqlalchemy.ext.asyncio import async_engine_from_config
+
+from bot_sdk.db.database import Base, make_sqlite_url
+
+config = context.config
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# Metadata for autogenerate
+# Add ORM models to Base to include them in migrations
+
+_target_metadata = Base.metadata
+
+
+def get_url() -> str:
+    env_url = os.getenv("DATABASE_URL")
+    if env_url:
+        return env_url
+    default_path = Path("bot_data/db.sqlite").resolve()
+    return make_sqlite_url(default_path)
+
+
+def run_migrations_offline() -> None:
+    url = get_url()
+    context.configure(
+        url=url,
+        target_metadata=_target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def do_run_migrations(connection: Connection) -> None:
+    context.configure(connection=connection, target_metadata=_target_metadata)
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    configuration = config.get_section(config.config_ini_section) or {}
+    configuration["sqlalchemy.url"] = get_url()
+
+    connectable = async_engine_from_config(
+        configuration,
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+        future=True,
+    )
+
+    async def run() -> None:
+        async with connectable.connect() as connection:
+            await connection.run_sync(do_run_migrations)
+
+    asyncio.run(run())
+
+
+def run_migrations() -> None:
+    if context.is_offline_mode():
+        run_migrations_offline()
+    else:
+        run_migrations_online()
+
+
+run_migrations()

bot_sdk/db/migrations/versions/0001_initial.py

@@ -0,0 +1,27 @@
+"""Initial empty migration placeholder.
+
+Revision ID: 0001_initial
+Revises: 
+Create Date: 2026-01-10
+"""
+
+from __future__ import annotations
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "0001_initial"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Add initial tables here when introducing ORM models.
+    pass
+
+
+def downgrade() -> None:
+    # Drop initial tables here when introducing ORM models.
+    pass

bot_sdk/db/repository.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Generic, Iterable, Optional, Sequence, TypeVar
+
+from sqlalchemy import Select, delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+ModelT = TypeVar("ModelT")
+
+__all__ = ["AsyncRepository"]
+
+
+class AsyncRepository(Generic[ModelT]):
+    """Lightweight base repository with common CRUD helpers.
+
+    Subclass and set ``model`` to your ORM mapped class.
+    Methods open their own unit-of-work on the provided session.
+    """
+
+    model: type[ModelT]
+
+    def __init__(self, session: AsyncSession):
+        self.session = session
+
+    async def get(self, pk: object) -> Optional[ModelT]:
+        return await self.session.get(self.model, pk)
+
+    async def add(self, instance: ModelT) -> ModelT:
+        self.session.add(instance)
+        await self.session.flush()
+        await self.session.refresh(instance)
+        return instance
+
+    async def delete(self, pk: object) -> bool:
+        instance = await self.get(pk)
+        if instance is None:
+            return False
+        await self.session.delete(instance)
+        await self.session.flush()
+        return True
+
+    async def list(self, *, limit: int = 100, offset: int = 0) -> Sequence[ModelT]:
+        stmt: Select[ModelT] = select(self.model).limit(limit).offset(offset)
+        rows = await self.session.execute(stmt)
+        return rows.scalars().all()