@aeriondyseti/vector-memory-mcp 0.2.2 → 0.2.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aeriondyseti/vector-memory-mcp",
-  "version": "0.2.2",
+  "version": "0.2.4",
   "description": "A zero-configuration RAG memory server for MCP clients",
   "type": "module",
   "main": "dist/index.js",
@@ -9,6 +9,7 @@
     "vector-memory-mcp": "dist/index.js"
   },
   "files": [
+    "src",
     "dist",
     "README.md",
     "LICENSE"
@@ -37,7 +38,7 @@
     "@lancedb/lancedb": "^0.22.3",
     "@modelcontextprotocol/sdk": "^1.0.0",
     "@huggingface/transformers": "^3.8.0",
-    "apache-arrow": "^21.1.0"
+    "apache-arrow": "^18.1.0"
   },
   "devDependencies": {
     "@types/bun": "latest",

package/src/config/index.ts

@@ -0,0 +1,29 @@
+import { join } from "path";
+import { homedir } from "os";
+
+export interface Config {
+  dbPath: string;
+  embeddingModel: string;
+  embeddingDimension: number;
+}
+
+const DEFAULT_DB_PATH = join(
+  homedir(),
+  ".local",
+  "share",
+  "vector-memory-mcp",
+  "memories.db"
+);
+
+const DEFAULT_EMBEDDING_MODEL = "Xenova/all-MiniLM-L6-v2";
+const DEFAULT_EMBEDDING_DIMENSION = 384;
+
+export function loadConfig(): Config {
+  return {
+    dbPath: process.env.VECTOR_MEMORY_DB_PATH ?? DEFAULT_DB_PATH,
+    embeddingModel: process.env.VECTOR_MEMORY_MODEL ?? DEFAULT_EMBEDDING_MODEL,
+    embeddingDimension: DEFAULT_EMBEDDING_DIMENSION,
+  };
+}
+
+export const config = loadConfig();

package/src/db/connection.ts

@@ -0,0 +1,11 @@
+import * as lancedb from "@lancedb/lancedb";
+import { mkdirSync } from "fs";
+import { dirname } from "path";
+
+export async function connectToDatabase(dbPath: string): Promise<lancedb.Connection> {
+  // Ensure directory exists
+  mkdirSync(dirname(dbPath), { recursive: true });
+
+  const db = await lancedb.connect(dbPath);
+  return db;
+}

package/src/db/memory.repository.ts

@@ -0,0 +1,94 @@
+import * as lancedb from "@lancedb/lancedb";
+import { TABLE_NAME, memorySchema } from "./schema.js";
+import {
+  type Memory,
+  type VectorRow,
+  DELETED_TOMBSTONE,
+} from "../types/memory.js";
+
+export class MemoryRepository {
+  constructor(private db: lancedb.Connection) {}
+
+  private async getTable() {
+    const names = await this.db.tableNames();
+    if (names.includes(TABLE_NAME)) {
+      return await this.db.openTable(TABLE_NAME);
+    }
+    // Create with empty data to initialize schema
+    return await this.db.createTable(TABLE_NAME, [], { schema: memorySchema });
+  }
+
+  async insert(memory: Memory): Promise<void> {
+    const table = await this.getTable();
+    await table.add([
+      {
+        id: memory.id,
+        vector: memory.embedding,
+        content: memory.content,
+        metadata: JSON.stringify(memory.metadata),
+        created_at: memory.createdAt.getTime(),
+        updated_at: memory.updatedAt.getTime(),
+        superseded_by: memory.supersededBy,
+      },
+    ]);
+  }
+
+  async findById(id: string): Promise<Memory | null> {
+    const table = await this.getTable();
+    const results = await table.query().where(`id = '${id}'`).limit(1).toArray();
+
+    if (results.length === 0) {
+      return null;
+    }
+
+    const row = results[0];
+    
+    // Handle Arrow Vector type conversion
+    // LanceDB returns an Arrow Vector object which is iterable but not an array
+    const vectorData = row.vector as any;
+    const embedding = Array.isArray(vectorData) 
+      ? vectorData 
+      : Array.from(vectorData) as number[];
+
+    return {
+      id: row.id as string,
+      content: row.content as string,
+      embedding,
+      metadata: JSON.parse(row.metadata as string),
+      createdAt: new Date(row.created_at as number),
+      updatedAt: new Date(row.updated_at as number),
+      supersededBy: row.superseded_by as string | null,
+    };
+  }
+
+  async markDeleted(id: string): Promise<boolean> {
+    const table = await this.getTable();
+    
+    // Verify existence first to match previous behavior (return false if not found)
+    const existing = await table.query().where(`id = '${id}'`).limit(1).toArray();
+    if (existing.length === 0) {
+      return false;
+    }
+
+    const now = Date.now();
+    await table.update({
+      where: `id = '${id}'`,
+      values: {
+        superseded_by: DELETED_TOMBSTONE,
+        updated_at: now,
+      },
+    });
+
+    return true;
+  }
+
+  async findSimilar(embedding: number[], limit: number): Promise<VectorRow[]> {
+    const table = await this.getTable();
+    const results = await table.vectorSearch(embedding).limit(limit).toArray();
+
+    return results.map((r) => ({
+      id: r.id as string,
+      distance: r._distance as number,
+    }));
+  }
+}

package/src/db/schema.ts

@@ -0,0 +1,33 @@
+import {
+  Schema,
+  Field,
+  FixedSizeList,
+  Float32,
+  Utf8,
+  Timestamp,
+  TimeUnit,
+} from "apache-arrow";
+
+export const TABLE_NAME = "memories";
+
+export const memorySchema = new Schema([
+  new Field("id", new Utf8(), false),
+  new Field(
+    "vector",
+    new FixedSizeList(384, new Field("item", new Float32())),
+    false
+  ),
+  new Field("content", new Utf8(), false),
+  new Field("metadata", new Utf8(), false), // JSON string
+  new Field(
+    "created_at",
+    new Timestamp(TimeUnit.MILLISECOND),
+    false
+  ),
+  new Field(
+    "updated_at",
+    new Timestamp(TimeUnit.MILLISECOND),
+    false
+  ),
+  new Field("superseded_by", new Utf8(), true), // Nullable
+]);

package/src/index.ts

@@ -0,0 +1,23 @@
+#!/usr/bin/env node
+
+import { config } from "./config/index.js";
+import { connectToDatabase } from "./db/connection.js";
+import { MemoryRepository } from "./db/memory.repository.js";
+import { EmbeddingsService } from "./services/embeddings.service.js";
+import { MemoryService } from "./services/memory.service.js";
+import { startServer } from "./mcp/server.js";
+
+async function main(): Promise<void> {
+  // Initialize database
+  const db = await connectToDatabase(config.dbPath);
+
+  // Initialize layers
+  const repository = new MemoryRepository(db);
+  const embeddings = new EmbeddingsService(config.embeddingModel, config.embeddingDimension);
+  const memoryService = new MemoryService(repository, embeddings);
+
+  // Start MCP server
+  await startServer(memoryService);
+}
+
+main().catch(console.error);

package/src/mcp/handlers.ts

@@ -0,0 +1,111 @@
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { MemoryService } from "../services/memory.service.js";
+
+export async function handleStoreMemory(
+  args: Record<string, unknown> | undefined,
+  service: MemoryService
+): Promise<CallToolResult> {
+  const content = args?.content as string;
+  const metadata = (args?.metadata as Record<string, unknown>) ?? {};
+  const memory = await service.store(content, metadata);
+
+  return {
+    content: [{ type: "text", text: `Memory stored with ID: ${memory.id}` }],
+  };
+}
+
+export async function handleDeleteMemory(
+  args: Record<string, unknown> | undefined,
+  service: MemoryService
+): Promise<CallToolResult> {
+  const id = args?.id as string;
+  const success = await service.delete(id);
+
+  return {
+    content: [
+      {
+        type: "text",
+        text: success
+          ? `Memory ${id} deleted successfully`
+          : `Memory ${id} not found`,
+      },
+    ],
+  };
+}
+
+export async function handleSearchMemories(
+  args: Record<string, unknown> | undefined,
+  service: MemoryService
+): Promise<CallToolResult> {
+  const query = args?.query as string;
+  const limit = (args?.limit as number) ?? 10;
+  const memories = await service.search(query, limit);
+
+  if (memories.length === 0) {
+    return {
+      content: [{ type: "text", text: "No memories found matching your query." }],
+    };
+  }
+
+  const results = memories.map((mem) => {
+    let result = `ID: ${mem.id}\nContent: ${mem.content}`;
+    if (Object.keys(mem.metadata).length > 0) {
+      result += `\nMetadata: ${JSON.stringify(mem.metadata)}`;
+    }
+    return result;
+  });
+
+  return {
+    content: [{ type: "text", text: results.join("\n\n---\n\n") }],
+  };
+}
+
+export async function handleGetMemory(
+  args: Record<string, unknown> | undefined,
+  service: MemoryService
+): Promise<CallToolResult> {
+  const id = args?.id as string;
+  const memory = await service.get(id);
+
+  if (!memory) {
+    return {
+      content: [{ type: "text", text: `Memory ${id} not found` }],
+    };
+  }
+
+  let result = `ID: ${memory.id}\nContent: ${memory.content}`;
+  if (Object.keys(memory.metadata).length > 0) {
+    result += `\nMetadata: ${JSON.stringify(memory.metadata)}`;
+  }
+  result += `\nCreated: ${memory.createdAt.toISOString()}`;
+  result += `\nUpdated: ${memory.updatedAt.toISOString()}`;
+  if (memory.supersededBy) {
+    result += `\nSuperseded by: ${memory.supersededBy}`;
+  }
+
+  return {
+    content: [{ type: "text", text: result }],
+  };
+}
+
+export async function handleToolCall(
+  name: string,
+  args: Record<string, unknown> | undefined,
+  service: MemoryService
+): Promise<CallToolResult> {
+  switch (name) {
+    case "store_memory":
+      return handleStoreMemory(args, service);
+    case "delete_memory":
+      return handleDeleteMemory(args, service);
+    case "search_memories":
+      return handleSearchMemories(args, service);
+    case "get_memory":
+      return handleGetMemory(args, service);
+    default:
+      return {
+        content: [{ type: "text", text: `Unknown tool: ${name}` }],
+        isError: true,
+      };
+  }
+}

package/src/mcp/server.ts

@@ -0,0 +1,34 @@
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import {
+  CallToolRequestSchema,
+  ListToolsRequestSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+
+import { tools } from "./tools.js";
+import { handleToolCall } from "./handlers.js";
+import type { MemoryService } from "../services/memory.service.js";
+
+export function createServer(memoryService: MemoryService): Server {
+  const server = new Server(
+    { name: "vector-memory-mcp", version: "0.2.0" },
+    { capabilities: { tools: {} } }
+  );
+
+  server.setRequestHandler(ListToolsRequestSchema, async () => {
+    return { tools };
+  });
+
+  server.setRequestHandler(CallToolRequestSchema, async (request) => {
+    const { name, arguments: args } = request.params;
+    return handleToolCall(name, args, memoryService);
+  });
+
+  return server;
+}
+
+export async function startServer(memoryService: MemoryService): Promise<void> {
+  const server = createServer(memoryService);
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+}

package/src/mcp/tools.ts

@@ -0,0 +1,80 @@
+import type { Tool } from "@modelcontextprotocol/sdk/types.js";
+
+export const storeMemoryTool: Tool = {
+  name: "store_memory",
+  description: "Store a new memory. Use this to save information for later recall.",
+  inputSchema: {
+    type: "object",
+    properties: {
+      content: {
+        type: "string",
+        description: "The text content to store as a memory",
+      },
+      metadata: {
+        type: "object",
+        description: "Optional key-value metadata to attach to the memory",
+        additionalProperties: true,
+      },
+    },
+    required: ["content"],
+  },
+};
+
+export const deleteMemoryTool: Tool = {
+  name: "delete_memory",
+  description:
+    "Delete a memory by its ID. The memory will be soft-deleted and won't appear in search results.",
+  inputSchema: {
+    type: "object",
+    properties: {
+      id: {
+        type: "string",
+        description: "The ID of the memory to delete",
+      },
+    },
+    required: ["id"],
+  },
+};
+
+export const searchMemoriesTool: Tool = {
+  name: "search_memories",
+  description:
+    "Search for memories using semantic similarity. Returns the most relevant memories for the given query.",
+  inputSchema: {
+    type: "object",
+    properties: {
+      query: {
+        type: "string",
+        description: "The search query to find relevant memories",
+      },
+      limit: {
+        type: "integer",
+        description: "Maximum number of results to return (default: 10)",
+        default: 10,
+      },
+    },
+    required: ["query"],
+  },
+};
+
+export const getMemoryTool: Tool = {
+  name: "get_memory",
+  description: "Retrieve a specific memory by its ID.",
+  inputSchema: {
+    type: "object",
+    properties: {
+      id: {
+        type: "string",
+        description: "The ID of the memory to retrieve",
+      },
+    },
+    required: ["id"],
+  },
+};
+
+export const tools: Tool[] = [
+  storeMemoryTool,
+  deleteMemoryTool,
+  searchMemoriesTool,
+  getMemoryTool,
+];

package/src/services/embeddings.service.ts

@@ -0,0 +1,66 @@
+import { pipeline, type FeatureExtractionPipeline } from "@huggingface/transformers";
+
+export class EmbeddingsService {
+  private modelName: string;
+  private extractor: FeatureExtractionPipeline | null = null;
+  private initPromise: Promise<void> | null = null;
+  private _dimension: number;
+
+  constructor(modelName: string, dimension: number) {
+    this.modelName = modelName;
+    this._dimension = dimension;
+  }
+
+  get dimension(): number {
+    return this._dimension;
+  }
+
+  private async init(): Promise<void> {
+    if (this.extractor) {
+      return; // Already initialized
+    }
+
+    if (!this.initPromise) {
+      this.initPromise = (async () => {
+        this.extractor = await pipeline("feature-extraction", this.modelName, {
+          dtype: "fp32",
+        });
+      })();
+    }
+    await this.initPromise;
+  }
+
+  async embed(text: string): Promise<number[]> {
+    await this.init();
+    if (!this.extractor) {
+      throw new Error("Extractor not initialized.");
+    }
+
+    const output = await this.extractor(text, {
+      pooling: "mean",
+      normalize: true,
+    });
+
+    // output.tolist() returns number[][] for single input, we want the first element
+    const embeddings = output.tolist() as number[][];
+    return embeddings[0];
+  }
+
+  async embedBatch(texts: string[]): Promise<number[][]> {
+    await this.init();
+    if (!this.extractor) {
+      throw new Error("Extractor not initialized.");
+    }
+
+    if (texts.length === 0) {
+      return [];
+    }
+
+    const output = await this.extractor(texts, {
+      pooling: "mean",
+      normalize: true,
+    });
+
+    return output.tolist() as number[][];
+  }
+}

package/src/services/memory.service.ts

@@ -0,0 +1,102 @@
+import { randomUUID } from "crypto";
+import type { Memory } from "../types/memory.js";
+import { DELETED_TOMBSTONE, isSuperseded } from "../types/memory.js";
+import type { MemoryRepository } from "../db/memory.repository.js";
+import type { EmbeddingsService } from "./embeddings.service.js";
+
+export class MemoryService {
+  constructor(
+    private repository: MemoryRepository,
+    private embeddings: EmbeddingsService
+  ) {}
+
+  async store(content: string, metadata: Record<string, unknown> = {}): Promise<Memory> {
+    const id = randomUUID();
+    const now = new Date();
+    const embedding = await this.embeddings.embed(content);
+
+    const memory: Memory = {
+      id,
+      content,
+      embedding,
+      metadata,
+      createdAt: now,
+      updatedAt: now,
+      supersededBy: null,
+    };
+
+    await this.repository.insert(memory);
+    return memory;
+  }
+
+  async get(id: string): Promise<Memory | null> {
+    return await this.repository.findById(id);
+  }
+
+  async delete(id: string): Promise<boolean> {
+    return await this.repository.markDeleted(id);
+  }
+
+  async search(query: string, limit: number = 10): Promise<Memory[]> {
+    const queryEmbedding = await this.embeddings.embed(query);
+    const fetchLimit = limit * 3;
+
+    const rows = await this.repository.findSimilar(queryEmbedding, fetchLimit);
+
+    const results: Memory[] = [];
+    const seenIds = new Set<string>();
+
+    for (const row of rows) {
+      let memory = await this.repository.findById(row.id);
+
+      if (!memory) {
+        continue;
+      }
+
+      if (isSuperseded(memory)) {
+        memory = await this.followSupersessionChain(row.id);
+        if (!memory) {
+          continue;
+        }
+      }
+
+      if (seenIds.has(memory.id)) {
+        continue;
+      }
+      seenIds.add(memory.id);
+
+      results.push(memory);
+      if (results.length >= limit) {
+        break;
+      }
+    }
+
+    return results;
+  }
+
+  private async followSupersessionChain(memoryId: string): Promise<Memory | null> {
+    const visited = new Set<string>();
+    let currentId: string | null = memoryId;
+
+    while (currentId && !visited.has(currentId)) {
+      visited.add(currentId);
+      const memory = await this.repository.findById(currentId);
+
+      if (!memory) {
+        return null;
+      }
+
+      if (memory.supersededBy === null) {
+        return memory;
+      }
+
+      if (memory.supersededBy === DELETED_TOMBSTONE) {
+        return null;
+      }
+
+      currentId = memory.supersededBy;
+    }
+
+    return null;
+  }
+}

package/src/types/memory.ts

@@ -0,0 +1,35 @@
+export const DELETED_TOMBSTONE = "DELETED";
+
+export interface Memory {
+  id: string;
+  content: string;
+  embedding: number[];
+  metadata: Record<string, unknown>;
+  createdAt: Date;
+  updatedAt: Date;
+  supersededBy: string | null;
+}
+
+export interface VectorRow {
+  id: string;
+  distance: number;
+}
+
+export function isDeleted(memory: Memory): boolean {
+  return memory.supersededBy === DELETED_TOMBSTONE;
+}
+
+export function isSuperseded(memory: Memory): boolean {
+  return memory.supersededBy !== null;
+}
+
+export function memoryToDict(memory: Memory): Record<string, unknown> {
+  return {
+    id: memory.id,
+    content: memory.content,
+    metadata: memory.metadata,
+    createdAt: memory.createdAt.toISOString(),
+    updatedAt: memory.updatedAt.toISOString(),
+    supersededBy: memory.supersededBy,
+  };
+}