fix graph import

.github/workflows/tests.yml

@@ -94,13 +94,7 @@ jobs:
           wget https://github.com/milvus-io/milvus/releases/download/v2.4.4/milvus-standalone-docker-compose.yml -O docker-compose.yml
           sudo docker compose up -d
           docker pull neo4j
-          docker run \
-            --name neo4j \
-            -d \
-            -p 7474:7474 \
-            -p 7687:7687 \
-            -e NEO4J_AUTH=none \
-            neo4j
+          docker run --name neo4j -d -p 7474:7474 -p 7687:7687 -e NEO4J_AUTH=none neo4j
 
       - name: Test with pytest
         run: |

src/cardinal/graph/neo4j.py

@@ -1,14 +1,20 @@
-import json
-from typing import Generic, Optional, Sequence, TypeVar
+from typing import Optional, Sequence, TypeVar
 from pydantic import BaseModel
-from neo4j import GraphDatabase
+
+from .schema import GraphStorage
+from .config import settings
+from ..utils.import_utils import is_neo4j_available
+
+
+if is_neo4j_available():
+    from neo4j import GraphDatabase
 
 T = TypeVar("T", bound=BaseModel)
 
-class Neo4j(Generic[T]):
+class Neo4j(GraphStorage[T]):
     def __init__(self, name: str) -> None:
         self.name = name
-        self.driver = GraphDatabase.driver("bolt://localhost:7687")
+        self.driver = GraphDatabase.driver(settings.neo4j_uri)
 
     def insert_node(self, key: Sequence[str], node: Sequence[T]) -> None:
         with self.driver.session() as session:

src/cardinal/utils/import_utils.py

@@ -22,6 +22,10 @@ def is_redis_available():
     return _is_package_available("redis")
 
 
+def is_neo4j_available():
+    return _is_package_available("neo4j")
+
+
 def is_tiktoken_available():
     return _is_package_available("tiktoken")