feat(kg): rename spec type names

Author: badmonster0

examples/docs_to_knowledge_graph/main.py

@@ -93,35 +93,31 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
         "document_node",
         cocoindex.storages.Neo4j(
             connection=conn_spec,
-            mapping=cocoindex.storages.NodeMapping(label="Document")),
+            mapping=cocoindex.storages.Nodes(label="Document")),
         primary_key_fields=["filename"],
     )
     # Declare reference Node to reference entity node in a relationship
     flow_builder.declare(
-        cocoindex.storages.Neo4jDeclarations(
+        cocoindex.storages.Neo4jDeclaration(
             connection=conn_spec,
-            referenced_nodes=[
-                cocoindex.storages.ReferencedNode(
-                    label="Entity",
-                    primary_key_fields=["value"],
-                )
-            ]
+            nodes_label="Entity",
+            primary_key_fields=["value"],
         )
     )
     entity_relationship.export(
         "entity_relationship",
         cocoindex.storages.Neo4j(
             connection=conn_spec,
-            mapping=cocoindex.storages.RelationshipMapping(
+            mapping=cocoindex.storages.Relationships(
                 rel_type="RELATIONSHIP",
-                source=cocoindex.storages.NodeReferenceMapping(
+                source=cocoindex.storages.NodeFromFields(
                     label="Entity",
                     fields=[
                         cocoindex.storages.TargetFieldMapping(
                             source="subject", target="value"),
                     ]
                 ),
-                target=cocoindex.storages.NodeReferenceMapping(
+                target=cocoindex.storages.NodeFromFields(
                     label="Entity",
                     fields=[
                         cocoindex.storages.TargetFieldMapping(
@@ -136,13 +132,13 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
         "entity_mention",
         cocoindex.storages.Neo4j(
             connection=conn_spec,
-            mapping=cocoindex.storages.RelationshipMapping(
+            mapping=cocoindex.storages.Relationships(
                 rel_type="MENTION",
-                source=cocoindex.storages.NodeReferenceMapping(
+                source=cocoindex.storages.NodeFromFields(
                     label="Document",
                     fields=[cocoindex.storages.TargetFieldMapping("filename")],
                 ),
-                target=cocoindex.storages.NodeReferenceMapping(
+                target=cocoindex.storages.NodeFromFields(
                     label="Entity",
                     fields=[cocoindex.storages.TargetFieldMapping(
                         source="entity", target="value")],

python/cocoindex/storages.py

@@ -37,7 +37,7 @@ class TargetFieldMapping:
     target: str | None = None
 
 @dataclass
-class NodeReferenceMapping:
+class NodeFromFields:
     """Spec for a referenced graph node, usually as part of a relationship."""
     label: str
     fields: list[TargetFieldMapping]
@@ -50,30 +50,37 @@ class ReferencedNode:
     vector_indexes: Sequence[index.VectorIndexDef] = ()
 
 @dataclass
-class NodeMapping:
+class Nodes:
     """Spec to map a row to a graph node."""
     kind = "Node"
 
     label: str
 
 @dataclass
-class RelationshipMapping:
+class Relationships:
     """Spec to map a row to a graph relationship."""
     kind = "Relationship"
 
     rel_type: str
-    source: NodeReferenceMapping
-    target: NodeReferenceMapping
+    source: NodeFromFields
+    target: NodeFromFields
+
+# For backwards compatibility only
+NodeMapping = Nodes
+RelationshipMapping = Relationships
+NodeReferenceMapping = NodeFromFields
 
 class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
 
     connection: AuthEntryReference
-    mapping: NodeMapping | RelationshipMapping
+    mapping: Nodes | Relationships
 
-class Neo4jDeclarations(op.DeclarationSpec):
+class Neo4jDeclaration(op.DeclarationSpec):
     """Declarations for Neo4j."""
 
     kind = "Neo4j"
     connection: AuthEntryReference
-    referenced_nodes: Sequence[ReferencedNode] = ()
+    nodes_label: str
+    primary_key_fields: Sequence[str]
+    vector_indexes: Sequence[index.VectorIndexDef] = ()

src/ops/storages/neo4j.rs

@@ -1,8 +1,6 @@
 use crate::prelude::*;
 
-use super::spec::{
-    GraphDeclarations, GraphElementMapping, NodeReferenceMapping, TargetFieldMapping,
-};
+use super::spec::{GraphDeclaration, GraphElementMapping, NodeFromFieldsSpec, TargetFieldMapping};
 use crate::setup::components::{self, State};
 use crate::setup::{ResourceSetupStatusCheck, SetupChangeType};
 use crate::{ops::sdk::*, setup::CombinedState};
@@ -32,7 +30,7 @@ pub struct Spec {
 pub struct Declaration {
     connection: spec::AuthEntryReference<ConnectionSpec>,
     #[serde(flatten)]
-    decl: GraphDeclarations,
+    decl: GraphDeclaration,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
@@ -994,7 +992,7 @@ struct DependentNodeLabelAnalyzer<'a> {
 
 impl<'a> DependentNodeLabelAnalyzer<'a> {
     fn new(
-        rel_end_spec: &'a NodeReferenceMapping,
+        rel_end_spec: &'a NodeFromFieldsSpec,
         index_options_map: &'a HashMap<String, IndexOptions>,
     ) -> Result<Self> {
         Ok(Self {
@@ -1104,8 +1102,7 @@ impl StorageFactoryBase for Factory {
             .chain(
                 declarations
                     .iter()
-                    .flat_map(|d| d.decl.referenced_nodes.iter())
-                    .map(|n| (n.label.clone(), n.index_options.clone())),
+                    .map(|n| (n.decl.nodes_label.clone(), n.decl.index_options.clone())),
             )
             .collect::<HashMap<String, IndexOptions>>();
         let mut label_value_field_types =
@@ -1212,37 +1209,41 @@ impl StorageFactoryBase for Factory {
             .collect::<Result<Vec<_>>>()?;
         let decl_output = declarations
             .into_iter()
-            .flat_map(|d| {
+            .map(|d| {
                 let label_value_field_types = &label_value_field_types;
-                d.decl.referenced_nodes.into_iter().map(move |n| {
-                    let setup_key = GraphElement {
-                        connection: d.connection.clone(),
-                        typ: ElementType::Node(n.label.clone()),
-                    };
-                    let primary_key_fields = n
-                        .index_options
-                        .primary_key_fields
-                        .as_ref()
+                let setup_key = GraphElement {
+                    connection: d.connection.clone(),
+                    typ: ElementType::Node(d.decl.nodes_label.clone()),
+                };
+                let primary_key_fields = d
+                    .decl
+                    .index_options
+                    .primary_key_fields
+                    .as_ref()
+                    .ok_or_else(|| {
+                        api_error!(
+                            "No primary key fields specified for node label `{}`",
+                            &d.decl.nodes_label
+                        )
+                    })?
+                    .iter()
+                    .map(|f| f.clone())
+                    .collect();
+                let setup_state = SetupState::new(
+                    &setup_key.typ,
+                    primary_key_fields,
+                    &d.decl.index_options,
+                    label_value_field_types
+                        .get(&d.decl.nodes_label)
                         .ok_or_else(|| {
                             api_error!(
-                                "No primary key fields specified for node label `{}`",
-                                &n.label
+                                "Data for nodes with label `{}` not provided",
+                                d.decl.nodes_label
                             )
-                        })?
-                        .iter()
-                        .map(|f| f.clone())
-                        .collect();
-                    let setup_state = SetupState::new(
-                        &setup_key.typ,
-                        primary_key_fields,
-                        &n.index_options,
-                        label_value_field_types.get(&n.label).ok_or_else(|| {
-                            api_error!("Data for nodes with label `{}` not provided", n.label)
                         })?,
-                        None,
-                    )?;
-                    Ok((setup_key, setup_state))
-                })
+                    None,
+                )?;
+                Ok((setup_key, setup_state))
             })
             .collect::<Result<Vec<_>>>()?;
         Ok((data_coll_output, decl_output))

src/ops/storages/spec.rs

@@ -17,39 +17,34 @@ impl TargetFieldMapping {
 }
 
 #[derive(Debug, Deserialize)]
-pub struct NodeReferenceMapping {
+pub struct NodeFromFieldsSpec {
     pub label: String,
     pub fields: Vec<TargetFieldMapping>,
 }
 
 #[derive(Debug, Deserialize)]
-pub struct NodeMapping {
+pub struct NodesSpec {
     pub label: String,
 }
 
 #[derive(Debug, Deserialize)]
-pub struct RelationshipMapping {
+pub struct RelationshipsSpec {
     pub rel_type: String,
-    pub source: NodeReferenceMapping,
-    pub target: NodeReferenceMapping,
+    pub source: NodeFromFieldsSpec,
+    pub target: NodeFromFieldsSpec,
 }
 
 #[derive(Debug, Deserialize)]
 #[serde(tag = "kind")]
 pub enum GraphElementMapping {
-    Relationship(RelationshipMapping),
-    Node(NodeMapping),
+    Relationship(RelationshipsSpec),
+    Node(NodesSpec),
 }
 
 #[derive(Debug, Deserialize)]
-pub struct ReferencedNodeSpec {
-    pub label: String,
+pub struct GraphDeclaration {
+    pub nodes_label: String,
 
     #[serde(flatten)]
     pub index_options: spec::IndexOptions,
 }
-
-#[derive(Debug, Deserialize)]
-pub struct GraphDeclarations {
-    pub referenced_nodes: Vec<ReferencedNodeSpec>,
-}