hobbs-iac/compose/hf-pdocker-01/voicecoach-v2/agent/store.py

"""SQLite session store — WAL mode, aiosqlite."""
import json
import os
import time
from typing import Optional

import aiosqlite

DB_PATH = os.getenv("DB_PATH", "/data/sessions.db")


async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("PRAGMA journal_mode=WAL")
        await db.executescript("""
            CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                started_at REAL NOT NULL,
                ended_at REAL,
                mode TEXT DEFAULT 'coaching',
                summary TEXT
            );

            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                vocal_energy REAL,
                vocal_emotion TEXT,
                ts REAL NOT NULL,
                FOREIGN KEY(session_id) REFERENCES sessions(id)
            );

            CREATE TABLE IF NOT EXISTS profile (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL,
                updated_at REAL NOT NULL
            );
        """)
        await db.commit()


async def create_session(session_id: str, mode: str = "coaching"):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR IGNORE INTO sessions (id, started_at, mode) VALUES (?, ?, ?)",
            (session_id, time.time(), mode),
        )
        await db.commit()


async def end_session(session_id: str, summary: Optional[str] = None):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE sessions SET ended_at=?, summary=? WHERE id=?",
            (time.time(), summary, session_id),
        )
        await db.commit()


async def add_message(
    session_id: str,
    role: str,
    content: str,
    vocal_energy: Optional[float] = None,
    vocal_emotion: Optional[str] = None,
):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO messages (session_id, role, content, vocal_energy, vocal_emotion, ts) VALUES (?,?,?,?,?,?)",
            (session_id, role, content, vocal_energy, vocal_emotion, time.time()),
        )
        await db.commit()


async def get_recent_context(limit: int = 3) -> str:
    """Return a brief context string from the last N sessions for injection into system prompt."""
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT id, started_at, mode, summary FROM sessions WHERE ended_at IS NOT NULL ORDER BY ended_at DESC LIMIT ?",
            (limit,),
        ) as cur:
            rows = await cur.fetchall()

    if not rows:
        return ""

    parts = ["Recent session context:"]
    for session_id, started_at, mode, summary in rows:
        import datetime
        dt = datetime.datetime.fromtimestamp(started_at).strftime("%Y-%m-%d")
        if summary:
            parts.append(f"- {dt} ({mode}): {summary}")
    return "\n".join(parts)


async def get_sessions(limit: int = 20) -> list[dict]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT id, started_at, ended_at, mode, summary FROM sessions ORDER BY started_at DESC LIMIT ?",
            (limit,),
        ) as cur:
            rows = await cur.fetchall()
    return [
        {
            "id": r[0],
            "started_at": r[1],
            "ended_at": r[2],
            "mode": r[3],
            "summary": r[4],
        }
        for r in rows
    ]


async def get_session_messages(session_id: str) -> list[dict]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT role, content, vocal_emotion, ts FROM messages WHERE session_id=? ORDER BY ts",
            (session_id,),
        ) as cur:
            rows = await cur.fetchall()
    return [{"role": r[0], "content": r[1], "vocal_emotion": r[2], "ts": r[3]} for r in rows]


async def get_profile(key: str) -> Optional[str]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT value FROM profile WHERE key=?", (key,)) as cur:
            row = await cur.fetchone()
    return row[0] if row else None


async def set_profile(key: str, value: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT OR REPLACE INTO profile (key, value, updated_at) VALUES (?,?,?)",
            (key, value, time.time()),
        )
        await db.commit()