src/ares/memory/sqlite_store.py

"""
SQLite petabyte-scale capable database (architected for scale, works locally).
Stores raw documents, metadata, embeddings pointer, cross-references.
"""
import sqlite3
import os
import json
import time
from typing import List, Dict, Optional
import numpy as np

class SQLiteStore:
    def __init__(self, db_path: str = "data/ares_knowledge.db"):
        os.makedirs(os.path.dirname(db_path) if os.path.dirname(db_path) else ".", exist_ok=True)
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.execute("PRAGMA journal_mode=WAL;")
        self.conn.execute("PRAGMA synchronous=NORMAL;")
        self._init_tables()

    def _init_tables(self):
        self.conn.execute("""
        CREATE TABLE IF NOT EXISTS documents (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            content TEXT NOT NULL,
            source TEXT,
            metadata TEXT,
            timestamp REAL,
            embedding_id INTEGER
        );
        """)
        self.conn.execute("""
        CREATE TABLE IF NOT EXISTS embeddings (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            doc_id INTEGER,
            vector BLOB,
            FOREIGN KEY(doc_id) REFERENCES documents(id)
        );
        """)
        self.conn.execute("""
        CREATE TABLE IF NOT EXISTS cross_refs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            doc_id INTEGER,
            ref_doc_id INTEGER,
            score REAL,
            relation TEXT,
            FOREIGN KEY(doc_id) REFERENCES documents(id),
            FOREIGN KEY(ref_doc_id) REFERENCES documents(id)
        );
        """)
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_docs_source ON documents(source);")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_emb_doc ON embeddings(doc_id);")
        self.conn.commit()

    def add_document(self, content: str, source: str = "unknown", metadata: Dict = None) -> int:
        cur = self.conn.cursor()
        cur.execute("INSERT INTO documents (content, source, metadata, timestamp) VALUES (?,?,?,?)",
                    (content, source, json.dumps(metadata or {}), time.time()))
        doc_id = cur.lastrowid
        self.conn.commit()
        return doc_id

    def add_documents_batch(self, docs: List[Dict]) -> List[int]:
        ids = []
        cur = self.conn.cursor()
        for d in docs:
            cur.execute("INSERT INTO documents (content, source, metadata, timestamp) VALUES (?,?,?,?)",
                        (d["content"], d.get("source","unknown"), json.dumps(d.get("metadata",{})), time.time()))
            ids.append(cur.lastrowid)
        self.conn.commit()
        return ids

    def add_embedding(self, doc_id: int, vector: np.ndarray) -> int:
        blob = vector.astype(np.float32).tobytes()
        cur = self.conn.cursor()
        cur.execute("INSERT INTO embeddings (doc_id, vector) VALUES (?,?)", (doc_id, blob))
        emb_id = cur.lastrowid
        # update doc pointer
        cur.execute("UPDATE documents SET embedding_id=? WHERE id=?", (emb_id, doc_id))
        self.conn.commit()
        return emb_id

    def get_all_embeddings(self):
        cur = self.conn.cursor()
        cur.execute("SELECT id, doc_id, vector FROM embeddings")
        rows = cur.fetchall()
        result = []
        for emb_id, doc_id, blob in rows:
            vec = np.frombuffer(blob, dtype=np.float32)
            result.append((emb_id, doc_id, vec))
        return result

    def get_document(self, doc_id: int) -> Optional[Dict]:
        cur = self.conn.cursor()
        cur.execute("SELECT id, content, source, metadata, timestamp FROM documents WHERE id=?", (doc_id,))
        row = cur.fetchone()
        if row:
            return {"id": row[0], "content": row[1], "source": row[2], "metadata": json.loads(row[3]), "timestamp": row[4]}
        return None

    def search_content(self, query: str, limit=10) -> List[Dict]:
        # Simple FTS fallback without vector
        cur = self.conn.cursor()
        cur.execute("SELECT id, content, source FROM documents WHERE content LIKE ? LIMIT ?", (f"%{query}%", limit))
        rows = cur.fetchall()
        return [{"id": r[0], "content": r[1], "source": r[2], "score": 1.0} for r in rows]

    def add_cross_ref(self, doc_id: int, ref_doc_id: int, score: float, relation="related"):
        cur = self.conn.cursor()
        cur.execute("INSERT INTO cross_refs (doc_id, ref_doc_id, score, relation) VALUES (?,?,?,?)",
                    (doc_id, ref_doc_id, score, relation))
        self.conn.commit()

    def stats(self):
        cur = self.conn.cursor()
        cur.execute("SELECT COUNT(*) FROM documents")
        doc_count = cur.fetchone()[0]
        cur.execute("SELECT COUNT(*) FROM embeddings")
        emb_count = cur.fetchone()[0]
        cur.execute("SELECT COUNT(*) FROM cross_refs")
        ref_count = cur.fetchone()[0]
        size = os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0
        return {"documents": doc_count, "embeddings": emb_count, "cross_refs": ref_count, "db_size_bytes": size}