@coreidentitylabs/open-graph-memory-mcp 1.0.1

package/src/storage/neo4j-store.ts

@@ -0,0 +1,564 @@
+// =============================================================================
+// Neo4j Graph Database Storage Backend
+// =============================================================================
+
+import neo4j, { Driver, Session } from "neo4j-driver";
+import type {
+  StorageBackend,
+  MemoryNode,
+  MemoryEdge,
+  ScoredNode,
+  Subgraph,
+  NodeFilter,
+  GraphStats,
+} from "../types.js";
+import { cosineSimilarity } from "../encoding/embedder.js";
+
+export class Neo4jStore implements StorageBackend {
+  private driver: Driver;
+  private uri: string;
+  private user: string;
+  private password: string;
+
+  constructor(uri: string, user: string, password: string) {
+    this.uri = uri;
+    this.user = user;
+    this.password = password;
+    this.driver = neo4j.driver(
+      this.uri,
+      neo4j.auth.basic(this.user, this.password),
+    );
+  }
+
+  private getSession(): Session {
+    return this.driver.session();
+  }
+
+  async initialize(): Promise<void> {
+    const session = this.getSession();
+    try {
+      // Create constraints and indexes
+      await session.run(
+        `CREATE CONSTRAINT memory_node_id IF NOT EXISTS FOR (n:MemoryNode) REQUIRE n.id IS UNIQUE`,
+      );
+      await session.run(
+        `CREATE CONSTRAINT memory_edge_id IF NOT EXISTS FOR ()-[r:MEMORY_EDGE]-() REQUIRE r.id IS UNIQUE`,
+      );
+      await session.run(
+        `CREATE INDEX memory_node_name IF NOT EXISTS FOR (n:MemoryNode) ON (n.name)`,
+      );
+      await session.run(
+        `CREATE INDEX memory_node_type IF NOT EXISTS FOR (n:MemoryNode) ON (n.type)`,
+      );
+
+      // Verify connection
+      const result = await session.run(
+        `MATCH (n:MemoryNode) RETURN count(n) AS count`,
+      );
+      const count = result.records[0]?.get("count")?.toNumber() ?? 0;
+      console.error(
+        `[open-memory] Connected to Neo4j at ${this.uri}. ${count} existing nodes.`,
+      );
+    } finally {
+      await session.close();
+    }
+  }
+
+  // -- Node Operations -------------------------------------------------------
+
+  async addNode(node: MemoryNode): Promise<void> {
+    const session = this.getSession();
+    try {
+      await session.run(
+        `CREATE (n:MemoryNode {
+          id: $id, name: $name, type: $type, description: $description,
+          embedding: $embedding, metadata: $metadata,
+          createdAt: $createdAt, updatedAt: $updatedAt,
+          validFrom: $validFrom, validUntil: $validUntil,
+          source: $source, accessCount: $accessCount,
+          lastAccessedAt: $lastAccessedAt
+        })`,
+        {
+          ...node,
+          embedding: node.embedding ?? [],
+          metadata: JSON.stringify(node.metadata),
+          validFrom: node.validFrom ?? null,
+          validUntil: node.validUntil ?? null,
+          source: node.source ?? null,
+          lastAccessedAt: node.lastAccessedAt ?? null,
+        },
+      );
+    } finally {
+      await session.close();
+    }
+  }
+
+  async updateNode(
+    id: string,
+    updates: Partial<MemoryNode>,
+  ): Promise<MemoryNode | null> {
+    const session = this.getSession();
+    try {
+      const setClause = buildSetClause(updates);
+      if (!setClause) return await this.getNode(id);
+
+      const result = await session.run(
+        `MATCH (n:MemoryNode {id: $id})
+         SET ${setClause}, n.updatedAt = $updatedAt
+         RETURN n`,
+        {
+          id,
+          updatedAt: new Date().toISOString(),
+          ...serializeUpdates(updates),
+        },
+      );
+      if (result.records.length === 0) return null;
+      return deserializeNode(result.records[0].get("n").properties);
+    } finally {
+      await session.close();
+    }
+  }
+
+  async deleteNode(id: string): Promise<boolean> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (n:MemoryNode {id: $id}) DETACH DELETE n RETURN count(n) AS deleted`,
+        { id },
+      );
+      return (result.records[0]?.get("deleted")?.toNumber() ?? 0) > 0;
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getNode(id: string): Promise<MemoryNode | null> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (n:MemoryNode {id: $id}) RETURN n`,
+        { id },
+      );
+      if (result.records.length === 0) return null;
+      return deserializeNode(result.records[0].get("n").properties);
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getNodeByName(name: string): Promise<MemoryNode | null> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (n:MemoryNode) WHERE toLower(n.name) = toLower($name) RETURN n LIMIT 1`,
+        { name },
+      );
+      if (result.records.length === 0) return null;
+      return deserializeNode(result.records[0].get("n").properties);
+    } finally {
+      await session.close();
+    }
+  }
+
+  async findNodes(filter: NodeFilter): Promise<MemoryNode[]> {
+    const session = this.getSession();
+    try {
+      const conditions: string[] = [];
+      const params: Record<string, unknown> = {};
+
+      if (filter.type) {
+        conditions.push("n.type = $type");
+        params.type = filter.type;
+      }
+      if (filter.source) {
+        conditions.push("n.source = $source");
+        params.source = filter.source;
+      }
+      if (filter.nameContains) {
+        conditions.push("toLower(n.name) CONTAINS toLower($nameContains)");
+        params.nameContains = filter.nameContains;
+      }
+      if (filter.createdAfter) {
+        conditions.push("n.createdAt >= $createdAfter");
+        params.createdAfter = filter.createdAfter;
+      }
+      if (filter.createdBefore) {
+        conditions.push("n.createdAt <= $createdBefore");
+        params.createdBefore = filter.createdBefore;
+      }
+
+      const whereClause =
+        conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
+      const result = await session.run(
+        `MATCH (n:MemoryNode) ${whereClause} RETURN n ORDER BY n.updatedAt DESC`,
+        params,
+      );
+      return result.records.map((r) => deserializeNode(r.get("n").properties));
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getAllNodes(
+    limit: number,
+    offset: number,
+  ): Promise<{ nodes: MemoryNode[]; total: number }> {
+    const session = this.getSession();
+    try {
+      const countResult = await session.run(
+        `MATCH (n:MemoryNode) RETURN count(n) AS total`,
+      );
+      const total = countResult.records[0]?.get("total")?.toNumber() ?? 0;
+
+      const result = await session.run(
+        `MATCH (n:MemoryNode) RETURN n ORDER BY n.updatedAt DESC SKIP $offset LIMIT $limit`,
+        { offset: neo4j.int(offset), limit: neo4j.int(limit) },
+      );
+      const nodes = result.records.map((r) =>
+        deserializeNode(r.get("n").properties),
+      );
+      return { nodes, total };
+    } finally {
+      await session.close();
+    }
+  }
+
+  async findNodesByEmbedding(
+    embedding: number[],
+    topK: number,
+  ): Promise<ScoredNode[]> {
+    // For now, fall back to client-side cosine similarity
+    // Neo4j vector index can be used in production for better perf
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (n:MemoryNode) WHERE size(n.embedding) > 0 RETURN n`,
+      );
+      const scored: ScoredNode[] = [];
+      for (const record of result.records) {
+        const node = deserializeNode(record.get("n").properties);
+        if (node.embedding && node.embedding.length > 0) {
+          const score = cosineSimilarity(embedding, node.embedding);
+          scored.push({ node, score, matchType: "semantic" });
+        }
+      }
+      scored.sort((a, b) => b.score - a.score);
+      return scored.slice(0, topK);
+    } finally {
+      await session.close();
+    }
+  }
+
+  // -- Edge Operations --------------------------------------------------------
+
+  async addEdge(edge: MemoryEdge): Promise<void> {
+    const session = this.getSession();
+    try {
+      await session.run(
+        `MATCH (s:MemoryNode {id: $source})
+         MATCH (t:MemoryNode {id: $target})
+         CREATE (s)-[r:MEMORY_EDGE {
+           id: $id, relation: $relation, description: $description,
+           weight: $weight, metadata: $metadata,
+           createdAt: $createdAt, updatedAt: $updatedAt
+         }]->(t)`,
+        {
+          ...edge,
+          metadata: JSON.stringify(edge.metadata),
+        },
+      );
+    } finally {
+      await session.close();
+    }
+  }
+
+  async updateEdge(
+    id: string,
+    updates: Partial<MemoryEdge>,
+  ): Promise<MemoryEdge | null> {
+    const session = this.getSession();
+    try {
+      const setClause = buildEdgeSetClause(updates);
+      if (!setClause) return await this.getEdge(id);
+
+      const result = await session.run(
+        `MATCH ()-[r:MEMORY_EDGE {id: $id}]-()
+         SET ${setClause}, r.updatedAt = $updatedAt
+         RETURN r, startNode(r).id AS sourceId, endNode(r).id AS targetId`,
+        {
+          id,
+          updatedAt: new Date().toISOString(),
+          ...serializeEdgeUpdates(updates),
+        },
+      );
+      if (result.records.length === 0) return null;
+      return deserializeEdge(result.records[0]);
+    } finally {
+      await session.close();
+    }
+  }
+
+  async deleteEdge(id: string): Promise<boolean> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH ()-[r:MEMORY_EDGE {id: $id}]-() DELETE r RETURN count(r) AS deleted`,
+        { id },
+      );
+      return (result.records[0]?.get("deleted")?.toNumber() ?? 0) > 0;
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getEdge(id: string): Promise<MemoryEdge | null> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (s)-[r:MEMORY_EDGE {id: $id}]->(t) RETURN r, s.id AS sourceId, t.id AS targetId`,
+        { id },
+      );
+      if (result.records.length === 0) return null;
+      return deserializeEdge(result.records[0]);
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getEdgesForNode(
+    nodeId: string,
+    direction: "in" | "out" | "both" = "both",
+  ): Promise<MemoryEdge[]> {
+    const session = this.getSession();
+    try {
+      let query: string;
+      if (direction === "out") {
+        query = `MATCH (s:MemoryNode {id: $nodeId})-[r:MEMORY_EDGE]->(t) RETURN r, s.id AS sourceId, t.id AS targetId`;
+      } else if (direction === "in") {
+        query = `MATCH (s)-[r:MEMORY_EDGE]->(t:MemoryNode {id: $nodeId}) RETURN r, s.id AS sourceId, t.id AS targetId`;
+      } else {
+        query = `MATCH (n:MemoryNode {id: $nodeId})-[r:MEMORY_EDGE]-(m) RETURN r, startNode(r).id AS sourceId, endNode(r).id AS targetId`;
+      }
+      const result = await session.run(query, { nodeId });
+      return result.records.map((rec) => deserializeEdge(rec));
+    } finally {
+      await session.close();
+    }
+  }
+
+  async getEdgesBetween(
+    sourceId: string,
+    targetId: string,
+  ): Promise<MemoryEdge[]> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH (s:MemoryNode)-[r:MEMORY_EDGE]-(t:MemoryNode)
+         WHERE (s.id = $sourceId AND t.id = $targetId)
+            OR (s.id = $targetId AND t.id = $sourceId)
+         RETURN r, startNode(r).id AS sourceId, endNode(r).id AS targetId`,
+        { sourceId, targetId },
+      );
+      return result.records.map((rec) => deserializeEdge(rec));
+    } finally {
+      await session.close();
+    }
+  }
+
+  // -- Graph Traversal --------------------------------------------------------
+
+  async getNeighborhood(nodeId: string, depth: number): Promise<Subgraph> {
+    const session = this.getSession();
+    try {
+      const result = await session.run(
+        `MATCH path = (start:MemoryNode {id: $nodeId})-[*0..${depth}]-(neighbor:MemoryNode)
+         WITH DISTINCT neighbor, relationships(path) AS rels
+         UNWIND rels AS r
+         WITH collect(DISTINCT neighbor) AS nodes,
+              collect(DISTINCT {rel: r, source: startNode(r).id, target: endNode(r).id}) AS edges
+         RETURN nodes, edges`,
+        { nodeId },
+      );
+
+      if (result.records.length === 0) {
+        // Return just the start node if it exists
+        const startNode = await this.getNode(nodeId);
+        return {
+          nodes: startNode ? [startNode] : [],
+          edges: [],
+        };
+      }
+
+      const record = result.records[0];
+      const rawNodes = record.get("nodes") as unknown[];
+      const rawEdges = record.get("edges") as unknown[];
+
+      const nodes: MemoryNode[] = (rawNodes || []).map((n: unknown) =>
+        deserializeNode(
+          (n as { properties: Record<string, unknown> }).properties,
+        ),
+      );
+      const edges: MemoryEdge[] = (rawEdges || []).map((e: unknown) => {
+        const edge = e as {
+          rel: { properties: Record<string, unknown> };
+          source: string;
+          target: string;
+        };
+        return {
+          id: edge.rel.properties.id as string,
+          source: edge.source,
+          target: edge.target,
+          relation: edge.rel.properties.relation as string,
+          description: edge.rel.properties.description as string,
+          weight: edge.rel.properties.weight as number,
+          metadata: JSON.parse(
+            (edge.rel.properties.metadata as string) || "{}",
+          ),
+          createdAt: edge.rel.properties.createdAt as string,
+          updatedAt: edge.rel.properties.updatedAt as string,
+        };
+      });
+
+      return { nodes, edges };
+    } finally {
+      await session.close();
+    }
+  }
+
+  // -- Stats & Lifecycle ------------------------------------------------------
+
+  async getStats(): Promise<GraphStats> {
+    const session = this.getSession();
+    try {
+      const nodesResult = await session.run(
+        `MATCH (n:MemoryNode) RETURN n.type AS type, count(n) AS count`,
+      );
+      const nodesByType: Record<string, number> = {};
+      let totalNodes = 0;
+      for (const record of nodesResult.records) {
+        const type = record.get("type") as string;
+        const count = (
+          record.get("count") as { toNumber(): number }
+        ).toNumber();
+        nodesByType[type] = count;
+        totalNodes += count;
+      }
+
+      const edgesResult = await session.run(
+        `MATCH ()-[r:MEMORY_EDGE]-() RETURN count(r) AS total`,
+      );
+      const totalEdges =
+        (
+          edgesResult.records[0]?.get("total") as { toNumber(): number }
+        )?.toNumber() ?? 0;
+
+      return {
+        totalNodes,
+        totalEdges,
+        nodesByType,
+        storageBackend: "neo4j",
+      };
+    } finally {
+      await session.close();
+    }
+  }
+
+  async close(): Promise<void> {
+    await this.driver.close();
+  }
+}
+
+// =============================================================================
+// Serialization Helpers
+// =============================================================================
+
+function deserializeNode(props: Record<string, unknown>): MemoryNode {
+  return {
+    id: props.id as string,
+    name: props.name as string,
+    type: props.type as MemoryNode["type"],
+    description: props.description as string,
+    embedding: (props.embedding as number[]) ?? [],
+    metadata:
+      typeof props.metadata === "string"
+        ? JSON.parse(props.metadata)
+        : ((props.metadata as Record<string, unknown>) ?? {}),
+    createdAt: props.createdAt as string,
+    updatedAt: props.updatedAt as string,
+    validFrom: (props.validFrom as string) ?? undefined,
+    validUntil: (props.validUntil as string) ?? undefined,
+    source: (props.source as string) ?? undefined,
+    accessCount:
+      typeof props.accessCount === "object" && props.accessCount !== null
+        ? (props.accessCount as { toNumber(): number }).toNumber()
+        : ((props.accessCount as number) ?? 0),
+    lastAccessedAt: (props.lastAccessedAt as string) ?? undefined,
+  };
+}
+
+function deserializeEdge(record: { get(key: string): unknown }): MemoryEdge {
+  const rel = record.get("r") as { properties: Record<string, unknown> };
+  const props = rel.properties;
+  return {
+    id: props.id as string,
+    source: record.get("sourceId") as string,
+    target: record.get("targetId") as string,
+    relation: props.relation as string,
+    description: props.description as string,
+    weight: props.weight as number,
+    metadata:
+      typeof props.metadata === "string"
+        ? JSON.parse(props.metadata)
+        : ((props.metadata as Record<string, unknown>) ?? {}),
+    createdAt: props.createdAt as string,
+    updatedAt: props.updatedAt as string,
+  };
+}
+
+function buildSetClause(updates: Partial<MemoryNode>): string {
+  const parts: string[] = [];
+  if (updates.name !== undefined) parts.push("n.name = $name");
+  if (updates.description !== undefined)
+    parts.push("n.description = $description");
+  if (updates.type !== undefined) parts.push("n.type = $type");
+  if (updates.embedding !== undefined) parts.push("n.embedding = $embedding");
+  if (updates.metadata !== undefined) parts.push("n.metadata = $metadata");
+  if (updates.validFrom !== undefined) parts.push("n.validFrom = $validFrom");
+  if (updates.validUntil !== undefined)
+    parts.push("n.validUntil = $validUntil");
+  if (updates.accessCount !== undefined)
+    parts.push("n.accessCount = $accessCount");
+  if (updates.lastAccessedAt !== undefined)
+    parts.push("n.lastAccessedAt = $lastAccessedAt");
+  return parts.join(", ");
+}
+
+function buildEdgeSetClause(updates: Partial<MemoryEdge>): string {
+  const parts: string[] = [];
+  if (updates.relation !== undefined) parts.push("r.relation = $relation");
+  if (updates.description !== undefined)
+    parts.push("r.description = $description");
+  if (updates.weight !== undefined) parts.push("r.weight = $weight");
+  if (updates.metadata !== undefined) parts.push("r.metadata = $metadata");
+  return parts.join(", ");
+}
+
+function serializeUpdates(
+  updates: Partial<MemoryNode>,
+): Record<string, unknown> {
+  const result: Record<string, unknown> = { ...updates };
+  if (updates.metadata !== undefined) {
+    result.metadata = JSON.stringify(updates.metadata);
+  }
+  return result;
+}
+
+function serializeEdgeUpdates(
+  updates: Partial<MemoryEdge>,
+): Record<string, unknown> {
+  const result: Record<string, unknown> = { ...updates };
+  if (updates.metadata !== undefined) {
+    result.metadata = JSON.stringify(updates.metadata);
+  }
+  return result;
+}