feat: support transform for FlowBuilder

lemorage · lemorage · commit 7501bc5160df · 2025-07-01T08:03:42.000+02:00
diff --git a/python/cocoindex/flow.py b/python/cocoindex/flow.py
@@ -5,25 +5,25 @@
 from __future__ import annotations
 
 import asyncio
-import re
-import inspect
 import datetime
 import functools
-
+import inspect
+import re
+from dataclasses import dataclass
+from enum import Enum
+from threading import Lock
 from typing import (
     Any,
     Callable,
+    Generic,
+    NamedTuple,
     Sequence,
     TypeVar,
-    Generic,
+    cast,
     get_args,
     get_origin,
-    NamedTuple,
-    cast,
 )
-from threading import Lock
-from enum import Enum
-from dataclasses import dataclass
+
 from rich.text import Text
 from rich.tree import Tree
 
@@ -32,8 +32,9 @@
 from . import op
 from . import setting
 from .convert import dump_engine_object, encode_engine_value, make_engine_value_decoder
-from .typing import encode_enriched_type
+from .op import FunctionSpec
 from .runtime import execution_context
+from .typing import encode_enriched_type
 
 
 class _NameBuilder:
@@ -91,6 +92,30 @@ def _spec_kind(spec: Any) -> str:
     return cast(str, spec.__class__.__name__)
 
 
+def _transform_helper(
+    flow_builder_state: _FlowBuilderState,
+    fn_spec: FunctionSpec,
+    transform_args: list[tuple[Any, str | None]],
+    name: str | None = None,
+) -> DataSlice[Any]:
+    if not isinstance(fn_spec, FunctionSpec):
+        raise ValueError("transform() can only be called on a CocoIndex function")
+
+    return _create_data_slice(
+        flow_builder_state,
+        lambda target_scope, name: flow_builder_state.engine_flow_builder.transform(
+            _spec_kind(fn_spec),
+            dump_engine_object(fn_spec),
+            transform_args,
+            target_scope,
+            flow_builder_state.field_name_builder.build_name(
+                name, prefix=_to_snake_case(_spec_kind(fn_spec)) + "_"
+            ),
+        ),
+        name,
+    )
+
+
 T = TypeVar("T")
 S = TypeVar("S")
 
@@ -190,31 +215,19 @@ def transform(
         """
         Apply a function to the data slice.
         """
-        if not isinstance(fn_spec, op.FunctionSpec):
-            raise ValueError("transform() can only be called on a CocoIndex function")
-
-        transform_args: list[tuple[Any, str | None]]
-        transform_args = [(self._state.engine_data_slice, None)]
+        transform_args: list[tuple[Any, str | None]] = [
+            (self._state.engine_data_slice, None)
+        ]
         transform_args += [
             (self._state.flow_builder_state.get_data_slice(v), None) for v in args
         ]
         transform_args += [
             (self._state.flow_builder_state.get_data_slice(v), k)
-            for (k, v) in kwargs.items()
+            for k, v in kwargs.items()
         ]
 
-        flow_builder_state = self._state.flow_builder_state
-        return _create_data_slice(
-            flow_builder_state,
-            lambda target_scope, name: flow_builder_state.engine_flow_builder.transform(
-                _spec_kind(fn_spec),
-                dump_engine_object(fn_spec),
-                transform_args,
-                target_scope,
-                flow_builder_state.field_name_builder.build_name(
-                    name, prefix=_to_snake_case(_spec_kind(fn_spec)) + "_"
-                ),
-            ),
+        return _transform_helper(
+            self._state.flow_builder_state, fn_spec, transform_args
         )
 
     def call(self, func: Callable[..., S], *args: Any, **kwargs: Any) -> S:
@@ -445,6 +458,24 @@ def add_source(
             name,
         )
 
+    def transform(
+        self, fn_spec: FunctionSpec, *args: Any, **kwargs: Any
+    ) -> DataSlice[Any]:
+        """
+        Apply a function to inputs, returning a DataSlice.
+        """
+        transform_args: list[tuple[Any, str | None]] = [
+            (self._state.get_data_slice(v), None) for v in args
+        ]
+        transform_args += [
+            (self._state.get_data_slice(v), k) for k, v in kwargs.items()
+        ]
+
+        if not transform_args:
+            raise ValueError("At least one input is required for transformation")
+
+        return _transform_helper(self._state, fn_spec, transform_args)
+
     def declare(self, spec: op.DeclarationSpec) -> None:
         """
         Add a declaration to the flow.
diff --git a/python/cocoindex/functions.py b/python/cocoindex/functions.py
@@ -1,30 +1,26 @@
 """All builtin functions."""
 
 import dataclasses
-from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeVar
+from typing import Annotated, Any, Literal
 
 import numpy as np
 from numpy.typing import NDArray
 
 from . import llm, op
-from .flow import DataSlice
 from .typing import TypeAttr, Vector
 
-# Libraries that are heavy to import. Lazily import them later.
-if TYPE_CHECKING:
-    import sentence_transformers
+# Check if sentence_transformers is available
+try:
+    import sentence_transformers  # type: ignore
 
-T = TypeVar("T")
+    _SENTENCE_TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    _SENTENCE_TRANSFORMERS_AVAILABLE = False
 
 
 class ParseJson(op.FunctionSpec):
     """Parse a text into a JSON object."""
 
-    def __call__(
-        self, *, text: DataSlice[T], language: str | None = "json"
-    ) -> DataSlice[T]:
-        return super().__call__(text=text, language=language)
-
 
 @dataclasses.dataclass
 class CustomLanguageSpec:
@@ -40,23 +36,6 @@ class SplitRecursively(op.FunctionSpec):
 
     custom_languages: list[CustomLanguageSpec] = dataclasses.field(default_factory=list)
 
-    def __call__(
-        self,
-        *,
-        text: DataSlice[T],
-        chunk_size: int,
-        min_chunk_size: int | None = None,
-        chunk_overlap: int | None = None,
-        language: DataSlice[T] | None = None,
-    ) -> DataSlice[T]:
-        return super().__call__(
-            text=text,
-            chunk_size=chunk_size,
-            language=language,
-            min_chunk_size=min_chunk_size,
-            chunk_overlap=chunk_overlap,
-        )
-
 
 class EmbedText(op.FunctionSpec):
     """Embed a text into a vector space."""
@@ -75,11 +54,6 @@ class ExtractByLlm(op.FunctionSpec):
     output_type: type
     instruction: str | None = None
 
-    def __call__(
-        self, *, text: DataSlice[T] | None = None, image: DataSlice[T] | None = None
-    ) -> DataSlice[T]:
-        return super().__call__(text=text, image=image)
-
 
 class SentenceTransformerEmbed(op.FunctionSpec):
     """
@@ -89,6 +63,10 @@ class SentenceTransformerEmbed(op.FunctionSpec):
 
         model: The name of the SentenceTransformer model to use.
         args: Additional arguments to pass to the SentenceTransformer constructor. e.g. {"trust_remote_code": True}
+
+    Note:
+        This function requires the optional sentence-transformers dependency.
+        Install it with: pip install 'cocoindex[embeddings]'
     """
 
     model: str
@@ -103,6 +81,14 @@ class SentenceTransformerEmbedExecutor:
     _model: "sentence_transformers.SentenceTransformer"
 
     def analyze(self, text: Any) -> type:
+        if not _SENTENCE_TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "sentence_transformers is required for SentenceTransformerEmbed function. "
+                "Install it with one of these commands:\n"
+                "  pip install 'cocoindex[embeddings]'\n"
+                "  pip install sentence-transformers"
+            )
+
         import sentence_transformers  # pylint: disable=import-outside-toplevel
 
         args = self.spec.args or {}
diff --git a/python/cocoindex/op.py b/python/cocoindex/op.py
@@ -9,8 +9,7 @@
 from typing import Any, Awaitable, Callable, Protocol, dataclass_transform
 
 from . import _engine  # type: ignore
-from .convert import dump_engine_object, encode_engine_value, make_engine_value_decoder
-from .flow import DataSlice, _create_data_slice, _spec_kind, _to_snake_case
+from .convert import encode_engine_value, make_engine_value_decoder
 from .typing import encode_enriched_type, resolve_forward_ref
 
 
@@ -49,51 +48,7 @@ class SourceSpec(metaclass=SpecMeta, category=OpCategory.SOURCE):  # pylint: dis
 
 
 class FunctionSpec(metaclass=SpecMeta, category=OpCategory.FUNCTION):  # pylint: disable=too-few-public-methods
-    """A function spec. Can be instantiated and called like a function: spec(args...).
-    For non-chain-style calls, use spec(args...) with at least one DataSlice argument.
-    For chain-style calls, use data_slice.transform(spec, args...).
-    """
-
-    def __call__(self, *args: Any, **kwargs: Any) -> DataSlice[Any]:
-        """Execute the function, returning a DataSlice."""
-
-        data_slice_args = [arg for arg in args if isinstance(arg, DataSlice)]
-        data_slice_kwargs = {
-            k: v for k, v in kwargs.items() if isinstance(v, DataSlice)
-        }
-        if not data_slice_args and not data_slice_kwargs:
-            raise ValueError(
-                "At least one DataSlice argument is required to provide flow context"
-            )
-
-        first_data_slice = (
-            data_slice_args[0]
-            if data_slice_args
-            else list(data_slice_kwargs.values())[0]
-        )
-        flow_builder_state = first_data_slice._state.flow_builder_state
-
-        transform_args: list[tuple[Any, str | None]] = [
-            (flow_builder_state.get_data_slice(v), None) for v in args if v is not None
-        ]
-        transform_args += [
-            (flow_builder_state.get_data_slice(v), k)
-            for k, v in kwargs.items()
-            if v is not None
-        ]
-
-        return _create_data_slice(
-            flow_builder_state,
-            lambda target_scope, name: flow_builder_state.engine_flow_builder.transform(
-                _spec_kind(self),
-                dump_engine_object(self),
-                transform_args,
-                target_scope,
-                flow_builder_state.field_name_builder.build_name(
-                    name, prefix=_to_snake_case(_spec_kind(self)) + "_"
-                ),
-            ),
-        )
+    """A function spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
 
 class TargetSpec(metaclass=SpecMeta, category=OpCategory.TARGET):  # pylint: disable=too-few-public-methods
@@ -355,7 +310,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
                 return fn(*args, **kwargs)
 
         class _Spec(FunctionSpec):
-            pass
+            def __call__(self, *args: Any, **kwargs: Any) -> Any:
+                return fn(*args, **kwargs)
 
         _Spec.__name__ = op_name
         _Spec.__doc__ = fn.__doc__
