Support exporting nodes directly in Neo4j. (#316)

Author: badmonster0

Cargo.toml

@@ -104,3 +104,4 @@ async-stream = "0.3.6"
 neo4rs = "0.8.0"
 bytes = "1.10.1"
 rand = "0.9.0"
+indoc = "2.0.6"

examples/docs_to_kg/main.py

@@ -21,7 +21,7 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
 
     conn_spec = cocoindex.add_auth_entry(
         "Neo4jConnection",
-        cocoindex.storages.Neo4jConnectionSpec(
+        cocoindex.storages.Neo4jConnection(
             uri="bolt://localhost:7687",
             user="neo4j",
             password="cocoindex",
@@ -70,38 +70,40 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
 
     relationships.export(
         "relationships",
-        cocoindex.storages.Neo4jRelationship(
+        cocoindex.storages.Neo4j(
             connection=conn_spec,
-            rel_type="RELATIONSHIP",
-            source=cocoindex.storages.Neo4jRelationshipEndSpec(
-                label="Entity",
-                fields=[
-                    cocoindex.storages.Neo4jFieldMapping(
-                        field_name="subject", node_field_name="value"),
-                    cocoindex.storages.Neo4jFieldMapping(
-                        field_name="subject_embedding", node_field_name="embedding"),
-                ]
-            ),
-            target=cocoindex.storages.Neo4jRelationshipEndSpec(
-                label="Entity",
-                fields=[
-                    cocoindex.storages.Neo4jFieldMapping(
-                        field_name="object", node_field_name="value"),
-                    cocoindex.storages.Neo4jFieldMapping(
-                        field_name="object_embedding", node_field_name="embedding"),
-                ]
-            ),
-            nodes={
-                "Entity": cocoindex.storages.Neo4jRelationshipNodeSpec(
-                    primary_key_fields=["value"],
-                    vector_indexes=[
-                        cocoindex.VectorIndexDef(
-                            field_name="embedding",
-                            metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
-                        ),
-                    ],
+            mapping=cocoindex.storages.Neo4jRelationship(
+                rel_type="RELATIONSHIP",
+                source=cocoindex.storages.Neo4jRelationshipEnd(
+                    label="Entity",
+                    fields=[
+                        cocoindex.storages.Neo4jFieldMapping(
+                            field_name="subject", node_field_name="value"),
+                        cocoindex.storages.Neo4jFieldMapping(
+                            field_name="subject_embedding", node_field_name="embedding"),
+                    ]
                 ),
-            },
+                target=cocoindex.storages.Neo4jRelationshipEnd(
+                    label="Entity",
+                    fields=[
+                        cocoindex.storages.Neo4jFieldMapping(
+                            field_name="object", node_field_name="value"),
+                        cocoindex.storages.Neo4jFieldMapping(
+                            field_name="object_embedding", node_field_name="embedding"),
+                    ]
+                ),
+                nodes={
+                    "Entity": cocoindex.storages.Neo4jRelationshipNode(
+                        primary_key_fields=["value"],
+                        vector_indexes=[
+                            cocoindex.VectorIndexDef(
+                                field_name="embedding",
+                                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+                            ),
+                        ],
+                    ),
+                },
+            ),
         ),
         primary_key_fields=["id"],
     )

python/cocoindex/convert.py

@@ -117,7 +117,10 @@ def dump_engine_object(v: Any) -> Any:
         nanos = int((total_secs - secs) * 1e9)
         return {'secs': secs, 'nanos': nanos}
     elif hasattr(v, '__dict__'):
-        return {k: dump_engine_object(v) for k, v in v.__dict__.items()}
+        s = {k: dump_engine_object(v) for k, v in v.__dict__.items()}
+        if hasattr(v, 'kind') and 'kind' not in s:
+            s['kind'] = v.kind
+        return s
     elif isinstance(v, (list, tuple)):
         return [dump_engine_object(item) for item in v]
     elif isinstance(v, dict):

python/cocoindex/storages.py

@@ -21,7 +21,7 @@ class Qdrant(op.StorageSpec):
     api_key: str | None = None
 
 @dataclass
-class Neo4jConnectionSpec:
+class Neo4jConnection:
     """Connection spec for Neo4j."""
     uri: str
     user: str
@@ -37,22 +37,36 @@ class Neo4jFieldMapping:
     node_field_name: str | None = None
 
 @dataclass
-class Neo4jRelationshipEndSpec:
+class Neo4jRelationshipEnd:
     """Spec for a Neo4j node type."""
     label: str
     fields: list[Neo4jFieldMapping]
 
 @dataclass
-class Neo4jRelationshipNodeSpec:
+class Neo4jRelationshipNode:
     """Spec for a Neo4j node type."""
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()
 
-class Neo4jRelationship(op.StorageSpec):
+@dataclass
+class Neo4jNode:
+    """Spec for a Neo4j node type."""
+    kind = "Node"
+
+    label: str
+
+@dataclass
+class Neo4jRelationship:
+    """Spec for a Neo4j relationship."""
+    kind = "Relationship"
+
+    rel_type: str
+    source: Neo4jRelationshipEnd
+    target: Neo4jRelationshipEnd
+    nodes: dict[str, Neo4jRelationshipNode]
+
+class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
 
     connection: AuthEntryReference
-    rel_type: str
-    source: Neo4jRelationshipEndSpec
-    target: Neo4jRelationshipEndSpec
-    nodes: dict[str, Neo4jRelationshipNodeSpec]
+    mapping: Neo4jNode | Neo4jRelationship

src/ops/registration.rs

@@ -15,8 +15,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
     Arc::new(storages::postgres::Factory::default()).register(registry)?;
     Arc::new(storages::qdrant::Factory::default()).register(registry)?;
 
-    let neo4j_pool = Arc::new(storages::neo4j::GraphPool::default());
-    storages::neo4j::RelationshipFactory::new(neo4j_pool).register(registry)?;
+    storages::neo4j::Factory::new().register(registry)?;
 
     Ok(())
 }