nexo-brain 2.0.0 → 2.2.0

package/src/migrate_embeddings 2.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""
+Migrate cognitive.db embeddings between models.
+
+Usage:
+  python migrate_embeddings.py upgrade   # 384 → 768 (bge-small → bge-base)
+  python migrate_embeddings.py rollback  # Restore from backup
+  python migrate_embeddings.py verify    # Check current embedding dims
+"""
+
+import os
+import shutil
+import sqlite3
+import sys
+import time
+import numpy as np
+
+NEXO_HOME = os.environ.get("NEXO_HOME", os.path.expanduser("~/.nexo"))
+_data_dir = os.path.join(NEXO_HOME, "data")
+os.makedirs(_data_dir, exist_ok=True)
+DB_PATH = os.path.join(_data_dir, "cognitive.db")
+BACKUP_PATH = DB_PATH + ".bak-384dims-pre-upgrade"
+
+MODELS = {
+    "small": ("BAAI/bge-small-en-v1.5", 384),
+    "base": ("BAAI/bge-base-en-v1.5", 768),
+}
+
+
+def verify():
+    """Check current embedding dimensions in the database."""
+    conn = sqlite3.connect(DB_PATH)
+    for table in ["stm_memories", "ltm_memories"]:
+        count = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
+        if count == 0:
+            print(f"  {table}: {count} rows (empty)")
+            continue
+        row = conn.execute(f"SELECT embedding FROM {table} LIMIT 1").fetchone()
+        vec = np.frombuffer(row[0], dtype=np.float32)
+        print(f"  {table}: {count} rows, embedding dim = {len(vec)}")
+    conn.close()
+
+
+def upgrade():
+    """Re-embed all memories from bge-small (384) to bge-base (768)."""
+    from fastembed import TextEmbedding
+
+    # Verify current state
+    print("Current state:")
+    verify()
+
+    # Verify backup exists
+    if not os.path.exists(BACKUP_PATH):
+        print(f"\nCreating backup at {BACKUP_PATH}")
+        shutil.copy2(DB_PATH, BACKUP_PATH)
+    else:
+        print(f"\nBackup already exists at {BACKUP_PATH}")
+
+    # Load new model
+    model_name, expected_dim = MODELS["base"]
+    print(f"\nLoading {model_name}...")
+    model = TextEmbedding(model_name)
+
+    conn = sqlite3.connect(DB_PATH)
+
+    for table in ["stm_memories", "ltm_memories"]:
+        rows = conn.execute(f"SELECT id, content FROM {table}").fetchall()
+        if not rows:
+            print(f"\n{table}: empty, skipping")
+            continue
+
+        print(f"\n{table}: re-embedding {len(rows)} memories...")
+        t0 = time.time()
+
+        # Batch embed for speed
+        contents = [r[1] for r in rows]
+        ids = [r[0] for r in rows]
+
+        embeddings = list(model.embed(contents))
+
+        for mem_id, emb in zip(ids, embeddings):
+            blob = np.array(emb, dtype=np.float32).tobytes()
+            conn.execute(f"UPDATE {table} SET embedding = ? WHERE id = ?", (blob, mem_id))
+
+        conn.commit()
+        elapsed = time.time() - t0
+        print(f"  Done: {len(rows)} memories in {elapsed:.1f}s ({elapsed/len(rows)*1000:.0f}ms/memory)")
+
+    conn.close()
+
+    print("\nAfter upgrade:")
+    verify()
+    print("\nUpgrade complete. Run 'verify' to confirm.")
+
+
+def rollback():
+    """Restore database from pre-upgrade backup."""
+    if not os.path.exists(BACKUP_PATH):
+        print(f"ERROR: Backup not found at {BACKUP_PATH}")
+        sys.exit(1)
+
+    print(f"Restoring from {BACKUP_PATH}...")
+    shutil.copy2(BACKUP_PATH, DB_PATH)
+    print("Restored. Current state:")
+    verify()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python migrate_embeddings.py [upgrade|rollback|verify]")
+        sys.exit(1)
+
+    cmd = sys.argv[1]
+    if cmd == "upgrade":
+        upgrade()
+    elif cmd == "rollback":
+        rollback()
+    elif cmd == "verify":
+        verify()
+    else:
+        print(f"Unknown command: {cmd}")
+        sys.exit(1)

package/src/plugin_loader 2.py

@@ -0,0 +1,202 @@
+"""Dynamic plugin loader for NEXO MCP server."""
+
+import importlib
+import importlib.util
+import os
+import signal
+import sys
+import time
+
+from db import get_db
+from fastmcp.tools import Tool
+
+SERVER_DIR = os.path.dirname(os.path.abspath(__file__))
+PLUGINS_DIR = os.path.join(SERVER_DIR, "plugins")
+
+# Personal plugins directory: NEXO_HOME/plugins/ (env var, defaults to ~/.nexo/)
+NEXO_HOME = os.environ.get("NEXO_HOME", os.path.expanduser("~/.nexo"))
+PERSONAL_PLUGINS_DIR = os.path.join(NEXO_HOME, "plugins")
+
+PLUGIN_LOAD_TIMEOUT = 10  # seconds per plugin
+
+
+class _PluginTimeout(Exception):
+    pass
+
+
+def _timeout_handler(signum, frame):
+    raise _PluginTimeout("Plugin loading timed out")
+
+
+def _ensure_src_in_path():
+    """Ensure server src/ is in sys.path so personal plugins can import db, cognitive, etc."""
+    if SERVER_DIR not in sys.path:
+        sys.path.insert(0, SERVER_DIR)
+
+
+def load_all_plugins(mcp) -> int:
+    """Load all plugins from repo and personal directories at startup. Returns total tools loaded."""
+    _ensure_src_in_path()
+    total = 0
+
+    # Collect plugins: repo first, personal overrides
+    plugin_map = {}  # filename -> (dir_path, source_label)
+
+    # 1. Repo plugins (base)
+    if os.path.isdir(PLUGINS_DIR):
+        for f in sorted(os.listdir(PLUGINS_DIR)):
+            if f.endswith(".py") and f != "__init__.py":
+                plugin_map[f] = (PLUGINS_DIR, "repo")
+
+    # 2. Personal plugins (override if same filename)
+    if os.path.isdir(PERSONAL_PLUGINS_DIR):
+        for f in sorted(os.listdir(PERSONAL_PLUGINS_DIR)):
+            if f.endswith(".py") and f != "__init__.py":
+                source = "personal (override)" if f in plugin_map else "personal"
+                plugin_map[f] = (PERSONAL_PLUGINS_DIR, source)
+
+    # Load all in sorted order
+    for f in sorted(plugin_map):
+        plugins_dir, source_label = plugin_map[f]
+        try:
+            old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+            signal.alarm(PLUGIN_LOAD_TIMEOUT)
+            try:
+                n = load_plugin(mcp, f, plugins_dir=plugins_dir)
+                total += n
+                print(f"[PLUGIN LOADED] {f} ({n} tools) from {source_label}: {plugins_dir}", file=sys.stderr)
+            finally:
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old_handler)
+        except _PluginTimeout:
+            print(f"[PLUGIN TIMEOUT] {f}: skipped after {PLUGIN_LOAD_TIMEOUT}s", file=sys.stderr)
+        except Exception as e:
+            print(f"[PLUGIN ERROR] {f}: {e}", file=sys.stderr)
+    return total
+
+
+def load_plugin(mcp, filename: str, plugins_dir: str | None = None) -> int:
+    """Load or reload a single plugin. Returns number of tools registered.
+
+    Args:
+        plugins_dir: Directory to load from. If None, searches repo PLUGINS_DIR first,
+                     then PERSONAL_PLUGINS_DIR. Personal plugins are loaded via
+                     importlib.util.spec_from_file_location.
+    """
+    if not filename.endswith(".py"):
+        filename += ".py"
+
+    if plugins_dir is not None:
+        filepath = os.path.join(plugins_dir, filename)
+        if not os.path.isfile(filepath):
+            raise FileNotFoundError(f"Plugin not found: {filepath}")
+    else:
+        # Search repo first, then personal
+        repo_path = os.path.join(PLUGINS_DIR, filename)
+        personal_path = os.path.join(PERSONAL_PLUGINS_DIR, filename)
+        if os.path.isfile(repo_path):
+            plugins_dir = PLUGINS_DIR
+            filepath = repo_path
+        elif os.path.isfile(personal_path):
+            plugins_dir = PERSONAL_PLUGINS_DIR
+            filepath = personal_path
+        else:
+            raise FileNotFoundError(
+                f"Plugin not found in repo ({PLUGINS_DIR}) or personal ({PERSONAL_PLUGINS_DIR}): {filename}"
+            )
+
+    module_name = f"plugins.{filename[:-3]}"
+
+    # For personal plugins (outside repo), use spec_from_file_location
+    if plugins_dir != PLUGINS_DIR:
+        _ensure_src_in_path()
+        spec = importlib.util.spec_from_file_location(module_name, filepath)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Cannot create module spec for {filepath}")
+        mod = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = mod
+        spec.loader.exec_module(mod)
+    elif module_name in sys.modules:
+        mod = importlib.reload(sys.modules[module_name])
+    else:
+        mod = importlib.import_module(module_name)
+
+    tools_list = getattr(mod, "TOOLS", [])
+    tool_names = []
+
+    for func, name, description in tools_list:
+        try:
+            mcp.local_provider.remove_tool(name)
+        except Exception:
+            pass
+        t = Tool.from_function(func, name=name, description=description)
+        mcp.add_tool(t)
+        tool_names.append(name)
+
+    source_label = "personal" if plugins_dir != PLUGINS_DIR else "repo"
+    _update_registry(filename, len(tool_names), ",".join(tool_names), source_label)
+
+    return len(tool_names)
+
+
+def remove_plugin(mcp, filename: str) -> list[str]:
+    """Unregister a plugin's tools from MCP and clean the registry.
+
+    Does NOT delete plugin files — only unregisters tools to avoid
+    accidental deletion of code from repo or personal directories.
+    """
+    if not filename.endswith(".py"):
+        filename += ".py"
+
+    conn = get_db()
+    row = conn.execute("SELECT tool_names FROM plugins WHERE filename = ?", (filename,)).fetchone()
+
+    removed = []
+    if row and row["tool_names"]:
+        for name in row["tool_names"].split(","):
+            name = name.strip()
+            if name:
+                try:
+                    mcp.local_provider.remove_tool(name)
+                    removed.append(name)
+                except Exception:
+                    pass
+
+    module_name = f"plugins.{filename[:-3]}"
+    sys.modules.pop(module_name, None)
+
+    conn = get_db()
+    conn.execute("DELETE FROM plugins WHERE filename = ?", (filename,))
+    conn.commit()
+
+    return removed
+
+
+def list_plugins() -> list[dict]:
+    """List all registered plugins with source info (repo/personal)."""
+    conn = get_db()
+    rows = conn.execute(
+        "SELECT filename, tools_count, tool_names, loaded_at, created_by FROM plugins ORDER BY filename"
+    ).fetchall()
+    result = []
+    for r in rows:
+        d = dict(r)
+        d["source"] = d.get("created_by", "repo")
+        result.append(d)
+    return result
+
+
+def _update_registry(filename: str, tools_count: int, tool_names: str, created_by: str):
+    """Insert or update plugin registry entry. Non-fatal on lock — tools still work."""
+    now = time.time()
+    try:
+        conn = get_db()
+        conn.execute(
+            "INSERT INTO plugins (filename, tools_count, tool_names, loaded_at, created_by) "
+            "VALUES (?, ?, ?, ?, ?) "
+            "ON CONFLICT(filename) DO UPDATE SET tools_count=?, tool_names=?, loaded_at=?, created_by=?",
+            (filename, tools_count, tool_names, now, created_by, tools_count, tool_names, now, created_by),
+        )
+        conn.commit()
+    except Exception as e:
+        print(f"[PLUGIN REGISTRY] Skipped update for {filename}: {e}")