feat(neo4j)!: add ReferencedNode as declaration (#374)

Author: badmonster0

examples/docs_to_knowledge_graph/main.py

@@ -92,13 +92,31 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
                     id=cocoindex.GeneratedField.UUID, entity=relationship["object"],
                     filename=doc["filename"], location=chunk["location"],
                 )
+
     document_node.export(
         "document_node",
         cocoindex.storages.Neo4j(
             connection=conn_spec,
             mapping=cocoindex.storages.NodeMapping(label="Document")),
         primary_key_fields=["filename"],
     )
+    flow_builder.declare(
+        cocoindex.storages.Neo4jDeclarations(
+            connection=conn_spec,
+            referenced_nodes=[
+                cocoindex.storages.ReferencedNode(
+                    label="Entity",
+                    primary_key_fields=["value"],
+                    vector_indexes=[
+                        cocoindex.VectorIndexDef(
+                            field_name="embedding",
+                            metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+                        ),
+                    ],
+                )
+            ]
+        )
+    )
     entity_relationship.export(
         "entity_relationship",
         cocoindex.storages.Neo4j(
@@ -123,17 +141,6 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
                             source="object_embedding", target="embedding"),
                     ]
                 ),
-                nodes_storage_spec={
-                    "Entity": cocoindex.storages.NodeStorageSpec(
-                        primary_key_fields=["value"],
-                        vector_indexes=[
-                            cocoindex.VectorIndexDef(
-                                field_name="embedding",
-                                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
-                            ),
-                        ],
-                    ),
-                },
             ),
         ),
         primary_key_fields=["id"],

python/cocoindex/flow.py

@@ -290,12 +290,6 @@ def export(self, name: str, target_spec: op.StorageSpec, /, *,
             name, _spec_kind(target_spec), dump_engine_object(target_spec),
             dump_engine_object(index_options), self._engine_data_collector, setup_by_user)
 
-    def declare(self, spec: op.DeclarationSpec):
-        """
-        Add a declaration to the flow.
-        """
-        self._flow_builder_state.engine_flow_builder.declare(dump_engine_object(spec))
-
 
 _flow_name_builder = _NameBuilder()
 
@@ -361,6 +355,12 @@ def add_source(self, spec: op.SourceSpec, /, *,
             name
         )
 
+    def declare(self, spec: op.DeclarationSpec):
+        """
+        Add a declaration to the flow.
+        """
+        self._state.engine_flow_builder.declare(dump_engine_object(spec))
+
 @dataclass
 class FlowLiveUpdaterOptions:
     """

python/cocoindex/op.py

@@ -7,7 +7,6 @@
 
 from typing import get_type_hints, Protocol, Any, Callable, Awaitable, dataclass_transform
 from enum import Enum
-from functools import partial
 
 from .typing import encode_enriched_type
 from .convert import to_engine_value, make_engine_value_converter
@@ -43,8 +42,6 @@ class StorageSpec(metaclass=SpecMeta, category=OpCategory.STORAGE): # pylint: di
 
 class DeclarationSpec(metaclass=SpecMeta, category=OpCategory.DECLARATION): # pylint: disable=too-few-public-methods
     """A declaration spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
-    kind: str
-
 class Executor(Protocol):
     """An executor for an operation."""
     op_category: OpCategory

python/cocoindex/storages.py

@@ -43,8 +43,9 @@ class NodeReferenceMapping:
     fields: list[TargetFieldMapping]
 
 @dataclass
-class NodeStorageSpec:
+class ReferencedNode:
     """Storage spec for a graph node."""
+    label: str
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()
 
@@ -63,10 +64,16 @@ class RelationshipMapping:
     rel_type: str
     source: NodeReferenceMapping
     target: NodeReferenceMapping
-    nodes_storage_spec: dict[str, NodeStorageSpec] | None = None
 
 class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
 
     connection: AuthEntryReference
     mapping: NodeMapping | RelationshipMapping
+
+class Neo4jDeclarations(op.DeclarationSpec):
+    """Declarations for Neo4j."""
+
+    kind = "Neo4j"
+    connection: AuthEntryReference
+    referenced_nodes: Sequence[ReferencedNode] = ()