@launchapp-dev/ao-memory-mcp 1.0.0 → 2.0.0

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@launchapp-dev/ao-memory-mcp",
-  "version": "1.0.0",
-  "description": "MCP server for agent memory management — structured recall, cross-project patterns, and lifecycle management for ao-cli agents",
+  "version": "2.0.0",
+  "description": "Cognitive memory MCP server for AI agents — semantic search, document RAG, knowledge graph, episodic memory, and hybrid retrieval",
   "type": "module",
   "bin": {
     "ao-memory-mcp": "./src/server.ts"
@@ -12,8 +12,10 @@
     "build": "tsc --noEmit -p tsconfig.json"
   },
   "dependencies": {
+    "@huggingface/transformers": "^3.0.0",
     "@modelcontextprotocol/sdk": "^1.0.0",
     "better-sqlite3": "^11.0.0",
+    "sqlite-vec": "^0.1.0",
     "zod": "^3.22.0"
   },
   "devDependencies": {

package/src/db.ts

@@ -7,6 +7,8 @@ import { fileURLToPath } from "node:url";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 
+let vecLoaded = false;
+
 export function resolveDbPath(cliDbPath?: string): string {
   if (cliDbPath) return cliDbPath;
   if (process.env.AO_MEMORY_DB) return process.env.AO_MEMORY_DB;
@@ -23,16 +25,41 @@ export function initDb(dbPath: string): Database.Database {
   return db;
 }
 
-export function contentHash(
-  entryType: string,
-  agentRole: string,
-  project: string,
-  title: string,
-  body: string
-): string {
-  return createHash("sha256")
-    .update(`${entryType}\0${agentRole}\0${project}\0${title}\0${body}`)
-    .digest("hex");
+export function initVec(db: Database.Database, dimensions: number) {
+  if (vecLoaded) return;
+  try {
+    const sqliteVec = require("sqlite-vec");
+    sqliteVec.load(db);
+    vecLoaded = true;
+  } catch {
+    try {
+      // Fallback: the package might export differently
+      const mod = require("sqlite-vec");
+      if (mod.default?.load) mod.default.load(db);
+      else if (mod.load) mod.load(db);
+      vecLoaded = true;
+    } catch (e) {
+      console.error("[ao-memory] sqlite-vec not available, vector search disabled");
+      return;
+    }
+  }
+
+  db.exec(`
+    CREATE VIRTUAL TABLE IF NOT EXISTS vec_memories USING vec0(
+      embedding float[${dimensions}] distance_metric=cosine
+    );
+    CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0(
+      embedding float[${dimensions}] distance_metric=cosine
+    );
+  `);
+}
+
+export function isVecAvailable(): boolean {
+  return vecLoaded;
+}
+
+export function contentHash(...parts: string[]): string {
+  return createHash("sha256").update(parts.join("\0")).digest("hex");
 }
 
 export function now(): string {
@@ -46,3 +73,34 @@ export function jsonResult(data: unknown) {
 export function errorResult(message: string) {
   return { content: [{ type: "text" as const, text: JSON.stringify({ error: message }) }], isError: true };
 }
+
+export function touchAccess(db: Database.Database, id: number) {
+  db.prepare("UPDATE memories SET last_accessed_at = ?, access_count = access_count + 1 WHERE id = ?").run(now(), id);
+}
+
+export function chunkText(text: string, maxChars: number = 1000, overlap: number = 100): { content: string; offset: number }[] {
+  if (text.length <= maxChars) {
+    return [{ content: text, offset: 0 }];
+  }
+
+  const chunks: { content: string; offset: number }[] = [];
+  let offset = 0;
+  while (offset < text.length) {
+    let end = Math.min(offset + maxChars, text.length);
+
+    if (end < text.length) {
+      const paraBreak = text.lastIndexOf("\n\n", end);
+      if (paraBreak > offset + maxChars * 0.3) end = paraBreak + 2;
+      else {
+        const sentBreak = text.lastIndexOf(". ", end);
+        if (sentBreak > offset + maxChars * 0.3) end = sentBreak + 2;
+      }
+    }
+
+    chunks.push({ content: text.slice(offset, end).trim(), offset });
+    offset = end - overlap;
+    if (offset >= text.length) break;
+  }
+
+  return chunks.filter(c => c.content.length > 0);
+}

package/src/embeddings.ts

@@ -0,0 +1,97 @@
+import type Database from "better-sqlite3";
+import { isVecAvailable } from "./db.ts";
+
+let extractor: any = null;
+
+const DEFAULT_MODEL = "nomic-ai/nomic-embed-text-v1.5";
+const NOMIC_DIMS = 768;
+const MINILM_DIMS = 384;
+
+export function getModelId(): string {
+  return process.env.AO_MEMORY_MODEL || DEFAULT_MODEL;
+}
+
+function isNomicModel(): boolean {
+  return getModelId().includes("nomic");
+}
+
+export function getDimensions(): number {
+  return isNomicModel() ? NOMIC_DIMS : MINILM_DIMS;
+}
+
+async function getExtractor() {
+  if (extractor) return extractor;
+  const { pipeline } = await import("@huggingface/transformers");
+  const model = getModelId();
+  console.error(`[ao-memory] Loading embedding model: ${model}`);
+  extractor = await pipeline("feature-extraction", model, { dtype: "q8" });
+  console.error(`[ao-memory] Model ready (${getDimensions()}d)`);
+  return extractor;
+}
+
+export async function embed(text: string, isQuery: boolean = false): Promise<Float32Array> {
+  const ext = await getExtractor();
+  const input = isNomicModel()
+    ? (isQuery ? `search_query: ${text}` : `search_document: ${text}`)
+    : text;
+  const output = await ext(input, { pooling: "mean", normalize: true });
+  return new Float32Array(output.data);
+}
+
+export function storeVector(db: Database.Database, table: string, rowid: number, embedding: Float32Array) {
+  if (!isVecAvailable()) return;
+  db.prepare(`INSERT OR REPLACE INTO ${table}(rowid, embedding) VALUES (?, ?)`).run(
+    BigInt(rowid), Buffer.from(embedding.buffer)
+  );
+}
+
+export function deleteVector(db: Database.Database, table: string, rowid: number) {
+  if (!isVecAvailable()) return;
+  db.prepare(`DELETE FROM ${table} WHERE rowid = ?`).run(BigInt(rowid));
+}
+
+export function searchVectors(db: Database.Database, table: string, queryEmbedding: Float32Array, limit: number = 20): { rowid: number; distance: number }[] {
+  if (!isVecAvailable()) return [];
+  return db.prepare(
+    `SELECT rowid, distance FROM ${table} WHERE embedding MATCH ? ORDER BY distance LIMIT ?`
+  ).all(Buffer.from(queryEmbedding.buffer), limit) as any[];
+}
+
+export function hybridSearch(
+  db: Database.Database,
+  ftsTable: string,
+  vecTable: string,
+  queryText: string,
+  queryEmbedding: Float32Array,
+  limit: number = 10,
+  alpha: number = 0.5
+): { rowid: number; score: number }[] {
+  const RRF_K = 60;
+  const scores = new Map<number, number>();
+
+  // FTS5 keyword results
+  try {
+    const ftsResults = db.prepare(
+      `SELECT rowid FROM ${ftsTable} WHERE ${ftsTable} MATCH ? LIMIT 30`
+    ).all(queryText) as any[];
+
+    ftsResults.forEach((r, i) => {
+      const id = Number(r.rowid);
+      scores.set(id, (scores.get(id) || 0) + (1 - alpha) * (1 / (RRF_K + i + 1)));
+    });
+  } catch {}
+
+  // Vector similarity results
+  if (isVecAvailable()) {
+    const vecResults = searchVectors(db, vecTable, queryEmbedding, 30);
+    vecResults.forEach((r, i) => {
+      const id = Number(r.rowid);
+      scores.set(id, (scores.get(id) || 0) + alpha * (1 / (RRF_K + i + 1)));
+    });
+  }
+
+  return [...scores.entries()]
+    .sort((a, b) => b[1] - a[1])
+    .slice(0, limit)
+    .map(([rowid, score]) => ({ rowid, score }));
+}

package/src/schema.sql

@@ -1,113 +1,194 @@
 PRAGMA journal_mode = WAL;
 PRAGMA foreign_keys = ON;
 
-CREATE TABLE IF NOT EXISTS memory_entries (
+-- ============================================================
+-- MEMORIES — unified store for semantic, episodic, procedural
+-- ============================================================
+CREATE TABLE IF NOT EXISTS memories (
+  id              INTEGER PRIMARY KEY AUTOINCREMENT,
+  memory_type     TEXT    NOT NULL,  -- semantic | episodic | procedural
+  scope           TEXT    NOT NULL DEFAULT 'project',  -- global | user | project | session
+  namespace       TEXT,              -- project name, user id, session id, etc.
+  agent_role      TEXT,              -- planner, reviewer, qa-tester, or any custom role
+  title           TEXT    NOT NULL,
+  content         TEXT    NOT NULL,
+  -- references
+  task_id         TEXT,
+  pr_number       INTEGER,
+  run_id          TEXT,
+  -- lifecycle
+  status          TEXT    NOT NULL DEFAULT 'active',  -- active | summarized | archived
+  confidence      REAL    NOT NULL DEFAULT 1.0,       -- 0.0-1.0, decays over time
+  superseded_by   INTEGER REFERENCES memories(id),
+  -- temporal
+  tags            TEXT    NOT NULL DEFAULT '[]',
+  metadata        TEXT    NOT NULL DEFAULT '{}',
+  created_at      TEXT    NOT NULL,
+  occurred_at     TEXT    NOT NULL,
+  updated_at      TEXT    NOT NULL,
+  last_accessed_at TEXT,
+  access_count    INTEGER NOT NULL DEFAULT 0,
+  content_hash    TEXT    NOT NULL
+);
+
+CREATE INDEX IF NOT EXISTS idx_mem_type       ON memories(memory_type);
+CREATE INDEX IF NOT EXISTS idx_mem_scope      ON memories(scope, namespace);
+CREATE INDEX IF NOT EXISTS idx_mem_role       ON memories(agent_role);
+CREATE INDEX IF NOT EXISTS idx_mem_status     ON memories(status);
+CREATE INDEX IF NOT EXISTS idx_mem_task       ON memories(task_id);
+CREATE INDEX IF NOT EXISTS idx_mem_occurred   ON memories(occurred_at);
+CREATE INDEX IF NOT EXISTS idx_mem_hash       ON memories(content_hash);
+CREATE INDEX IF NOT EXISTS idx_mem_ns_type    ON memories(namespace, memory_type);
+CREATE INDEX IF NOT EXISTS idx_mem_ns_role    ON memories(namespace, agent_role);
+CREATE INDEX IF NOT EXISTS idx_mem_confidence ON memories(confidence);
+CREATE INDEX IF NOT EXISTS idx_mem_accessed   ON memories(last_accessed_at);
+
+-- FTS5 for memories
+CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5(
+  title, content,
+  content=memories, content_rowid=id
+);
+CREATE TRIGGER IF NOT EXISTS mem_fts_i AFTER INSERT ON memories BEGIN
+  INSERT INTO memories_fts(rowid, title, content) VALUES (new.id, new.title, new.content);
+END;
+CREATE TRIGGER IF NOT EXISTS mem_fts_u AFTER UPDATE ON memories BEGIN
+  INSERT INTO memories_fts(memories_fts, rowid, title, content) VALUES ('delete', old.id, old.title, old.content);
+  INSERT INTO memories_fts(rowid, title, content) VALUES (new.id, new.title, new.content);
+END;
+CREATE TRIGGER IF NOT EXISTS mem_fts_d AFTER DELETE ON memories BEGIN
+  INSERT INTO memories_fts(memories_fts, rowid, title, content) VALUES ('delete', old.id, old.title, old.content);
+END;
+
+-- ============================================================
+-- DOCUMENTS — source documents ingested for RAG
+-- ============================================================
+CREATE TABLE IF NOT EXISTS documents (
+  id          INTEGER PRIMARY KEY AUTOINCREMENT,
+  namespace   TEXT,
+  title       TEXT    NOT NULL,
+  source      TEXT,              -- file path, URL, or identifier
+  mime_type   TEXT    DEFAULT 'text/plain',
+  content     TEXT    NOT NULL,  -- full original content
+  metadata    TEXT    NOT NULL DEFAULT '{}',
+  created_at  TEXT    NOT NULL,
+  updated_at  TEXT    NOT NULL
+);
+
+CREATE INDEX IF NOT EXISTS idx_doc_ns ON documents(namespace);
+
+-- CHUNKS — document chunks with embeddings
+CREATE TABLE IF NOT EXISTS chunks (
   id           INTEGER PRIMARY KEY AUTOINCREMENT,
-  entry_type   TEXT    NOT NULL,
-  agent_role   TEXT    NOT NULL,
-  project      TEXT    NOT NULL,
-  title        TEXT    NOT NULL,
-  body         TEXT    NOT NULL,
-  task_id      TEXT,
-  pr_number    INTEGER,
-  run_id       TEXT,
-  status       TEXT    NOT NULL DEFAULT 'active',
-  tags         TEXT    NOT NULL DEFAULT '[]',
+  document_id  INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
+  chunk_index  INTEGER NOT NULL,
+  content      TEXT    NOT NULL,
+  char_offset  INTEGER NOT NULL DEFAULT 0,
+  char_length  INTEGER NOT NULL DEFAULT 0,
   metadata     TEXT    NOT NULL DEFAULT '{}',
-  created_at   TEXT    NOT NULL,
-  occurred_at  TEXT    NOT NULL,
-  updated_at   TEXT    NOT NULL,
-  content_hash TEXT    NOT NULL
+  created_at   TEXT    NOT NULL
+);
+
+CREATE INDEX IF NOT EXISTS idx_chunk_doc ON chunks(document_id);
+
+-- FTS5 for chunks
+CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
+  content,
+  content=chunks, content_rowid=id
 );
+CREATE TRIGGER IF NOT EXISTS chunk_fts_i AFTER INSERT ON chunks BEGIN
+  INSERT INTO chunks_fts(rowid, content) VALUES (new.id, new.content);
+END;
+CREATE TRIGGER IF NOT EXISTS chunk_fts_u AFTER UPDATE ON chunks BEGIN
+  INSERT INTO chunks_fts(chunks_fts, rowid, content) VALUES ('delete', old.id, old.content);
+  INSERT INTO chunks_fts(rowid, content) VALUES (new.id, new.content);
+END;
+CREATE TRIGGER IF NOT EXISTS chunk_fts_d AFTER DELETE ON chunks BEGIN
+  INSERT INTO chunks_fts(chunks_fts, rowid, content) VALUES ('delete', old.id, old.content);
+END;
 
-CREATE TABLE IF NOT EXISTS memory_summaries (
-  id            INTEGER PRIMARY KEY AUTOINCREMENT,
-  agent_role    TEXT    NOT NULL,
-  project       TEXT    NOT NULL,
-  entry_type    TEXT,
-  title         TEXT    NOT NULL,
-  body          TEXT    NOT NULL,
-  entry_count   INTEGER NOT NULL,
-  date_from     TEXT    NOT NULL,
-  date_to       TEXT    NOT NULL,
-  entry_ids     TEXT    NOT NULL,
-  created_at    TEXT    NOT NULL
+-- ============================================================
+-- KNOWLEDGE GRAPH — entities and relations
+-- ============================================================
+CREATE TABLE IF NOT EXISTS entities (
+  id          INTEGER PRIMARY KEY AUTOINCREMENT,
+  name        TEXT    NOT NULL,
+  entity_type TEXT    NOT NULL,  -- project, person, technology, concept, file, etc.
+  namespace   TEXT,
+  description TEXT,
+  metadata    TEXT    NOT NULL DEFAULT '{}',
+  created_at  TEXT    NOT NULL,
+  updated_at  TEXT    NOT NULL,
+  UNIQUE(name, entity_type, namespace)
 );
 
-CREATE TABLE IF NOT EXISTS memory_patterns (
+CREATE INDEX IF NOT EXISTS idx_ent_type ON entities(entity_type);
+CREATE INDEX IF NOT EXISTS idx_ent_ns   ON entities(namespace);
+
+CREATE TABLE IF NOT EXISTS relations (
   id               INTEGER PRIMARY KEY AUTOINCREMENT,
-  pattern_type     TEXT    NOT NULL,
-  title            TEXT    NOT NULL,
-  description      TEXT    NOT NULL,
-  projects         TEXT    NOT NULL DEFAULT '[]',
-  agent_roles      TEXT    NOT NULL DEFAULT '[]',
-  entry_ids        TEXT    NOT NULL DEFAULT '[]',
-  occurrence_count INTEGER NOT NULL DEFAULT 1,
-  status           TEXT    NOT NULL DEFAULT 'active',
-  first_seen       TEXT    NOT NULL,
-  last_seen        TEXT    NOT NULL,
-  resolved_at      TEXT,
+  source_entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+  relation_type    TEXT    NOT NULL,  -- uses, depends_on, created_by, part_of, related_to, etc.
+  target_entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+  weight           REAL    NOT NULL DEFAULT 1.0,
+  memory_id        INTEGER REFERENCES memories(id),  -- evidence link
+  metadata         TEXT    NOT NULL DEFAULT '{}',
   created_at       TEXT    NOT NULL,
-  updated_at       TEXT    NOT NULL
+  UNIQUE(source_entity_id, relation_type, target_entity_id)
 );
 
--- Indexes for memory_entries
-CREATE INDEX IF NOT EXISTS idx_me_entry_type   ON memory_entries(entry_type);
-CREATE INDEX IF NOT EXISTS idx_me_agent_role   ON memory_entries(agent_role);
-CREATE INDEX IF NOT EXISTS idx_me_project      ON memory_entries(project);
-CREATE INDEX IF NOT EXISTS idx_me_status       ON memory_entries(status);
-CREATE INDEX IF NOT EXISTS idx_me_task_id      ON memory_entries(task_id);
-CREATE INDEX IF NOT EXISTS idx_me_occurred_at  ON memory_entries(occurred_at);
-CREATE INDEX IF NOT EXISTS idx_me_content_hash ON memory_entries(content_hash);
-CREATE INDEX IF NOT EXISTS idx_me_role_project ON memory_entries(agent_role, project);
-CREATE INDEX IF NOT EXISTS idx_me_proj_type    ON memory_entries(project, entry_type);
-CREATE INDEX IF NOT EXISTS idx_me_proj_date    ON memory_entries(project, occurred_at);
-
--- Indexes for memory_summaries
-CREATE INDEX IF NOT EXISTS idx_ms_role_project ON memory_summaries(agent_role, project);
-
--- Indexes for memory_patterns
-CREATE INDEX IF NOT EXISTS idx_mp_status       ON memory_patterns(status);
-CREATE INDEX IF NOT EXISTS idx_mp_pattern_type ON memory_patterns(pattern_type);
-
--- FTS5 for memory_entries
-CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
-  title,
-  body,
-  content=memory_entries,
-  content_rowid=id
+CREATE INDEX IF NOT EXISTS idx_rel_source ON relations(source_entity_id);
+CREATE INDEX IF NOT EXISTS idx_rel_target ON relations(target_entity_id);
+CREATE INDEX IF NOT EXISTS idx_rel_type   ON relations(relation_type);
+
+-- ============================================================
+-- EPISODES — conversation/run history
+-- ============================================================
+CREATE TABLE IF NOT EXISTS episodes (
+  id          INTEGER PRIMARY KEY AUTOINCREMENT,
+  session_id  TEXT    NOT NULL,
+  namespace   TEXT,
+  agent_role  TEXT,
+  role        TEXT    NOT NULL,  -- user | assistant | system
+  content     TEXT    NOT NULL,
+  summary     TEXT,
+  metadata    TEXT    NOT NULL DEFAULT '{}',
+  created_at  TEXT    NOT NULL
 );
 
-CREATE TRIGGER IF NOT EXISTS memory_fts_insert AFTER INSERT ON memory_entries BEGIN
-  INSERT INTO memory_fts(rowid, title, body) VALUES (new.id, new.title, new.body);
-END;
+CREATE INDEX IF NOT EXISTS idx_ep_session ON episodes(session_id);
+CREATE INDEX IF NOT EXISTS idx_ep_ns      ON episodes(namespace);
 
-CREATE TRIGGER IF NOT EXISTS memory_fts_update AFTER UPDATE ON memory_entries BEGIN
-  INSERT INTO memory_fts(memory_fts, rowid, title, body) VALUES ('delete', old.id, old.title, old.body);
-  INSERT INTO memory_fts(rowid, title, body) VALUES (new.id, new.title, new.body);
+-- FTS5 for episodes
+CREATE VIRTUAL TABLE IF NOT EXISTS episodes_fts USING fts5(
+  content, summary,
+  content=episodes, content_rowid=id
+);
+CREATE TRIGGER IF NOT EXISTS ep_fts_i AFTER INSERT ON episodes BEGIN
+  INSERT INTO episodes_fts(rowid, content, summary) VALUES (new.id, new.content, new.summary);
 END;
-
-CREATE TRIGGER IF NOT EXISTS memory_fts_delete AFTER DELETE ON memory_entries BEGIN
-  INSERT INTO memory_fts(memory_fts, rowid, title, body) VALUES ('delete', old.id, old.title, old.body);
+CREATE TRIGGER IF NOT EXISTS ep_fts_u AFTER UPDATE ON episodes BEGIN
+  INSERT INTO episodes_fts(episodes_fts, rowid, content, summary) VALUES ('delete', old.id, old.content, old.summary);
+  INSERT INTO episodes_fts(rowid, content, summary) VALUES (new.id, new.content, new.summary);
 END;
-
--- FTS5 for memory_summaries
-CREATE VIRTUAL TABLE IF NOT EXISTS memory_summaries_fts USING fts5(
-  title,
-  body,
-  content=memory_summaries,
-  content_rowid=id
-);
-
-CREATE TRIGGER IF NOT EXISTS ms_fts_insert AFTER INSERT ON memory_summaries BEGIN
-  INSERT INTO memory_summaries_fts(rowid, title, body) VALUES (new.id, new.title, new.body);
+CREATE TRIGGER IF NOT EXISTS ep_fts_d AFTER DELETE ON episodes BEGIN
+  INSERT INTO episodes_fts(episodes_fts, rowid, content, summary) VALUES ('delete', old.id, old.content, old.content);
 END;
 
-CREATE TRIGGER IF NOT EXISTS ms_fts_update AFTER UPDATE ON memory_summaries BEGIN
-  INSERT INTO memory_summaries_fts(memory_summaries_fts, rowid, title, body) VALUES ('delete', old.id, old.title, old.body);
-  INSERT INTO memory_summaries_fts(rowid, title, body) VALUES (new.id, new.title, new.body);
-END;
+-- ============================================================
+-- SUMMARIES — rolled-up digests
+-- ============================================================
+CREATE TABLE IF NOT EXISTS summaries (
+  id          INTEGER PRIMARY KEY AUTOINCREMENT,
+  scope       TEXT    NOT NULL,
+  namespace   TEXT,
+  agent_role  TEXT,
+  title       TEXT    NOT NULL,
+  content     TEXT    NOT NULL,
+  entry_count INTEGER NOT NULL,
+  date_from   TEXT    NOT NULL,
+  date_to     TEXT    NOT NULL,
+  entry_ids   TEXT    NOT NULL,  -- JSON array
+  created_at  TEXT    NOT NULL
+);
 
-CREATE TRIGGER IF NOT EXISTS ms_fts_delete AFTER DELETE ON memory_summaries BEGIN
-  INSERT INTO memory_summaries_fts(memory_summaries_fts, rowid, title, body) VALUES ('delete', old.id, old.title, old.body);
-END;
+CREATE INDEX IF NOT EXISTS idx_sum_ns ON summaries(namespace, agent_role);

package/src/server.ts

@@ -5,38 +5,45 @@ import {
   ListToolsRequestSchema,
   CallToolRequestSchema,
 } from "@modelcontextprotocol/sdk/types.js";
-import { resolveDbPath, initDb, errorResult } from "./db.ts";
+import { resolveDbPath, initDb, initVec, errorResult } from "./db.ts";
+import { getDimensions } from "./embeddings.ts";
 import { storeTools, handleStore } from "./tools/store.ts";
 import { recallTools, handleRecall } from "./tools/recall.ts";
 import { statsTools, handleStats } from "./tools/stats.ts";
 import { contextTools, handleContext } from "./tools/context.ts";
 import { summarizeTools, handleSummarize } from "./tools/summarize.ts";
-import { patternTools, handlePatterns } from "./tools/patterns.ts";
+import { documentTools, handleDocuments } from "./tools/documents.ts";
+import { knowledgeTools, handleKnowledge } from "./tools/knowledge.ts";
+import { episodeTools, handleEpisodes } from "./tools/episodes.ts";
 
 // Parse CLI args
 const args = process.argv.slice(2);
 let dbPath: string | undefined;
 for (let i = 0; i < args.length; i++) {
-  if (args[i] === "--db" && args[i + 1]) {
-    dbPath = args[++i];
-  }
+  if (args[i] === "--db" && args[i + 1]) dbPath = args[++i];
 }
 
 const db = initDb(resolveDbPath(dbPath));
+initVec(db, getDimensions());
 
 const allTools = [
   ...storeTools,
   ...recallTools,
-  ...statsTools,
+  ...documentTools,
+  ...knowledgeTools,
+  ...episodeTools,
   ...contextTools,
   ...summarizeTools,
-  ...patternTools,
+  ...statsTools,
 ];
 
-const handlers = [handleStore, handleRecall, handleStats, handleContext, handleSummarize, handlePatterns];
+const handlers: Array<(db: any, name: string, args: any) => any> = [
+  handleStore, handleRecall, handleDocuments, handleKnowledge,
+  handleEpisodes, handleContext, handleSummarize, handleStats,
+];
 
 const server = new Server(
-  { name: "ao-memory-mcp", version: "1.0.0" },
+  { name: "ao-memory-mcp", version: "2.0.0" },
   { capabilities: { tools: {} } }
 );
 
@@ -49,7 +56,11 @@ server.setRequestHandler(CallToolRequestSchema, async (req) => {
 
   for (const handler of handlers) {
     const result = handler(db, name, input || {});
-    if (result) return result;
+    if (result !== null) {
+      // Handle async results (embed operations)
+      if (result instanceof Promise) return await result;
+      return result;
+    }
   }
 
   return errorResult(`Unknown tool: ${name}`);