From e0aa7c9a61ea2de40d519c6f4b249032395a686f Mon Sep 17 00:00:00 2001
From: Hande <159312713+hande-k@users.noreply.github.com>
Date: Mon, 2 Jun 2025 22:38:23 +0200
Subject: [PATCH] fix: memgraph adapter (#899)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
---
 .github/workflows/test_memgraph.yml           | 26 ++++++++++---------
 .../graph/memgraph/memgraph_adapter.py        | 21 ++++++++++++++-
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/test_memgraph.yml b/.github/workflows/test_memgraph.yml
index b7ea9d837..c9134bc1f 100644
--- a/.github/workflows/test_memgraph.yml
+++ b/.github/workflows/test_memgraph.yml
@@ -1,9 +1,9 @@
 name: test | memgraph
 
-# on:
-#  workflow_dispatch:
-#  pull_request:
-#    types: [labeled, synchronize]
+on:
+  workflow_dispatch:
+  pull_request:
+    types: [labeled, synchronize]
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
@@ -17,13 +17,15 @@ jobs:
     name: test
     runs-on: ubuntu-22.04
 
-    defaults:
-      run:
-        shell: bash
+    services:
+      memgraph:
+        image: memgraph/memgraph-mage:latest
+        ports:
+          - 7687:7687
 
     steps:
       - name: Check out
-        uses: actions/checkout@master
+        uses: actions/checkout@v3
 
       - name: Setup Python
         uses: actions/setup-python@v5
@@ -38,7 +40,7 @@ jobs:
           installer-parallel: true
 
       - name: Install dependencies
-        run: poetry install -E memgraph --no-interaction
+        run: poetry install -E neo4j
 
       - name: Run default Memgraph
         env:
@@ -51,7 +53,7 @@ jobs:
           EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
           EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
           EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
-          GRAPH_DATABASE_URL: ${{ secrets.MEMGRAPH_API_URL }}
-          GRAPH_DATABASE_PASSWORD: ${{ secrets.MEMGRAPH_API_KEY }}
-          GRAPH_DATABASE_USERNAME: " "
+          GRAPH_DATABASE_URL: "bolt://localhost:7687"
+          GRAPH_DATABASE_PASSWORD: "memgraph"
+          GRAPH_DATABASE_USERNAME: "memgraph"
         run: poetry run python ./cognee/tests/test_memgraph.py
diff --git a/cognee/infrastructure/databases/graph/memgraph/memgraph_adapter.py b/cognee/infrastructure/databases/graph/memgraph/memgraph_adapter.py
index 75c0722a8..bd2cffeb8 100644
--- a/cognee/infrastructure/databases/graph/memgraph/memgraph_adapter.py
+++ b/cognee/infrastructure/databases/graph/memgraph/memgraph_adapter.py
@@ -614,7 +614,7 @@ class MemgraphAdapter(GraphDBInterface):
 
             return [result["successor"] for result in results]
 
-    async def get_neighbours(self, node_id: str) -> List[Dict[str, Any]]:
+    async def get_neighbors(self, node_id: str) -> List[Dict[str, Any]]:
         """
         Get both predecessors and successors of a node.
 
@@ -634,6 +634,25 @@ class MemgraphAdapter(GraphDBInterface):
 
         return predecessors + successors
 
+    async def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
+        """Get a single node by ID."""
+        query = """
+        MATCH (node {id: $node_id})
+        RETURN node
+        """
+        results = await self.query(query, {"node_id": node_id})
+        return results[0]["node"] if results else None
+
+    async def get_nodes(self, node_ids: List[str]) -> List[Dict[str, Any]]:
+        """Get multiple nodes by their IDs."""
+        query = """
+        UNWIND $node_ids AS id
+        MATCH (node {id: id})
+        RETURN node
+        """
+        results = await self.query(query, {"node_ids": node_ids})
+        return [result["node"] for result in results]
+
     async def get_connections(self, node_id: UUID) -> list:
         """
         Retrieve connections for a given node, including both predecessors and successors.