cocoindex 0.1.81__cp312-cp312-manylinux_2_28_x86_64.whl → 0.1.83__cp312-cp312-manylinux_2_28_x86_64.whl

cocoindex/cli.py

@@ -4,9 +4,7 @@ import datetime
 import importlib.util
 import os
 import signal
-import sys
 import threading
-import types
 from types import FrameType
 from typing import Any, Iterable
 
@@ -20,6 +18,8 @@ from rich.table import Table
 from . import flow, lib, setting
 from .setup import flow_names_with_setup
 from .runtime import execution_context
+from .subprocess_exec import add_user_app
+from .user_app_loader import load_user_app
 
 # Create ServerSettings lazily upon first call, as environment variables may be loaded from files, etc.
 COCOINDEX_HOST = "https://cocoindex.io"
@@ -76,50 +76,9 @@ def _get_app_ref_from_specifier(
     return app_ref
 
 
-def _load_user_app(app_target: str) -> types.ModuleType:
-    """
-    Loads the user's application, which can be a file path or an installed module name.
-    Exits on failure.
-    """
-    if not app_target:
-        raise click.ClickException("Application target not provided.")
-
-    looks_like_path = os.sep in app_target or app_target.lower().endswith(".py")
-
-    if looks_like_path:
-        if not os.path.isfile(app_target):
-            raise click.ClickException(f"Application file path not found: {app_target}")
-        app_path = os.path.abspath(app_target)
-        app_dir = os.path.dirname(app_path)
-        module_name = os.path.splitext(os.path.basename(app_path))[0]
-
-        if app_dir not in sys.path:
-            sys.path.insert(0, app_dir)
-        try:
-            spec = importlib.util.spec_from_file_location(module_name, app_path)
-            if spec is None:
-                raise ImportError(f"Could not create spec for file: {app_path}")
-            module = importlib.util.module_from_spec(spec)
-            sys.modules[spec.name] = module
-            if spec.loader is None:
-                raise ImportError(f"Could not create loader for file: {app_path}")
-            spec.loader.exec_module(module)
-            return module
-        except (ImportError, FileNotFoundError, PermissionError) as e:
-            raise click.ClickException(f"Failed importing file '{app_path}': {e}")
-        finally:
-            if app_dir in sys.path and sys.path[0] == app_dir:
-                sys.path.pop(0)
-
-    # Try as module
-    try:
-        return importlib.import_module(app_target)
-    except ImportError as e:
-        raise click.ClickException(f"Failed to load module '{app_target}': {e}")
-    except Exception as e:
-        raise click.ClickException(
-            f"Unexpected error importing module '{app_target}': {e}"
-        )
+def _load_user_app(app_target: str) -> None:
+    load_user_app(app_target)
+    add_user_app(app_target)
 
 
 def _initialize_cocoindex_in_process() -> None:

cocoindex/convert.py

@@ -9,26 +9,26 @@ import datetime
 import inspect
 import warnings
 from enum import Enum
-from typing import Any, Callable, Mapping, get_origin
+from typing import Any, Callable, Mapping, Type, get_origin
 
 import numpy as np
 
 from .typing import (
     KEY_FIELD_NAME,
     TABLE_TYPES,
-    analyze_type_info,
-    encode_enriched_type,
-    is_namedtuple_type,
-    is_struct_type,
-    AnalyzedTypeInfo,
     AnalyzedAnyType,
+    AnalyzedBasicType,
     AnalyzedDictType,
     AnalyzedListType,
-    AnalyzedBasicType,
+    AnalyzedStructType,
+    AnalyzedTypeInfo,
     AnalyzedUnionType,
     AnalyzedUnknownType,
-    AnalyzedStructType,
+    analyze_type_info,
+    encode_enriched_type,
+    is_namedtuple_type,
     is_numpy_number_type,
+    is_struct_type,
 )
 
 
@@ -50,34 +50,6 @@ class ChildFieldPath:
         self._field_path.pop()
 
 
-def encode_engine_value(value: Any) -> Any:
-    """Encode a Python value to an engine value."""
-    if dataclasses.is_dataclass(value):
-        return [
-            encode_engine_value(getattr(value, f.name))
-            for f in dataclasses.fields(value)
-        ]
-    if is_namedtuple_type(type(value)):
-        return [encode_engine_value(getattr(value, name)) for name in value._fields]
-    if isinstance(value, np.number):
-        return value.item()
-    if isinstance(value, np.ndarray):
-        return value
-    if isinstance(value, (list, tuple)):
-        return [encode_engine_value(v) for v in value]
-    if isinstance(value, dict):
-        if not value:
-            return {}
-
-        first_val = next(iter(value.values()))
-        if is_struct_type(type(first_val)):  # KTable
-            return [
-                [encode_engine_value(k)] + encode_engine_value(v)
-                for k, v in value.items()
-            ]
-    return value
-
-
 _CONVERTIBLE_KINDS = {
     ("Float32", "Float64"),
     ("LocalDateTime", "OffsetDateTime"),
@@ -91,6 +63,118 @@ def _is_type_kind_convertible_to(src_type_kind: str, dst_type_kind: str) -> bool
     )
 
 
+# Pre-computed type info for missing/Any type annotations
+ANY_TYPE_INFO = analyze_type_info(inspect.Parameter.empty)
+
+
+def make_engine_value_encoder(type_info: AnalyzedTypeInfo) -> Callable[[Any], Any]:
+    """
+    Create an encoder closure for a specific type.
+    """
+    variant = type_info.variant
+
+    if isinstance(variant, AnalyzedUnknownType):
+        raise ValueError(f"Type annotation `{type_info.core_type}` is unsupported")
+
+    if isinstance(variant, AnalyzedListType):
+        elem_type_info = (
+            analyze_type_info(variant.elem_type) if variant.elem_type else ANY_TYPE_INFO
+        )
+        if isinstance(elem_type_info.variant, AnalyzedStructType):
+            elem_encoder = make_engine_value_encoder(elem_type_info)
+
+            def encode_struct_list(value: Any) -> Any:
+                return None if value is None else [elem_encoder(v) for v in value]
+
+            return encode_struct_list
+
+    if isinstance(variant, AnalyzedDictType):
+        if not variant.value_type:
+            return lambda value: value
+
+        value_type_info = analyze_type_info(variant.value_type)
+        if isinstance(value_type_info.variant, AnalyzedStructType):
+
+            def encode_struct_dict(value: Any) -> Any:
+                if not isinstance(value, dict):
+                    return value
+                if not value:
+                    return []
+
+                sample_key, sample_val = next(iter(value.items()))
+                key_type, val_type = type(sample_key), type(sample_val)
+
+                # Handle KTable case
+                if value and is_struct_type(val_type):
+                    key_encoder = (
+                        make_engine_value_encoder(analyze_type_info(key_type))
+                        if is_struct_type(key_type)
+                        else make_engine_value_encoder(ANY_TYPE_INFO)
+                    )
+                    value_encoder = make_engine_value_encoder(
+                        analyze_type_info(val_type)
+                    )
+                    return [
+                        [key_encoder(k)] + value_encoder(v) for k, v in value.items()
+                    ]
+                return {key_encoder(k): value_encoder(v) for k, v in value.items()}
+
+            return encode_struct_dict
+
+    if isinstance(variant, AnalyzedStructType):
+        struct_type = variant.struct_type
+
+        if dataclasses.is_dataclass(struct_type):
+            fields = dataclasses.fields(struct_type)
+            field_encoders = [
+                make_engine_value_encoder(analyze_type_info(f.type)) for f in fields
+            ]
+            field_names = [f.name for f in fields]
+
+            def encode_dataclass(value: Any) -> Any:
+                if not dataclasses.is_dataclass(value):
+                    return value
+                return [
+                    encoder(getattr(value, name))
+                    for encoder, name in zip(field_encoders, field_names)
+                ]
+
+            return encode_dataclass
+
+        elif is_namedtuple_type(struct_type):
+            annotations = struct_type.__annotations__
+            field_names = list(getattr(struct_type, "_fields", ()))
+            field_encoders = [
+                make_engine_value_encoder(
+                    analyze_type_info(annotations[name])
+                    if name in annotations
+                    else ANY_TYPE_INFO
+                )
+                for name in field_names
+            ]
+
+            def encode_namedtuple(value: Any) -> Any:
+                if not is_namedtuple_type(type(value)):
+                    return value
+                return [
+                    encoder(getattr(value, name))
+                    for encoder, name in zip(field_encoders, field_names)
+                ]
+
+            return encode_namedtuple
+
+    def encode_basic_value(value: Any) -> Any:
+        if isinstance(value, np.number):
+            return value.item()
+        if isinstance(value, np.ndarray):
+            return value
+        if isinstance(value, (list, tuple)):
+            return [encode_basic_value(v) for v in value]
+        return value
+
+    return encode_basic_value
+
+
 def make_engine_value_decoder(
     field_path: list[str],
     src_type: dict[str, Any],

cocoindex/flow.py

@@ -9,15 +9,6 @@ import datetime
 import functools
 import inspect
 import re
-
-from .validation import (
-    validate_flow_name,
-    NamingError,
-    validate_full_flow_name,
-    validate_target_name,
-)
-from .typing import analyze_type_info
-
 from dataclasses import dataclass
 from enum import Enum
 from threading import Lock
@@ -25,13 +16,13 @@ from typing import (
     Any,
     Callable,
     Generic,
+    Iterable,
     NamedTuple,
     Sequence,
     TypeVar,
     cast,
     get_args,
     get_origin,
-    Iterable,
 )
 
 from rich.text import Text
@@ -41,11 +32,20 @@ from . import _engine  # type: ignore
 from . import index
 from . import op
 from . import setting
-from .convert import dump_engine_object, encode_engine_value, make_engine_value_decoder
+from .convert import (
+    dump_engine_object,
+    make_engine_value_decoder,
+    make_engine_value_encoder,
+)
 from .op import FunctionSpec
 from .runtime import execution_context
 from .setup import SetupChangeBundle
-from .typing import encode_enriched_type
+from .typing import analyze_type_info, encode_enriched_type
+from .validation import (
+    validate_flow_name,
+    validate_full_flow_name,
+    validate_target_name,
+)
 
 
 class _NameBuilder:
@@ -105,18 +105,26 @@ def _spec_kind(spec: Any) -> str:
 
 def _transform_helper(
     flow_builder_state: _FlowBuilderState,
-    fn_spec: FunctionSpec,
+    fn_spec: FunctionSpec | Callable[..., Any],
     transform_args: list[tuple[Any, str | None]],
     name: str | None = None,
 ) -> DataSlice[Any]:
-    if not isinstance(fn_spec, FunctionSpec):
+    if isinstance(fn_spec, FunctionSpec):
+        kind = _spec_kind(fn_spec)
+        spec = fn_spec
+    elif callable(fn_spec) and (
+        op_kind := getattr(fn_spec, "__cocoindex_op_kind__", None)
+    ):
+        kind = op_kind
+        spec = op.EmptyFunctionSpec()
+    else:
         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),
+            kind,
+            dump_engine_object(spec),
             transform_args,
             target_scope,
             flow_builder_state.field_name_builder.build_name(
@@ -245,7 +253,7 @@ class DataSlice(Generic[T]):
             f(scope)
 
     def transform(
-        self, fn_spec: op.FunctionSpec, *args: Any, **kwargs: Any
+        self, fn_spec: op.FunctionSpec | Callable[..., Any], *args: Any, **kwargs: Any
     ) -> DataSlice[Any]:
         """
         Apply a function to the data slice.
@@ -513,7 +521,7 @@ class FlowBuilder:
         )
 
     def transform(
-        self, fn_spec: FunctionSpec, *args: Any, **kwargs: Any
+        self, fn_spec: FunctionSpec | Callable[..., Any], *args: Any, **kwargs: Any
     ) -> DataSlice[Any]:
         """
         Apply a function to inputs, returning a DataSlice.
@@ -978,6 +986,12 @@ class TransformFlowInfo(NamedTuple):
     result_decoder: Callable[[Any], T]
 
 
+class FlowArgInfo(NamedTuple):
+    name: str
+    type_hint: Any
+    encoder: Callable[[Any], Any]
+
+
 class TransformFlow(Generic[T]):
     """
     A transient transformation flow that transforms in-memory data.
@@ -985,8 +999,7 @@ class TransformFlow(Generic[T]):
 
     _flow_fn: Callable[..., DataSlice[T]]
     _flow_name: str
-    _flow_arg_types: list[Any]
-    _param_names: list[str]
+    _args_info: list[FlowArgInfo]
 
     _lazy_lock: asyncio.Lock
     _lazy_flow_info: TransformFlowInfo | None = None
@@ -994,7 +1007,6 @@ class TransformFlow(Generic[T]):
     def __init__(
         self,
         flow_fn: Callable[..., DataSlice[T]],
-        flow_arg_types: Sequence[Any],
         /,
         name: str | None = None,
     ):
@@ -1002,9 +1014,32 @@ class TransformFlow(Generic[T]):
         self._flow_name = _transform_flow_name_builder.build_name(
             name, prefix="_transform_flow_"
         )
-        self._flow_arg_types = list(flow_arg_types)
         self._lazy_lock = asyncio.Lock()
 
+        sig = inspect.signature(flow_fn)
+        args_info = []
+        for param_name, param in sig.parameters.items():
+            if param.kind not in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            ):
+                raise ValueError(
+                    f"Parameter `{param_name}` is not a parameter can be passed by name"
+                )
+            value_type_annotation: type | None = _get_data_slice_annotation_type(
+                param.annotation
+            )
+            if value_type_annotation is None:
+                raise ValueError(
+                    f"Parameter `{param_name}` for {flow_fn} has no value type annotation. "
+                    "Please use `cocoindex.DataSlice[T]` where T is the type of the value."
+                )
+            encoder = make_engine_value_encoder(
+                analyze_type_info(value_type_annotation)
+            )
+            args_info.append(FlowArgInfo(param_name, value_type_annotation, encoder))
+        self._args_info = args_info
+
     def __call__(self, *args: Any, **kwargs: Any) -> DataSlice[T]:
         return self._flow_fn(*args, **kwargs)
 
@@ -1024,31 +1059,15 @@ class TransformFlow(Generic[T]):
 
     async def _build_flow_info_async(self) -> TransformFlowInfo:
         flow_builder_state = _FlowBuilderState(self._flow_name)
-        sig = inspect.signature(self._flow_fn)
-        if len(sig.parameters) != len(self._flow_arg_types):
-            raise ValueError(
-                f"Number of parameters in the flow function ({len(sig.parameters)}) "
-                f"does not match the number of argument types ({len(self._flow_arg_types)})"
-            )
-
         kwargs: dict[str, DataSlice[T]] = {}
-        for (param_name, param), param_type in zip(
-            sig.parameters.items(), self._flow_arg_types
-        ):
-            if param.kind not in (
-                inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                inspect.Parameter.KEYWORD_ONLY,
-            ):
-                raise ValueError(
-                    f"Parameter `{param_name}` is not a parameter can be passed by name"
-                )
-            encoded_type = encode_enriched_type(param_type)
+        for arg_info in self._args_info:
+            encoded_type = encode_enriched_type(arg_info.type_hint)
             if encoded_type is None:
-                raise ValueError(f"Parameter `{param_name}` has no type annotation")
+                raise ValueError(f"Parameter `{arg_info.name}` has no type annotation")
             engine_ds = flow_builder_state.engine_flow_builder.add_direct_input(
-                param_name, encoded_type
+                arg_info.name, encoded_type
             )
-            kwargs[param_name] = DataSlice(
+            kwargs[arg_info.name] = DataSlice(
                 _DataSliceState(flow_builder_state, engine_ds)
             )
 
@@ -1061,13 +1080,12 @@ class TransformFlow(Generic[T]):
                 execution_context.event_loop
             )
         )
-        self._param_names = list(sig.parameters.keys())
 
         engine_return_type = (
             _data_slice_state(output).engine_data_slice.data_type().schema()
         )
         python_return_type: type[T] | None = _get_data_slice_annotation_type(
-            sig.return_annotation
+            inspect.signature(self._flow_fn).return_annotation
         )
         result_decoder = make_engine_value_decoder(
             [], engine_return_type["type"], analyze_type_info(python_return_type)
@@ -1099,13 +1117,14 @@ class TransformFlow(Generic[T]):
         """
         flow_info = await self._flow_info_async()
         params = []
-        for i, arg in enumerate(self._param_names):
+        for i, arg_info in enumerate(self._args_info):
             if i < len(args):
-                params.append(encode_engine_value(args[i]))
+                arg = args[i]
             elif arg in kwargs:
-                params.append(encode_engine_value(kwargs[arg]))
+                arg = kwargs[arg]
             else:
                 raise ValueError(f"Parameter {arg} is not provided")
+            params.append(arg_info.encoder(arg))
         engine_result = await flow_info.engine_flow.evaluate_async(params)
         return flow_info.result_decoder(engine_result)
 
@@ -1116,27 +1135,7 @@ def transform_flow() -> Callable[[Callable[..., DataSlice[T]]], TransformFlow[T]
     """
 
     def _transform_flow_wrapper(fn: Callable[..., DataSlice[T]]) -> TransformFlow[T]:
-        sig = inspect.signature(fn)
-        arg_types = []
-        for param_name, param in sig.parameters.items():
-            if param.kind not in (
-                inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                inspect.Parameter.KEYWORD_ONLY,
-            ):
-                raise ValueError(
-                    f"Parameter `{param_name}` is not a parameter can be passed by name"
-                )
-            value_type_annotation: type[T] | None = _get_data_slice_annotation_type(
-                param.annotation
-            )
-            if value_type_annotation is None:
-                raise ValueError(
-                    f"Parameter `{param_name}` for {fn} has no value type annotation. "
-                    "Please use `cocoindex.DataSlice[T]` where T is the type of the value."
-                )
-            arg_types.append(value_type_annotation)
-
-        _transform_flow = TransformFlow(fn, arg_types)
+        _transform_flow = TransformFlow(fn)
         functools.update_wrapper(_transform_flow, fn)
         return _transform_flow
 

cocoindex/functions.py

@@ -89,7 +89,7 @@ class SentenceTransformerEmbedExecutor:
     spec: SentenceTransformerEmbed
     _model: Any | None = None
 
-    def analyze(self, _text: Any) -> type:
+    def analyze(self) -> type:
         try:
             # Only import sentence_transformers locally when it's needed, as its import is very slow.
             import sentence_transformers  # pylint: disable=import-outside-toplevel
@@ -245,7 +245,7 @@ class ColPaliEmbedImageExecutor:
     spec: ColPaliEmbedImage
     _model_info: ColPaliModelInfo
 
-    def analyze(self, _img_bytes: Any) -> type:
+    def analyze(self) -> type:
         # Get shared model and dimension
         self._model_info = _get_colpali_model_and_processor(self.spec.model)
 
@@ -321,7 +321,7 @@ class ColPaliEmbedQueryExecutor:
     spec: ColPaliEmbedQuery
     _model_info: ColPaliModelInfo
 
-    def analyze(self, _query: Any) -> type:
+    def analyze(self) -> type:
         # Get shared model and dimension
         self._model_info = _get_colpali_model_and_processor(self.spec.model)