@mrtrinhvn/ag-kit 1.3.2 → 1.4.0

package/template/.agent/scripts/memory_mcp_server.py

@@ -0,0 +1,423 @@
+#!/usr/bin/env python3
+"""
+memory_mcp_server.py — Ag-Kit Memory MCP Server (Option B - Vector Ready)
+
+Exposes the Two-Tier Memory (graph.db) as native MCP tools to Antigravity AI.
+AI can call these tools directly without any shell commands or bootstrap.
+It uses Ollama locally to automatically generate vector embeddings for semantic search.
+
+Tools:
+  memory_save     — Save a hot node immediately (auto embeds)
+  memory_search   — Search nodes by keyword + cosine similarity
+  memory_link     — Create edge between two nodes
+  memory_graph    — Show node connections
+  memory_hot      — List hot (working memory) nodes
+  memory_cold     — List cold (consolidated) nodes
+  memory_consolidate — Merge old hot nodes into cold summaries
+  memory_status   — DB stats (node counts, tiers)
+
+Protocol: MCP stdio (JSON-RPC 2.0)
+Run: python3 .agent/scripts/memory_mcp_server.py [--db PATH]
+"""
+
+import sys
+import json
+import sqlite3
+import argparse
+import urllib.request
+import struct
+import math
+from pathlib import Path
+from datetime import datetime, timedelta
+
+# ─── Default DB path: resolve relative to this script ──────────────────────
+DEFAULT_DB = Path(__file__).parent.parent / "memory" / "graph.db"
+VALID_CATEGORIES = {"user_pref", "decision", "error", "pattern", "context", "general"}
+
+
+# ─── Option B: Vector Embeddings (Ollama) ─────────────────────────────────────
+
+def get_embedding(text: str) -> bytes:
+    try:
+        req = urllib.request.Request(
+            "http://localhost:11434/api/embed",
+            json.dumps({"model": "nomic-embed-text", "input": text}).encode('utf-8')
+        )
+        with urllib.request.urlopen(req, timeout=10) as res:
+            data = json.loads(res.read().decode('utf-8'))
+            floats = data.get("embeddings", [[]])[0]
+            if not floats: return None
+            return struct.pack(f"{len(floats)}f", *floats)
+    except Exception:
+        return None
+
+def cosine_similarity(v1: bytes, v2: bytes) -> float:
+    if not v1 or not v2: return 0.0
+    try:
+        f1 = struct.unpack(f"{len(v1)//4}f", v1)
+        f2 = struct.unpack(f"{len(v2)//4}f", v2)
+        dot = sum(a * b for a, b in zip(f1, f2))
+        norm1 = math.sqrt(sum(a * a for a in f1))
+        norm2 = math.sqrt(sum(b * b for b in f2))
+        if norm1 == 0 or norm2 == 0: return 0.0
+        return dot / (norm1 * norm2)
+    except Exception:
+        return 0.0
+
+
+# ─── DB ────────────────────────────────────────────────────────────────────
+
+def get_db(db_path: Path) -> sqlite3.Connection:
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA journal_mode=WAL")
+    _ensure_schema(conn)
+    return conn
+
+
+def _ensure_schema(conn):
+    conn.executescript("""
+        CREATE TABLE IF NOT EXISTS nodes (
+            id          INTEGER PRIMARY KEY AUTOINCREMENT,
+            content     TEXT NOT NULL,
+            category    TEXT NOT NULL DEFAULT 'general',
+            energy      INTEGER NOT NULL DEFAULT 100,
+            tier        TEXT NOT NULL DEFAULT 'hot',
+            source_ids  TEXT,
+            embedding   BLOB,
+            created_at  TEXT NOT NULL,
+            updated_at  TEXT NOT NULL
+        );
+        CREATE TABLE IF NOT EXISTS edges (
+            id          INTEGER PRIMARY KEY AUTOINCREMENT,
+            from_node   INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
+            to_node     INTEGER NOT NULL REFERENCES nodes(id) ON DELETE CASCADE,
+            relation    TEXT NOT NULL DEFAULT 'related_to',
+            created_at  TEXT NOT NULL
+        );
+        CREATE INDEX IF NOT EXISTS idx_nodes_tier     ON nodes(tier);
+        CREATE INDEX IF NOT EXISTS idx_nodes_category ON nodes(category);
+        CREATE INDEX IF NOT EXISTS idx_nodes_energy   ON nodes(energy);
+        CREATE INDEX IF NOT EXISTS idx_edges_from     ON edges(from_node);
+    """)
+    cols = [r[1] for r in conn.execute("PRAGMA table_info(nodes)").fetchall()]
+    for col, typedef in [("tier", "TEXT DEFAULT 'hot'"), ("source_ids", "TEXT"), ("embedding", "BLOB")]:
+        if col not in cols:
+            conn.execute(f"ALTER TABLE nodes ADD COLUMN {col} {typedef}")
+    conn.commit()
+
+
+# ─── Tool Implementations ──────────────────────────────────────────────────
+
+def tool_memory_save(args: dict, db: Path) -> str:
+    content = args.get("content", "").strip()
+    if not content: return "❌ content is required"
+    category = args.get("category", "general")
+    if category not in VALID_CATEGORIES: category = "general"
+    
+    now = datetime.now().isoformat()
+    emb = get_embedding(content)
+    
+    conn = get_db(db)
+    cur = conn.execute(
+        "INSERT INTO nodes (content, category, energy, tier, embedding, created_at, updated_at) VALUES (?, ?, 100, 'hot', ?, ?, ?)",
+        (content, category, emb, now, now)
+    )
+    conn.commit()
+    conn.close()
+    
+    emb_status = "✨[Vector]" if emb else "⚠️[No-Vector]"
+    return f"✅ {emb_status} Saved hot node #{cur.lastrowid} [{category}]: {content[:80]}"
+
+
+def tool_memory_search(args: dict, db: Path) -> str:
+    keyword = args.get("keyword", "").strip()
+    if not keyword: return "❌ keyword is required"
+    tier = args.get("tier")
+    limit = int(args.get("limit", 8))
+    tier_filter = f"AND tier = '{tier}'" if tier in ("hot", "cold") else ""
+    
+    conn = get_db(db)
+    
+    # 1. Exact match
+    exact_rows = conn.execute(
+        f"SELECT id, tier, category, energy, content, created_at FROM nodes WHERE content LIKE ? {tier_filter} ORDER BY energy DESC LIMIT ?",
+        (f"%{keyword}%", limit)
+    ).fetchall()
+    
+    results = []
+    seen = set()
+    for r in exact_rows:
+        seen.add(r["id"])
+        results.append({"row": r, "score": 1.0, "match": "exact"})
+
+    # 2. Semantic search
+    query_emb = get_embedding(keyword)
+    if query_emb:
+        all_nodes = conn.execute(f"SELECT id, tier, category, energy, content, created_at, embedding FROM nodes WHERE embedding IS NOT NULL {tier_filter}").fetchall()
+        for r in all_nodes:
+            if r["id"] in seen: continue
+            sim = cosine_similarity(query_emb, r["embedding"])
+            if sim > 0.60:
+                results.append({"row": dict(r), "score": sim, "match": f"semantic {sim:.2f}"})
+
+    results.sort(key=lambda x: (x["score"], x["row"]["energy"]), reverse=True)
+    results = results[:limit]
+
+    if not results:
+        conn.close()
+        return f"🔍 No results for '{keyword}'"
+
+    ids = [res["row"]["id"] for res in results]
+    if ids:
+        conn.execute(
+            f"UPDATE nodes SET energy=MIN(100,energy+10), updated_at=? WHERE id IN ({','.join('?'*len(ids))})",
+            [datetime.now().isoformat()] + ids
+        )
+        conn.commit()
+    conn.close()
+
+    lines = [f"🔍 Found {len(results)} result(s) for '{keyword}':\n"]
+    for res in results:
+        r = res["row"]
+        icon = "🔥" if r["tier"] == "hot" else "❄️"
+        lines.append(f"{icon} #{r['id']} [{r['category']}] ⚡{r['energy']} ({res['match']}) — {r['content'][:150]}")
+    return "\n".join(lines)
+
+
+def tool_memory_link(args: dict, db: Path) -> str:
+    from_id = args.get("from_id")
+    to_id = args.get("to_id")
+    relation = args.get("relation", "related_to")
+    if not from_id or not to_id: return "❌ from_id and to_id are required"
+    now = datetime.now().isoformat()
+    conn = get_db(db)
+    a = conn.execute("SELECT content FROM nodes WHERE id=?", (from_id,)).fetchone()
+    b = conn.execute("SELECT content FROM nodes WHERE id=?", (to_id,)).fetchone()
+    if not a or not b:
+        conn.close()
+        return f"❌ Node #{from_id} or #{to_id} not found"
+    conn.execute(
+        "INSERT INTO edges (from_node, to_node, relation, created_at) VALUES (?, ?, ?, ?)",
+        (from_id, to_id, relation, now)
+    )
+    conn.commit()
+    conn.close()
+    return f"🔗 #{from_id} --[{relation}]--> #{to_id}\n  From: {a['content'][:60]}\n  To: {b['content'][:60]}"
+
+
+def tool_memory_graph(args: dict, db: Path) -> str:
+    node_id = args.get("node_id")
+    if not node_id: return "❌ node_id is required"
+    conn = get_db(db)
+    node = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
+    if not node:
+        conn.close()
+        return f"❌ Node #{node_id} not found"
+    out = [f"🕸️ Node #{node['id']} [{node['tier'].upper()}] [{node['category']}] ⚡{node['energy']}"]
+    out.append(f"   {node['content'][:200]}\n")
+    outgoing = conn.execute("SELECT e.relation, e.to_node, n.content FROM edges e JOIN nodes n ON e.to_node=n.id WHERE e.from_node=?", (node_id,)).fetchall()
+    if outgoing:
+        out.append(f"→ OUTGOING ({len(outgoing)}):")
+        for e in outgoing: out.append(f"  --[{e['relation']}]--> #{e['to_node']}: {e['content'][:80]}")
+    incoming = conn.execute("SELECT e.relation, e.from_node, n.content FROM edges e JOIN nodes n ON e.from_node=n.id WHERE e.to_node=?", (node_id,)).fetchall()
+    if incoming:
+        out.append(f"\n← INCOMING ({len(incoming)}):")
+        for e in incoming: out.append(f"  #{e['from_node']}: {e['content'][:80]} --[{e['relation']}]-->")
+    conn.close()
+    return "\n".join(out)
+
+
+def tool_memory_hot(args: dict, db: Path) -> str:
+    limit = int(args.get("limit", 10))
+    conn = get_db(db)
+    rows = conn.execute(
+        "SELECT id, category, energy, content, embedding, created_at FROM nodes WHERE tier='hot' ORDER BY energy DESC, updated_at DESC LIMIT ?",
+        (limit,)
+    ).fetchall()
+    conn.close()
+    if not rows: return "🔥 No hot nodes yet."
+    lines = [f"🔥 HOT nodes ({len(rows)}):\n"]
+    for r in rows:
+        lines.append(f"  {'✨' if r['embedding'] else '  '}#{r['id']} [{r['category']}] ⚡{r['energy']} 📅{r['created_at'][:10]}")
+        lines.append(f"     {r['content'][:120]}")
+    return "\n".join(lines)
+
+
+def tool_memory_cold(args: dict, db: Path) -> str:
+    limit = int(args.get("limit", 10))
+    conn = get_db(db)
+    rows = conn.execute(
+        "SELECT id, category, energy, content, source_ids, updated_at FROM nodes WHERE tier='cold' ORDER BY energy DESC LIMIT ?",
+        (limit,)
+    ).fetchall()
+    conn.close()
+    if not rows: return "❄️ No cold nodes yet."
+    lines = [f"❄️ COLD nodes ({len(rows)} consolidated):\n"]
+    for r in rows:
+        srcs = len(json.loads(r["source_ids"] or "[]"))
+        lines.append(f"  #{r['id']} [{r['category']}] ⚡{r['energy']} (from {srcs} hot nodes)")
+        lines.append(f"     {r['content'][:200]}")
+    return "\n".join(lines)
+
+
+def tool_memory_consolidate(args: dict, db: Path) -> str:
+    days = int(args.get("days", 7))
+    category = args.get("category")
+    threshold = (datetime.now() - timedelta(days=days)).isoformat()
+    cat_filter = f"AND category='{category}'" if category else ""
+    
+    conn = get_db(db)
+    hot_rows = conn.execute(
+        f"SELECT id, content, category FROM nodes WHERE tier='hot' AND updated_at<? {cat_filter} ORDER BY category, energy DESC",
+        (threshold,)
+    ).fetchall()
+    
+    if not hot_rows:
+        conn.close()
+        return f"⚡ No hot nodes older than {days} days to consolidate."
+    
+    from collections import defaultdict
+    groups = defaultdict(list)
+    for r in hot_rows:
+        groups[r["category"]].append(r)
+        
+    results = []
+    now = datetime.now().isoformat()
+    for cat, nodes in groups.items():
+        node_ids = [r["id"] for r in nodes]
+        merged = "\n".join(f"• {r['content'][:200]}" for r in nodes[:20])
+        summary = f"[Consolidated {len(nodes)} nodes — {cat}]\n{merged}"
+        
+        emb = get_embedding(summary)
+        
+        cur = conn.execute(
+            "INSERT INTO nodes (content, category, energy, tier, source_ids, embedding, created_at, updated_at) VALUES (?, ?, 80, 'cold', ?, ?, ?, ?)",
+            (summary, cat, json.dumps(node_ids), emb, now, now)
+        )
+        cold_id = cur.lastrowid
+        for nid in node_ids:
+            conn.execute("INSERT INTO edges (from_node, to_node, relation, created_at) VALUES (?, ?, 'consolidated_into', ?)", (nid, cold_id, now))
+            conn.execute("UPDATE nodes SET energy=MAX(0, energy-40) WHERE id=?", (nid,))
+        results.append(f"❄️ [{cat}] merged {len(nodes)} hot nodes → cold #{cold_id}")
+    
+    conn.commit()
+    conn.close()
+    return "\n".join(results) + f"\n\n✅ Done. Summary vectors calculated."
+
+
+def tool_memory_status(args: dict, db: Path) -> str:
+    conn = get_db(db)
+    total = conn.execute("SELECT COUNT(*) FROM nodes").fetchone()[0]
+    hot = conn.execute("SELECT COUNT(*) FROM nodes WHERE tier='hot'").fetchone()[0]
+    cold = conn.execute("SELECT COUNT(*) FROM nodes WHERE tier='cold'").fetchone()[0]
+    vecs = conn.execute("SELECT COUNT(*) FROM nodes WHERE embedding IS NOT NULL").fetchone()[0]
+    edges = conn.execute("SELECT COUNT(*) FROM edges").fetchone()[0]
+    cats = conn.execute("SELECT category, COUNT(*) n FROM nodes GROUP BY category ORDER BY n DESC").fetchall()
+    conn.close()
+    lines = [
+        f"🧠 Memory Graph Status",
+        f"   DB: {db}",
+        f"   Nodes: {total} total (🔥{hot} hot | ❄️{cold} cold | ✨{vecs} vector embeddings)",
+        f"   Edges: {edges}",
+        f"   Categories: " + ", ".join(f"{r['category']}({r['n']})" for r in cats)
+    ]
+    return "\n".join(lines)
+
+
+# ─── MCP Protocol (stdio JSON-RPC 2.0) ────────────────────────────────────
+
+TOOLS = {
+    "memory_save": {
+        "description": "Save an important piece of knowledge as a hot memory node. Will automatically generate vector embeddings.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "content": {"type": "string", "description": "The knowledge to save"},
+                "category": {"type": "string", "enum": list(VALID_CATEGORIES), "description": "Category of knowledge"}
+            },
+            "required": ["content"]
+        }
+    },
+    "memory_search": {
+        "description": "Search memory graph by keyword. Uses Semantic Cosine Similarity search under the hood.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "keyword": {"type": "string", "description": "Keyword to search"},
+                "tier": {"type": "string", "enum": ["hot", "cold"], "description": "Search only in hot or cold tier"},
+                "limit": {"type": "integer", "description": "Max results (default 8)"}
+            },
+            "required": ["keyword"]
+        }
+    },
+    "memory_link": {
+        "description": "Create a relationship (edge) between two memory nodes.",
+        "inputSchema": {"type": "object", "properties": {"from_id": {"type": "integer"}, "to_id": {"type": "integer"}, "relation": {"type": "string"}}, "required": ["from_id", "to_id"]}
+    },
+    "memory_graph": {"description": "Show all edges (incoming and outgoing) for a node.", "inputSchema": {"type": "object", "properties": {"node_id": {"type": "integer"}}, "required": ["node_id"]}},
+    "memory_hot": {"description": "List hot (working memory) nodes.", "inputSchema": {"type": "object", "properties": {"limit": {"type": "integer"}}}},
+    "memory_cold": {"description": "List cold (consolidated long-term memory) nodes.", "inputSchema": {"type": "object", "properties": {"limit": {"type": "integer"}}}},
+    "memory_consolidate": {"description": "Merge old hot nodes into cold consolidated summaries.", "inputSchema": {"type": "object", "properties": {"days": {"type": "integer"}, "category": {"type": "string"}}}},
+    "memory_status": {"description": "Show memory graph statistics.", "inputSchema": {"type": "object", "properties": {}}}
+}
+
+TOOL_FNS = {
+    "memory_save": tool_memory_save, "memory_search": tool_memory_search,
+    "memory_link": tool_memory_link, "memory_graph": tool_memory_graph,
+    "memory_hot": tool_memory_hot, "memory_cold": tool_memory_cold,
+    "memory_consolidate": tool_memory_consolidate, "memory_status": tool_memory_status,
+}
+
+def send(obj: dict):
+    sys.stdout.write(json.dumps(obj) + "\n")
+    sys.stdout.flush()
+
+def handle(req: dict, db: Path):
+    method = req.get("method", "")
+    req_id = req.get("id")
+
+    if method == "initialize":
+        send({"jsonrpc": "2.0", "id": req_id, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "ag-kit-memory", "version": "1.0.1"}}})
+    elif method == "tools/list":
+        send({"jsonrpc": "2.0", "id": req_id, "result": {"tools": [{"name": name, **info} for name, info in TOOLS.items()]}})
+    elif method == "tools/call":
+        name = req.get("params", {}).get("name", "")
+        args = req.get("params", {}).get("arguments", {})
+        fn = TOOL_FNS.get(name)
+        if fn:
+            try:
+                result_text = fn(args, db)
+                send({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": result_text}], "isError": False}})
+            except Exception as e:
+                send({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": f"❌ Error: {e}"}], "isError": True}})
+        else:
+            send({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Tool not found: {name}"}})
+    elif method == "notifications/initialized":
+        pass
+    else:
+        if req_id is not None:
+            send({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}})
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--db", default=str(DEFAULT_DB))
+    args = p.parse_args()
+    db = Path(args.db)
+    
+    # Ensure DB ready
+    conn = get_db(db)
+    conn.close()
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line: continue
+        try:
+            req = json.loads(line)
+            handle(req, db)
+        except json.JSONDecodeError: pass
+        except Exception as e: sys.stderr.write(f"[memory_mcp] Error: {e}\n")
+
+if __name__ == "__main__":
+    main()