"""SQLite-backed conversation store with async support and in-memory LRU cache.""" from __future__ import annotations import json import logging import os import uuid from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any import aiosqlite logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Data classes (mirrors portal.py ChatMessage / Conversation) # --------------------------------------------------------------------------- @dataclass class ChatMessage: role: str # "user" or "assistant" content: str timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) metadata: dict = field(default_factory=dict) @dataclass class Conversation: id: str messages: list[ChatMessage] = field(default_factory=list) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) # --------------------------------------------------------------------------- # Schema # --------------------------------------------------------------------------- _SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS conversations ( id TEXT PRIMARY KEY, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, conversation_id TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, timestamp TEXT NOT NULL, metadata TEXT DEFAULT '{}', FOREIGN KEY (conversation_id) REFERENCES conversations(id) ); CREATE INDEX IF NOT EXISTS idx_messages_conv_id ON messages(conversation_id); """ class SqliteConversationStore: """SQLite-backed conversation store with an in-memory LRU cache. Drop-in replacement for the in-memory ConversationStore in portal.py. Data is persisted to SQLite so it survives server restarts. An in-memory LRU cache of recent conversations provides fast access. """ def __init__( self, db_path: str | Path | None = None, max_conversations: int = 1000, ) -> None: if db_path is None: db_path = os.path.expanduser("~/.agentkit/conversations.db") self._db_path = str(db_path) self._max = max_conversations self._cache: OrderedDict[str, Conversation] = OrderedDict() self._db: aiosqlite.Connection | None = None # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ async def _ensure_db(self) -> aiosqlite.Connection: """Lazily open the database and ensure schema exists.""" if self._db is None: db_dir = os.path.dirname(self._db_path) if db_dir: os.makedirs(db_dir, exist_ok=True) self._db = await aiosqlite.connect(self._db_path) self._db.row_factory = aiosqlite.Row await self._db.execute("PRAGMA journal_mode=WAL") await self._db.executescript(_SCHEMA_SQL) await self._db.commit() return self._db async def _close_db(self) -> None: """Close the database connection.""" if self._db is not None: await self._db.close() self._db = None @staticmethod def _dt_to_str(dt: datetime) -> str: return dt.isoformat() @staticmethod def _str_to_dt(s: str) -> datetime: return datetime.fromisoformat(s) def _touch_cache(self, conv_id: str) -> None: """Move conversation to the end of the LRU cache (most recently used).""" if conv_id in self._cache: self._cache.move_to_end(conv_id) def _evict_if_needed(self) -> None: """Evict oldest entry from cache if over limit (does NOT delete from SQLite).""" while len(self._cache) > self._max: self._cache.popitem(last=False) # ------------------------------------------------------------------ # Public API (matches ConversationStore interface) # ------------------------------------------------------------------ async def get_or_create(self, conversation_id: str | None = None) -> Conversation: """Get existing conversation or create a new one. If the conversation is in cache, return it directly. Otherwise, try to load from SQLite. If not found, create new. Messages are loaded lazily (only when get_history is called). """ db = await self._ensure_db() if conversation_id and conversation_id in self._cache: conv = self._cache[conversation_id] conv.updated_at = datetime.now(timezone.utc) self._touch_cache(conv.id) return conv # Try loading from SQLite if conversation_id: cursor = await db.execute( "SELECT id, created_at, updated_at FROM conversations WHERE id = ?", (conversation_id,), ) row = await cursor.fetchone() if row: conv = Conversation( id=row["id"], created_at=self._str_to_dt(row["created_at"]), updated_at=self._str_to_dt(row["updated_at"]), ) self._cache[conv.id] = conv self._evict_if_needed() return conv # Create new cid = conversation_id or str(uuid.uuid4()) now = datetime.now(timezone.utc) conv = Conversation(id=cid, created_at=now, updated_at=now) await db.execute( "INSERT INTO conversations (id, created_at, updated_at) VALUES (?, ?, ?)", (conv.id, self._dt_to_str(conv.created_at), self._dt_to_str(conv.updated_at)), ) await db.commit() self._cache[conv.id] = conv self._evict_if_needed() return conv async def add_message( self, conversation_id: str, role: str, content: str, metadata: dict | None = None, ) -> ChatMessage: """Add a message to a conversation (both in-memory cache and SQLite).""" db = await self._ensure_db() # Ensure conversation exists in cache if conversation_id not in self._cache: # Try loading from SQLite cursor = await db.execute( "SELECT id, created_at, updated_at FROM conversations WHERE id = ?", (conversation_id,), ) row = await cursor.fetchone() if row: conv = Conversation( id=row["id"], created_at=self._str_to_dt(row["created_at"]), updated_at=self._str_to_dt(row["updated_at"]), ) self._cache[conv.id] = conv else: raise KeyError(f"Conversation '{conversation_id}' not found") conv = self._cache[conversation_id] msg = ChatMessage( role=role, content=content, metadata=metadata or {}, ) conv.messages.append(msg) conv.updated_at = datetime.now(timezone.utc) self._touch_cache(conv.id) # Persist to SQLite await db.execute( "INSERT INTO messages (conversation_id, role, content, timestamp, metadata) " "VALUES (?, ?, ?, ?, ?)", ( conversation_id, role, content, self._dt_to_str(msg.timestamp), json.dumps(msg.metadata, default=str), ), ) await db.execute( "UPDATE conversations SET updated_at = ? WHERE id = ?", (self._dt_to_str(conv.updated_at), conversation_id), ) await db.commit() return msg async def get_history(self, conversation_id: str, limit: int = 50) -> list[ChatMessage]: """Get recent messages for a conversation. Loads from SQLite to ensure completeness (in-memory cache may only have a subset of messages after a restart). """ db = await self._ensure_db() cursor = await db.execute( "SELECT role, content, timestamp, metadata FROM messages " "WHERE conversation_id = ? ORDER BY id DESC LIMIT ?", (conversation_id, limit), ) rows = await cursor.fetchall() # Reverse because we fetched DESC but want chronological order messages: list[ChatMessage] = [] for row in reversed(rows): meta: dict[str, Any] = {} try: meta = json.loads(row["metadata"]) if row["metadata"] else {} except (json.JSONDecodeError, TypeError): pass messages.append( ChatMessage( role=row["role"], content=row["content"], timestamp=self._str_to_dt(row["timestamp"]), metadata=meta, ) ) return messages async def list_conversations(self, limit: int = 20) -> list[Conversation]: """List recent conversations ordered by updated_at (most recent first).""" db = await self._ensure_db() cursor = await db.execute( "SELECT id, created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ?", (limit,), ) rows = await cursor.fetchall() result: list[Conversation] = [] for row in rows: conv_id = row["id"] # Use cached version if available (may have in-memory messages) if conv_id in self._cache: result.append(self._cache[conv_id]) else: conv = Conversation( id=conv_id, created_at=self._str_to_dt(row["created_at"]), updated_at=self._str_to_dt(row["updated_at"]), ) # Load first user message from SQLite for title derivation try: title_cursor = await db.execute( "SELECT content FROM messages " "WHERE conversation_id = ? AND role = 'user' " "ORDER BY timestamp ASC LIMIT 1", (conv_id,), ) title_row = await title_cursor.fetchone() if title_row: conv.messages.append(ChatMessage(role="user", content=title_row["content"])) except Exception: pass result.append(conv) return result async def delete_conversation(self, conversation_id: str) -> bool: """Delete a conversation and all its messages. Returns ``True`` if the conversation existed and was deleted, ``False`` if it was not found. """ db = await self._ensure_db() cursor = await db.execute( "SELECT id FROM conversations WHERE id = ?", (conversation_id,) ) row = await cursor.fetchone() if row is None: return False # ponytail: wrap both DELETEs in a try/except with rollback — # previously the second DELETE failure would leave orphaned # messages (conversation row gone, messages lingering) because # the first DELETE already auto-committed in autocommit mode. # aiosqlite uses autocommit=False by default but explicit rollback # makes the failure path safe and observable. try: await db.execute( "DELETE FROM messages WHERE conversation_id = ?", (conversation_id,) ) await db.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,)) await db.commit() except Exception: await db.rollback() logger.exception("Failed to delete conversation %s; rolled back", conversation_id) raise self._cache.pop(conversation_id, None) return True async def restore_from_store( self, max_sessions: int = 50, max_messages_per_session: int = 100, ) -> None: """No-op for SQLite store — data is already persisted in the database.""" # Nothing to do; all data lives in SQLite and is loaded on demand. async def get_first_user_message(self, conversation_id: str) -> ChatMessage | None: """Get the first user message of a conversation (authoritative source for title). Reads directly from SQLite so the title is correct even when the in-memory cache has not been populated or has an empty `messages` list (e.g. after a server restart). """ db = await self._ensure_db() try: cursor = await db.execute( "SELECT role, content, timestamp, metadata FROM messages " "WHERE conversation_id = ? AND role = 'user' AND content != '' " "ORDER BY id ASC LIMIT 1", (conversation_id,), ) row = await cursor.fetchone() if not row: return None meta: dict[str, Any] = {} try: meta = json.loads(row["metadata"]) if row["metadata"] else {} except (json.JSONDecodeError, TypeError): pass return ChatMessage( role=row["role"], content=row["content"], timestamp=self._str_to_dt(row["timestamp"]), metadata=meta, ) except Exception as e: logger.warning(f"get_first_user_message failed for {conversation_id}: {e}") return None