feat(attachment): add interface for target attachment

Author: georgeh0

python/cocoindex/flow.py

@@ -405,6 +405,7 @@ def export(
         /,
         *,
         primary_key_fields: Sequence[str],
+        attachments: Sequence[op.TargetAttachmentSpec],
         vector_indexes: Sequence[index.VectorIndexDef] = (),
         vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
         setup_by_user: bool = False,
@@ -436,6 +437,7 @@ def export(
             target_name,
             _spec_kind(target_spec),
             dump_engine_object(target_spec),
+            dump_engine_object(attachments),
             dump_engine_object(index_options),
             self._engine_data_collector,
             setup_by_user,

python/cocoindex/op.py

@@ -46,6 +46,7 @@ class OpCategory(Enum):
     SOURCE = "source"
     TARGET = "target"
     DECLARATION = "declaration"
+    TARGET_ATTACHMENT = "target_attachment"
 
 
 @dataclass_transform()
@@ -81,6 +82,10 @@ class TargetSpec(metaclass=SpecMeta, category=OpCategory.TARGET):  # pylint: dis
     """A target spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
 
+class TargetAttachmentSpec(metaclass=SpecMeta, category=OpCategory.TARGET_ATTACHMENT):  # pylint: disable=too-few-public-methods
+    """A target attachment spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
+
+
 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, ...)"""
 

src/base/spec.rs

@@ -493,6 +493,10 @@ impl fmt::Display for IndexOptions {
 pub struct ExportOpSpec {
     pub collector_name: FieldName,
     pub target: OpSpec,
+
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub attachments: Vec<OpSpec>,
+
     pub index_options: IndexOptions,
     pub setup_by_user: bool,
 }

src/builder/flow_builder.rs

@@ -549,12 +549,13 @@ impl FlowBuilder {
         Ok(())
     }
 
-    #[pyo3(signature = (name, kind, op_spec, index_options, input, setup_by_user=false))]
+    #[pyo3(signature = (name, kind, op_spec, attachments, index_options, input, setup_by_user=false))]
     pub fn export(
         &mut self,
         name: String,
         kind: String,
         op_spec: py::Pythonized<serde_json::Map<String, serde_json::Value>>,
+        attachments: py::Pythonized<Vec<spec::OpSpec>>,
         index_options: py::Pythonized<spec::IndexOptions>,
         input: &DataCollector,
         setup_by_user: bool,
@@ -574,6 +575,7 @@ impl FlowBuilder {
             spec: spec::ExportOpSpec {
                 collector_name: input.name.clone(),
                 target: spec,
+                attachments: attachments.into_inner(),
                 index_options: index_options.into_inner(),
                 setup_by_user,
             },

src/ops/interface.rs

@@ -316,9 +316,42 @@ pub trait TargetFactory: Send + Sync {
     ) -> Result<()>;
 }
 
+pub struct TargetAttachmentState {
+    pub setup_key: serde_json::Value,
+    pub setup_state: serde_json::Value,
+}
+
+#[async_trait]
+pub trait TargetAttachmentFactory: Send + Sync {
+    /// Normalize the key. e.g. the JSON format may change (after code change, e.g. new optional field or field ordering), even if the underlying value is not changed.
+    /// This should always return the canonical serialized form.
+    fn normalize_setup_key(&self, key: &serde_json::Value) -> Result<serde_json::Value>;
+
+    fn get_state(
+        &self,
+        target_name: &str,
+        target_spec: &spec::OpSpec,
+        attachment_spec: serde_json::Value,
+    ) -> Result<TargetAttachmentState>;
+
+    fn describe(&self, key: &serde_json::Value, state: &serde_json::Value) -> Result<String>;
+
+    fn has_update(&self, old: &serde_json::Value, new: &serde_json::Value) -> bool {
+        old != new
+    }
+
+    async fn apply_setup_change(
+        &self,
+        key: &serde_json::Value,
+        old: Option<serde_json::Value>,
+        new: Option<serde_json::Value>,
+    ) -> Result<()>;
+}
+
 #[derive(Clone)]
 pub enum ExecutorFactory {
     Source(Arc<dyn SourceFactory + Send + Sync>),
     SimpleFunction(Arc<dyn SimpleFunctionFactory + Send + Sync>),
     ExportTarget(Arc<dyn TargetFactory + Send + Sync>),
+    TargetAttachment(Arc<dyn TargetAttachmentFactory + Send + Sync>),
 }

src/ops/registry.rs

@@ -8,6 +8,8 @@ pub struct ExecutorFactoryRegistry {
     function_factories:
         HashMap<String, Arc<dyn super::interface::SimpleFunctionFactory + Send + Sync>>,
     target_factories: HashMap<String, Arc<dyn super::interface::TargetFactory + Send + Sync>>,
+    target_attachment_factories:
+        HashMap<String, Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>>,
 }
 
 impl Default for ExecutorFactoryRegistry {
@@ -22,6 +24,7 @@ impl ExecutorFactoryRegistry {
             source_factories: HashMap::new(),
             function_factories: HashMap::new(),
             target_factories: HashMap::new(),
+            target_attachment_factories: HashMap::new(),
         }
     }
 
@@ -61,6 +64,18 @@ impl ExecutorFactoryRegistry {
                     }
                 }
             }
+            ExecutorFactory::TargetAttachment(target_attachment_factory) => {
+                match self.target_attachment_factories.entry(name) {
+                    std::collections::hash_map::Entry::Occupied(entry) => Err(anyhow::anyhow!(
+                        "Target attachment factory with name already exists: {}",
+                        entry.key()
+                    )),
+                    std::collections::hash_map::Entry::Vacant(entry) => {
+                        entry.insert(target_attachment_factory);
+                        Ok(())
+                    }
+                }
+            }
         }
     }