"""SQLite chat session storage.""" import json import sqlite3 import uuid from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, List, Optional from config import settings class ChatStore: def __init__(self, db_path: Optional[str] = None): self.db_path = db_path or settings.CHAT_DB_PATH Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) def _connect(self) -> sqlite3.Connection: conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row return conn def _schema_sql(self) -> tuple[str, str, str]: sessions = """ CREATE TABLE IF NOT EXISTS sessions ( id TEXT PRIMARY KEY, title TEXT NOT NULL, created_at TEXT NOT NULL ) """ messages = """ CREATE TABLE IF NOT EXISTS messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, sources_json TEXT, confidence REAL, created_at TEXT NOT NULL, FOREIGN KEY (session_id) REFERENCES sessions(id) ) """ session_documents = """ CREATE TABLE IF NOT EXISTS session_documents ( session_id TEXT NOT NULL, document_id TEXT NOT NULL, PRIMARY KEY (session_id, document_id), FOREIGN KEY (session_id) REFERENCES sessions(id) ) """ return sessions, messages, session_documents def init_db(self) -> None: sessions, messages, session_documents = self._schema_sql() with sqlite3.connect(self.db_path) as db: db.execute(sessions) db.execute(messages) db.execute(session_documents) db.commit() def create_session( self, title: str = "New Chat", document_ids: Optional[List[str]] = None ) -> Dict[str, Any]: session_id = str(uuid.uuid4()) now = datetime.now(timezone.utc).isoformat() with sqlite3.connect(self.db_path) as db: db.execute( "INSERT INTO sessions (id, title, created_at) VALUES (?, ?, ?)", (session_id, title, now), ) if document_ids: for doc_id in document_ids: db.execute( "INSERT INTO session_documents (session_id, document_id) VALUES (?, ?)", (session_id, doc_id), ) db.commit() return {"id": session_id, "title": title, "created_at": now} def list_sessions(self) -> List[Dict[str, Any]]: with self._connect() as db: cursor = db.execute( "SELECT id, title, created_at FROM sessions ORDER BY created_at DESC" ) return [dict(row) for row in cursor.fetchall()] def get_messages(self, session_id: str) -> List[Dict[str, Any]]: with self._connect() as db: cursor = db.execute( """ SELECT id, role, content, sources_json, confidence, created_at FROM messages WHERE session_id = ? ORDER BY created_at ASC """, (session_id,), ) result = [] for row in cursor.fetchall(): msg = dict(row) if msg.get("sources_json"): msg["sources"] = json.loads(msg["sources_json"]) else: msg["sources"] = [] del msg["sources_json"] result.append(msg) return result def add_message( self, session_id: str, role: str, content: str, sources: Optional[List[Dict]] = None, confidence: Optional[float] = None, ) -> Dict[str, Any]: msg_id = str(uuid.uuid4()) now = datetime.now(timezone.utc).isoformat() sources_json = json.dumps(sources or []) with sqlite3.connect(self.db_path) as db: db.execute( """ INSERT INTO messages (id, session_id, role, content, sources_json, confidence, created_at) VALUES (?, ?, ?, ?, ?, ?, ?) """, (msg_id, session_id, role, content, sources_json, confidence, now), ) db.commit() return { "id": msg_id, "session_id": session_id, "role": role, "content": content, "sources": sources or [], "confidence": confidence, "created_at": now, } def export_session(self, session_id: str) -> str: with self._connect() as db: cursor = db.execute( "SELECT title, created_at FROM sessions WHERE id = ?", (session_id,), ) session = cursor.fetchone() if not session: return "" messages = self.get_messages(session_id) lines = [ "FinSight AI Chat Export", f"Session: {session['title']}", f"Created: {session['created_at']}", "=" * 60, "", ] for msg in messages: role_label = "You" if msg["role"] == "user" else "FinSight AI" lines.append(f"## {role_label} ({msg['created_at']})") lines.append(msg["content"]) if msg.get("confidence"): lines.append(f"\n[Confidence: {msg['confidence']}/10]") if msg.get("sources"): lines.append("\nSources:") for i, src in enumerate(msg["sources"], 1): lines.append( f" {i}. {src.get('document_name', '')} p.{src.get('page_number', '?')} " f"({src.get('section', '')})" ) lines.append("") return "\n".join(lines)