File size: 5,118 Bytes
701cf7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """
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}
|