NexusOS/synapse/memory/store.py

from pathlib import Path
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
import sqlite3
import json
import os
import time

# -----------------------------
# Models
# -----------------------------
class MemoryItem(BaseModel):
    id: str
    section: str = "General"
    text: str
    tags: List[str] = []

class MessageItem(BaseModel):
    role: str  # "user" or "assistant"
    content: str
    timestamp: float
    model: Optional[str] = None
    tokens: Optional[int] = None

class ConversationItem(BaseModel):
    id: str
    messages: List[MessageItem] = []
    created_at: float
    updated_at: float

    @property
    def preview(self) -> str:
        for msg in self.messages:
            if msg.role == "user":
                return msg.content[:80]
        return "Empty conversation"

    @property
    def timestamp(self) -> float:
        return self.created_at

# -----------------------------
# Persistent Store
# -----------------------------
class PersistentMemoryStore:
    def __init__(self, db_path: Path):
        self.db_path = db_path
        os.makedirs(self.db_path.parent, exist_ok=True)
        self._ensure_tables()
        self._cache: Dict[str, MemoryItem] = self._load_all_memory()

    # -----------------------------
    # Internal helpers
    # -----------------------------
    def _connect(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL;")
        return conn

    def _ensure_tables(self):
        conn = self._connect()
        cur = conn.cursor()

        cur.execute("""
            CREATE TABLE IF NOT EXISTS memory (
                id TEXT PRIMARY KEY,
                section TEXT NOT NULL DEFAULT 'General',
                text TEXT NOT NULL,
                tags TEXT
            )
        """)
        # Migrate: add section column if it doesn't exist yet
        try:
            cur.execute("ALTER TABLE memory ADD COLUMN section TEXT NOT NULL DEFAULT 'General'")
        except Exception:
            pass

        cur.execute("""
            CREATE TABLE IF NOT EXISTS conversations (
                id TEXT PRIMARY KEY,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL
            )
        """)

        cur.execute("""
            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 REAL NOT NULL,
                model TEXT,
                tokens INTEGER,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)
        for col, typedef in (("model", "TEXT"), ("tokens", "INTEGER")):
            try:
                cur.execute(f"ALTER TABLE messages ADD COLUMN {col} {typedef}")
            except Exception:
                pass

        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_messages_conversation_id
            ON messages (conversation_id)
        """)

        cur.execute("""
            CREATE TABLE IF NOT EXISTS settings (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )
        """)

        conn.commit()
        conn.close()

    # -----------------------------
    # Loaders
    # -----------------------------
    def _load_all_memory(self) -> Dict[str, MemoryItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT id, section, text, tags FROM memory")
        rows = cur.fetchall()
        conn.close()

        cache = {}
        for row in rows:
            try:
                tags = json.loads(row["tags"]) if row["tags"] else []
            except Exception:
                tags = []
            cache[row["id"]] = MemoryItem(
                id=row["id"],
                section=row["section"] or "General",
                text=row["text"],
                tags=tags,
            )

        return cache

    # -----------------------------
    # Memory API
    # -----------------------------
    def add(self, item: MemoryItem):
        self._cache[item.id] = item
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT OR REPLACE INTO memory (id, section, text, tags) VALUES (?, ?, ?, ?)",
                (item.id, item.section or "General", item.text, json.dumps(item.tags))
            )
            conn.commit()
        finally:
            conn.close()

    def update(self, item: MemoryItem):
        self.add(item)

    def delete(self, item_id: str):
        self._cache.pop(item_id, None)
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute("DELETE FROM memory WHERE id = ?", (item_id,))
            conn.commit()
        finally:
            conn.close()

    def get(self, item_id: str) -> Optional[MemoryItem]:
        return self._cache.get(item_id)

    def all(self) -> List[MemoryItem]:
        return list(self._cache.values())

    # -----------------------------
    # Conversation API
    # -----------------------------
    def create_conversation(self, conversation_id: str) -> ConversationItem:
        now = time.time()
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT OR IGNORE INTO conversations (id, created_at, updated_at) VALUES (?, ?, ?)",
                (conversation_id, now, now)
            )
            conn.commit()
        finally:
            conn.close()
        return ConversationItem(id=conversation_id, created_at=now, updated_at=now)

    def add_message(self, conversation_id: str, role: str, content: str, model: str = None, tokens: int = None):
        now = time.time()
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT INTO messages (conversation_id, role, content, timestamp, model, tokens) VALUES (?, ?, ?, ?, ?, ?)",
                (conversation_id, role, content, now, model, tokens)
            )
            cur.execute(
                "UPDATE conversations SET updated_at = ? WHERE id = ?",
                (now, conversation_id)
            )
            conn.commit()
        finally:
            conn.close()

    def get_conversation(self, conversation_id: str) -> Optional[ConversationItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT * FROM conversations WHERE id = ?", (conversation_id,))
        row = cur.fetchone()
        if not row:
            conn.close()
            return None
        cur.execute(
            "SELECT role, content, timestamp, model, tokens FROM messages WHERE conversation_id = ? ORDER BY timestamp ASC",
            (conversation_id,)
        )
        msg_rows = cur.fetchall()
        conn.close()

        messages = [MessageItem(role=r["role"], content=r["content"], timestamp=r["timestamp"], model=r["model"], tokens=r["tokens"]) for r in msg_rows]
        return ConversationItem(
            id=row["id"],
            messages=messages,
            created_at=row["created_at"],
            updated_at=row["updated_at"]
        )

    def all_conversations(self) -> List[ConversationItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("""
            SELECT c.id, c.created_at, c.updated_at,
                   m.role, m.content, m.timestamp, m.model, m.tokens
            FROM conversations c
            LEFT JOIN messages m ON m.conversation_id = c.id
            ORDER BY c.updated_at DESC, m.timestamp ASC
        """)
        rows = cur.fetchall()
        conn.close()

        convs: Dict[str, ConversationItem] = {}
        order: list[str] = []
        for row in rows:
            cid = row["id"]
            if cid not in convs:
                convs[cid] = ConversationItem(
                    id=cid,
                    created_at=row["created_at"],
                    updated_at=row["updated_at"],
                )
                order.append(cid)
            if row["role"] is not None:
                convs[cid].messages.append(
                    MessageItem(role=row["role"], content=row["content"], timestamp=row["timestamp"], model=row["model"], tokens=row["tokens"])
                )
        return [convs[cid] for cid in order]

    def delete_conversation(self, conversation_id: str):
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
            cur.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
            conn.commit()
        finally:
            conn.close()

    # -----------------------------
    # Search API
    # -----------------------------
    def search_conversations(self, query: str, limit: int = 3) -> List[dict]:
        """Return up to `limit` conversations that contain the query string,
        with full user+assistant exchange pairs around each match."""
        if not query or not query.strip():
            return []
        q = query.strip().lower()
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("""
            SELECT DISTINCT c.id, c.created_at, c.updated_at
            FROM conversations c
            JOIN messages m ON m.conversation_id = c.id
            WHERE LOWER(m.content) LIKE ?
            ORDER BY c.updated_at DESC
            LIMIT ?
        """, (f"%{q}%", limit))
        rows = cur.fetchall()

        results = []
        for row in rows:
            # Load all messages in order so we can find complete exchange pairs
            cur.execute("""
                SELECT role, content FROM messages
                WHERE conversation_id = ?
                ORDER BY timestamp ASC
            """, (row["id"],))
            all_msgs = [{"role": r["role"], "content": r["content"]} for r in cur.fetchall()]

            # For each matching message, collect the full user+assistant pair around it
            seen_pairs: set = set()
            matches = []
            for i, msg in enumerate(all_msgs):
                if q not in msg["content"].lower():
                    continue
                if msg["role"] == "user":
                    start, end = i, i + 1 if i + 1 < len(all_msgs) else i
                else:
                    start, end = (i - 1 if i > 0 else i), i
                if (start, end) in seen_pairs:
                    continue
                seen_pairs.add((start, end))
                for m in all_msgs[start:end + 1]:
                    matches.append({"role": m["role"], "content": m["content"][:500]})
                if len(seen_pairs) >= 2:
                    break

            if matches:
                results.append({
                    "id": row["id"],
                    "updated_at": row["updated_at"],
                    "matches": matches,
                })

        conn.close()
        return results

    # -----------------------------
    # Settings API
    # -----------------------------
    _SETTINGS_DEFAULTS: Dict[str, Any] = {
        "model": "",
        "temperature": 0.7,
        "system_prompt": "",
        "timeout": 120,
    }

    def get_settings(self) -> Dict[str, Any]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT key, value FROM settings")
        rows = cur.fetchall()
        conn.close()
        result = dict(self._SETTINGS_DEFAULTS)
        for row in rows:
            try:
                result[row["key"]] = json.loads(row["value"])
            except Exception:
                result[row["key"]] = row["value"]
        return result

    def update_settings(self, data: Dict[str, Any]):
        conn = self._connect()
        try:
            cur = conn.cursor()
            for key, value in data.items():
                if key in self._SETTINGS_DEFAULTS:
                    cur.execute(
                        "INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)",
                        (key, json.dumps(value))
                    )
            conn.commit()
        finally:
            conn.close()

# -----------------------------
# Store Instance
# -----------------------------
from ..nexus_config import MEMORY_DB

DB_PATH = MEMORY_DB
store = PersistentMemoryStore(DB_PATH)