memory-os/setup/setup_db.py

315 lines
10 KiB
Python

#!/usr/bin/env python3
"""Memory OS — Database Setup Script
Creates the SQLite databases and schemas required by the Icarus plugin.
Run this once during initial installation, or after deleting the databases.
Usage:
python setup_db.py # uses default paths
python setup_db.py --state-db /custom/path/state.db
python setup_db.py --memory-db /custom/path/memory_store.db
python setup_db.py --dry-run # print SQL without executing
Environment variables:
STATE_DB_PATH — path to state.db (default: ~/.hermes/state.db)
MEMORY_STORE_PATH — path to memory_store.db (default: ~/.hermes/memory_store.db)
"""
import argparse
import os
import sqlite3
import sys
from pathlib import Path
# ── Default paths ──────────────────────────────────────────────────────────
DEFAULT_HERMES_HOME = Path.home() / ".hermes"
DEFAULT_STATE_DB = DEFAULT_HERMES_HOME / "state.db"
DEFAULT_MEMORY_DB = DEFAULT_HERMES_HOME / "memory_store.db"
# ── Schema: state.db ───────────────────────────────────────────────────────
STATE_SCHEMA = """
-- Sessions table — one row per Hermes session
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
source TEXT NOT NULL,
user_id TEXT,
model TEXT,
model_config TEXT,
system_prompt TEXT,
parent_session_id TEXT,
started_at REAL NOT NULL,
ended_at REAL,
end_reason TEXT,
message_count INTEGER DEFAULT 0,
tool_call_count INTEGER DEFAULT 0,
input_tokens INTEGER DEFAULT 0,
output_tokens INTEGER DEFAULT 0,
cache_read_tokens INTEGER DEFAULT 0,
cache_write_tokens INTEGER DEFAULT 0,
reasoning_tokens INTEGER DEFAULT 0,
billing_provider TEXT,
billing_base_url TEXT,
billing_mode TEXT,
estimated_cost_usd REAL,
actual_cost_usd REAL,
cost_status TEXT,
cost_source TEXT,
pricing_version TEXT,
title TEXT,
api_call_count INTEGER DEFAULT 0,
handoff_state TEXT,
handoff_platform TEXT,
handoff_error TEXT,
cwd TEXT,
rewind_count INTEGER NOT NULL DEFAULT 0,
archived INTEGER NOT NULL DEFAULT 0,
FOREIGN KEY (parent_session_id) REFERENCES sessions(id)
);
-- Messages table — one row per user/assistant/tool message
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL REFERENCES sessions(id),
role TEXT NOT NULL,
content TEXT,
tool_call_id TEXT,
tool_calls TEXT,
tool_name TEXT,
timestamp REAL NOT NULL,
token_count INTEGER,
finish_reason TEXT,
reasoning TEXT,
reasoning_content TEXT,
reasoning_details TEXT,
codex_reasoning_items TEXT,
codex_message_items TEXT,
platform_message_id TEXT,
observed INTEGER DEFAULT 0,
active INTEGER NOT NULL DEFAULT 1
);
-- FTS5 full-text search over message content (used by session_search)
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
content,
content=messages,
content_rowid=id
);
-- Trigram FTS5 index for substring/phrase matching
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts_trigram USING fts5(
content,
tokenize='trigram',
content=messages,
content_rowid=id
);
-- Lineage tracking — generation provenance for context enhancement
CREATE TABLE IF NOT EXISTS lineage (
lineage_id TEXT PRIMARY KEY,
session_id TEXT,
query TEXT,
retrieved_chunk_ids TEXT,
generation_model TEXT,
generation_context_hash TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
-- Reflection budget — rate-limiting for micro-reflections
CREATE TABLE IF NOT EXISTS reflection_budget (
hour_window TEXT PRIMARY KEY,
count INTEGER DEFAULT 0,
tokens_used INTEGER DEFAULT 0
);
-- Compression locks — concurrency control for session compression
CREATE TABLE IF NOT EXISTS compression_locks (
session_id TEXT PRIMARY KEY,
holder TEXT NOT NULL,
acquired_at REAL NOT NULL,
expires_at REAL NOT NULL
);
-- Schema version tracking
CREATE TABLE IF NOT EXISTS schema_version (
version INTEGER NOT NULL
);
-- State metadata key-value store
CREATE TABLE IF NOT EXISTS state_meta (
key TEXT PRIMARY KEY,
value TEXT
);
"""
# ── Schema: memory_store.db ────────────────────────────────────────────────
MEMORY_SCHEMA = """
-- Entities — named entities extracted from facts
CREATE TABLE IF NOT EXISTS entities (
entity_id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
entity_type TEXT DEFAULT 'unknown',
aliases TEXT DEFAULT '',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
-- Facts — durable structured memory with trust scoring
CREATE TABLE IF NOT EXISTS facts (
fact_id INTEGER PRIMARY KEY AUTOINCREMENT,
content TEXT NOT NULL UNIQUE,
category TEXT DEFAULT 'general',
tags TEXT DEFAULT '',
trust_score REAL DEFAULT 0.5,
retrieval_count INTEGER DEFAULT 0,
helpful_count INTEGER DEFAULT 0,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
hrr_vector BLOB
);
-- Fact-entity junction (many-to-many)
CREATE TABLE IF NOT EXISTS fact_entities (
fact_id INTEGER REFERENCES facts(fact_id),
entity_id INTEGER REFERENCES entities(entity_id),
PRIMARY KEY (fact_id, entity_id)
);
-- FTS5 full-text search over facts
CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts
USING fts5(content, tags, content=facts, content_rowid=fact_id);
-- Memory banks — algebraically composable fact groups (HRR vectors)
CREATE TABLE IF NOT EXISTS memory_banks (
bank_id INTEGER PRIMARY KEY AUTOINCREMENT,
bank_name TEXT NOT NULL UNIQUE,
vector BLOB NOT NULL,
dim INTEGER NOT NULL,
fact_count INTEGER DEFAULT 0,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
"""
def create_database(db_path: Path, schema: str, dry_run: bool = False) -> bool:
"""Create database and execute schema. Returns True on success."""
print(f"\n📘 Database: {db_path}")
if dry_run:
print(" [DRY-RUN] Would execute schema:")
for line in schema.strip().split("\n"):
if line.strip():
print(f" {line.strip()[:100]}")
return True
db_path.parent.mkdir(parents=True, exist_ok=True)
try:
conn = sqlite3.connect(str(db_path))
conn.executescript(schema)
conn.commit()
conn.close()
print(f" ✅ Created ({db_path.stat().st_size:,} bytes)")
# Verify tables
conn = sqlite3.connect(str(db_path))
tables = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
).fetchall()
conn.close()
print(f" 📋 Tables: {', '.join(t[0] for t in tables)}")
return True
except sqlite3.Error as e:
print(f" ❌ Error: {e}", file=sys.stderr)
return False
def insert_schema_version(db_path: Path, version: int = 1) -> None:
"""Record the schema version if the table exists and is empty."""
try:
conn = sqlite3.connect(str(db_path))
existing = conn.execute(
"SELECT COUNT(*) FROM schema_version"
).fetchone()[0]
if existing == 0:
conn.execute(
"INSERT INTO schema_version (version) VALUES (?)", (version,)
)
conn.commit()
conn.close()
except sqlite3.Error:
pass # table might not exist (not state.db)
def main():
parser = argparse.ArgumentParser(
description="Memory OS — Database Setup"
)
parser.add_argument(
"--state-db",
default=os.environ.get(
"STATE_DB_PATH", str(DEFAULT_STATE_DB)
),
help=f"Path to state.db (default: {DEFAULT_STATE_DB})",
)
parser.add_argument(
"--memory-db",
default=os.environ.get(
"MEMORY_STORE_PATH", str(DEFAULT_MEMORY_DB)
),
help=f"Path to memory_store.db (default: {DEFAULT_MEMORY_DB})",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Print SQL without executing",
)
args = parser.parse_args()
state_path = Path(args.state_db)
memory_path = Path(args.memory_db)
print("🧠 Memory OS — Database Setup")
print(f" Hermes home: {DEFAULT_HERMES_HOME}")
print(f" State DB: {state_path}")
print(f" Memory DB: {memory_path}")
if args.dry_run:
print("\n⚠️ DRY-RUN mode — no changes will be made.\n")
# ── Create state.db ──
ok_state = create_database(state_path, STATE_SCHEMA, args.dry_run)
if ok_state and not args.dry_run:
insert_schema_version(state_path, version=1)
# ── Create memory_store.db ──
ok_memory = create_database(memory_path, MEMORY_SCHEMA, args.dry_run)
# ── Summary ──
print()
if args.dry_run:
print("🏁 Dry-run complete. Run without --dry-run to apply.")
elif ok_state and ok_memory:
print("✅ Both databases created successfully.")
print()
print(" Next steps:")
print(f" 1. Ensure Qdrant is running: docker compose up -d qdrant")
print(f" 2. Create the Qdrant collection:")
print(f" python scripts/bulk_wiki_ingest.py (creates collection automatically)")
print(f" OR manually:")
print(f" curl -X PUT http://localhost:6333/collections/knowledge_base \\")
print(f" -H 'Content-Type: application/json' \\")
print(f" -d '{{\"vectors\":{{\"dense\":{{\"size\":4096,\"distance\":\"Cosine\"}}}},\"sparse_vectors\":{{\"sparse\":{{}}}}}}'")
sys.exit(0)
else:
print("❌ Some databases failed to create. Check errors above.")
sys.exit(1)
if __name__ == "__main__":
main()