371 lines
14 KiB
Python
371 lines
14 KiB
Python
"""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
|