Harshavard21/FinRAG

1"""2app/db/chat_store.py3=====================4SQLite-backed chat history store.5Zero dependencies beyond Python stdlib — single .db file on disk.6 7Schema:8    conversations — one row per chat session (id, title, created_at, mode)9    messages      — one row per message (id, conversation_id, role, content, metadata_json)10"""11 12from __future__ import annotations13 14import json15import sqlite316import uuid17from datetime import datetime18from pathlib import Path19from typing import Optional20 21DB_PATH = Path("data/chat_history.db")22 23 24def _conn() -> sqlite3.Connection:25    DB_PATH.parent.mkdir(parents=True, exist_ok=True)26    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)27    conn.row_factory = sqlite3.Row28    return conn29 30 31def init_db() -> None:32    """Create tables if they don't exist."""33    with _conn() as conn:34        conn.executescript("""35            CREATE TABLE IF NOT EXISTS conversations (36                id          TEXT PRIMARY KEY,37                title       TEXT NOT NULL DEFAULT 'New Chat',38                mode        TEXT NOT NULL DEFAULT 'chat',39                created_at  TEXT NOT NULL,40                updated_at  TEXT NOT NULL41            );42 43            CREATE TABLE IF NOT EXISTS messages (44                id              INTEGER PRIMARY KEY AUTOINCREMENT,45                conversation_id TEXT NOT NULL,46                role            TEXT NOT NULL,47                content         TEXT NOT NULL,48                metadata        TEXT,49                created_at      TEXT NOT NULL,50                FOREIGN KEY (conversation_id) REFERENCES conversations(id)51            );52        """)53 54 55# ------------------------------------------------------------------ #56# Conversations57# ------------------------------------------------------------------ #58 59def new_conversation(mode: str = "chat") -> str:60    """Create a new conversation, return its ID."""61    cid = str(uuid.uuid4())62    now = datetime.utcnow().isoformat()63    with _conn() as conn:64        conn.execute(65            "INSERT INTO conversations (id, title, mode, created_at, updated_at) VALUES (?,?,?,?,?)",66            (cid, "New Chat", mode, now, now),67        )68    return cid69 70 71def update_conversation_title(cid: str, title: str) -> None:72    """Set conversation title from first user message (truncated)."""73    short = title[:60] + "…" if len(title) > 60 else title74    now = datetime.utcnow().isoformat()75    with _conn() as conn:76        conn.execute(77            "UPDATE conversations SET title=?, updated_at=? WHERE id=?",78            (short, now, cid),79        )80 81 82def touch_conversation(cid: str) -> None:83    now = datetime.utcnow().isoformat()84    with _conn() as conn:85        conn.execute("UPDATE conversations SET updated_at=? WHERE id=?", (now, cid))86 87 88def delete_conversation(cid: str) -> None:89    with _conn() as conn:90        conn.execute("DELETE FROM messages WHERE conversation_id=?", (cid,))91        conn.execute("DELETE FROM conversations WHERE id=?", (cid,))92 93 94def list_conversations(limit: int = 40) -> list[dict]:95    """Return conversations newest-first."""96    with _conn() as conn:97        rows = conn.execute(98            "SELECT id, title, mode, created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ?",99            (limit,),100        ).fetchall()101    return [dict(r) for r in rows]102 103 104# ------------------------------------------------------------------ #105# Messages106# ------------------------------------------------------------------ #107 108def add_message(109    cid: str,110    role: str,111    content: str,112    metadata: Optional[dict] = None,113) -> None:114    """Append a message to a conversation."""115    now = datetime.utcnow().isoformat()116    meta_str = json.dumps(metadata) if metadata else None117    with _conn() as conn:118        conn.execute(119            "INSERT INTO messages (conversation_id, role, content, metadata, created_at) VALUES (?,?,?,?,?)",120            (cid, role, content, meta_str, now),121        )122    touch_conversation(cid)123 124 125def get_messages(cid: str) -> list[dict]:126    """Return all messages for a conversation, oldest-first."""127    with _conn() as conn:128        rows = conn.execute(129            "SELECT role, content, metadata, created_at FROM messages WHERE conversation_id=? ORDER BY id ASC",130            (cid,),131        ).fetchall()132    result = []133    for r in rows:134        msg = dict(r)135        msg["metadata"] = json.loads(msg["metadata"]) if msg["metadata"] else {}136        result.append(msg)137    return result138 139 140def group_conversations_by_date(conversations: list[dict]) -> dict[str, list[dict]]:141    """Group conversation list into Today / Yesterday / Earlier."""142    from datetime import date, timedelta143    today = date.today()144    yesterday = today - timedelta(days=1)145    week_ago = today - timedelta(days=7)146 147    groups: dict[str, list] = {"Today": [], "Yesterday": [], "This Week": [], "Earlier": []}148    for c in conversations:149        try:150            d = datetime.fromisoformat(c["updated_at"]).date()151        except Exception:152            d = today153 154        if d == today:155            groups["Today"].append(c)156        elif d == yesterday:157            groups["Yesterday"].append(c)158        elif d >= week_ago:159            groups["This Week"].append(c)160        else:161            groups["Earlier"].append(c)162 163    return {k: v for k, v in groups.items() if v}   # drop empty groups164 165 166# Initialise on import167init_db()168