db.py

"""
db.py — SQLite persistence for chat history
Stored in /data/chat.db (persistent HF Space volume)
"""

import sqlite3
import json
import os
from datetime import datetime

DB_PATH = os.environ.get("DB_PATH", "/data/chat.db")


def get_conn():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    with get_conn() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS conversations (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id     TEXT    NOT NULL,
                title       TEXT    NOT NULL DEFAULT 'Nouvelle conversation',
                model_id    TEXT,
                created_at  TEXT    NOT NULL,
                updated_at  TEXT    NOT NULL
            );

            CREATE TABLE IF NOT EXISTS messages (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id INTEGER NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
                role            TEXT    NOT NULL CHECK(role IN ('user','assistant','system')),
                content         TEXT    NOT NULL,
                created_at      TEXT    NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_conv_user   ON conversations(user_id, updated_at DESC);
            CREATE INDEX IF NOT EXISTS idx_msg_conv    ON messages(conversation_id, id ASC);
        """)


# ── Conversations ──────────────────────────────────────────────

def list_conversations(user_id: str) -> list[dict]:
    with get_conn() as conn:
        rows = conn.execute(
            """SELECT c.id, c.title, c.model_id, c.updated_at,
                      COUNT(m.id) as msg_count
               FROM conversations c
               LEFT JOIN messages m ON m.conversation_id = c.id
               WHERE c.user_id = ?
               GROUP BY c.id
               ORDER BY c.updated_at DESC
               LIMIT 100""",
            (user_id,),
        ).fetchall()
    return [dict(r) for r in rows]


def create_conversation(user_id: str, title: str = "Nouvelle conversation", model_id: str = "") -> int:
    now = datetime.utcnow().isoformat()
    with get_conn() as conn:
        cur = conn.execute(
            "INSERT INTO conversations(user_id, title, model_id, created_at, updated_at) VALUES(?,?,?,?,?)",
            (user_id, title, model_id, now, now),
        )
        return cur.lastrowid


def rename_conversation(conv_id: int, user_id: str, title: str):
    now = datetime.utcnow().isoformat()
    with get_conn() as conn:
        conn.execute(
            "UPDATE conversations SET title=?, updated_at=? WHERE id=? AND user_id=?",
            (title[:80], now, conv_id, user_id),
        )


def delete_conversation(conv_id: int, user_id: str):
    with get_conn() as conn:
        conn.execute(
            "DELETE FROM conversations WHERE id=? AND user_id=?",
            (conv_id, user_id),
        )


def update_conversation_model(conv_id: int, model_id: str):
    now = datetime.utcnow().isoformat()
    with get_conn() as conn:
        conn.execute(
            "UPDATE conversations SET model_id=?, updated_at=? WHERE id=?",
            (model_id, now, conv_id),
        )


# ── Messages ───────────────────────────────────────────────────

def get_messages(conv_id: int) -> list[dict]:
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT role, content, created_at FROM messages WHERE conversation_id=? ORDER BY id ASC",
            (conv_id,),
        ).fetchall()
    return [dict(r) for r in rows]


def add_message(conv_id: int, role: str, content: str):
    now = datetime.utcnow().isoformat()
    with get_conn() as conn:
        conn.execute(
            "INSERT INTO messages(conversation_id, role, content, created_at) VALUES(?,?,?,?)",
            (conv_id, role, content, now),
        )
        conn.execute(
            "UPDATE conversations SET updated_at=? WHERE id=?",
            (now, conv_id),
        )


def auto_title_from_message(msg: str) -> str:
    """Generate a short title from the first user message."""
    title = msg.strip().replace("\n", " ")
    return title[:60] + ("…" if len(title) > 60 else "")