feat(convert): support explicit type hints in engine value encoding

Author: lemorage

python/cocoindex/convert.py

@@ -28,69 +28,147 @@
     is_struct_type,
 )
 
+_CONVERTIBLE_KINDS = {
+    ("Float32", "Float64"),
+    ("LocalDateTime", "OffsetDateTime"),
+}
 
-def encode_engine_value(
-    value: Any, in_struct: bool = False, type_hint: Type[Any] | str | None = None
+
+def _is_type_kind_convertible_to(src_type_kind: str, dst_type_kind: str) -> bool:
+    return (
+        src_type_kind == dst_type_kind
+        or (src_type_kind, dst_type_kind) in _CONVERTIBLE_KINDS
+    )
+
+
+def _encode_engine_value_core(
+    value: Any,
+    in_struct: bool = False,
+    type_hint: Type[Any] | str | None = None,
+    type_variant: AnalyzedTypeInfo | None = None,
 ) -> Any:
-    """Encode a Python value to an engine value."""
+    """Core encoding logic for converting Python values to engine values."""
+
     if dataclasses.is_dataclass(value):
+        fields = dataclasses.fields(value)
         return [
             encode_engine_value(
-                getattr(value, f.name), in_struct=True, type_hint=f.type
+                getattr(value, f.name),
+                in_struct=True,
+                type_hint=f.type,
             )
-            for f in dataclasses.fields(value)
+            for f in fields
         ]
+
     if is_namedtuple_type(type(value)):
         annotations = type(value).__annotations__
         return [
             encode_engine_value(
-                getattr(value, name), in_struct=True, type_hint=annotations.get(name)
+                getattr(value, name),
+                in_struct=True,
+                type_hint=annotations.get(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, in_struct) for v in value]
+        if (
+            type_variant
+            and isinstance(type_variant.variant, AnalyzedListType)
+            and type_variant.variant.elem_type
+        ):
+            elem_encoder = make_engine_value_encoder(type_variant.variant.elem_type)
+            return [elem_encoder(v) for v in value]
+        else:
+            return [encode_engine_value(v, in_struct) for v in value]
+
     if isinstance(value, dict):
+        # Determine if this is a JSON type
         is_json_type = False
-        if type_hint:
-            type_info = analyze_type_info(type_hint)
+        if type_variant and isinstance(type_variant.variant, AnalyzedBasicType):
+            is_json_type = type_variant.variant.kind == "Json"
+        elif type_hint:
+            hint_type_info = analyze_type_info(type_hint)
             is_json_type = (
-                isinstance(type_info.variant, AnalyzedBasicType)
-                and type_info.variant.kind == "Json"
+                isinstance(hint_type_info.variant, AnalyzedBasicType)
+                and hint_type_info.variant.kind == "Json"
             )
 
-        # For empty dicts, check type hints if in a struct context
-        # when no contexts are provided, return an empty dict as default
-        # TODO: always pass in the type annotation to make this robust
+        # Handle empty dict
         if not value:
             if in_struct:
                 return value if is_json_type else []
-            return {}
+            return {} if is_json_type else value
 
+        # Handle KTable
         first_val = next(iter(value.values()))
-        if is_struct_type(type(first_val)):  # KTable
+        if is_struct_type(type(first_val)):
             return [
                 [encode_engine_value(k, in_struct)] + encode_engine_value(v, in_struct)
                 for k, v in value.items()
             ]
+
+        # Handle regular dict
+        if (
+            type_variant
+            and isinstance(type_variant.variant, AnalyzedDictType)
+            and type_variant.variant.value_type
+        ):
+            value_encoder = make_engine_value_encoder(type_variant.variant.value_type)
+            return {k: value_encoder(v) for k, v in value.items()}
+
     return value
 
 
-_CONVERTIBLE_KINDS = {
-    ("Float32", "Float64"),
-    ("LocalDateTime", "OffsetDateTime"),
-}
+def make_engine_value_encoder(type_annotation: Any) -> Callable[[Any], Any]:
+    """
+    Make an encoder from a Python value to an engine value.
 
+    Args:
+        type_annotation: The type annotation of the Python value.
 
-def _is_type_kind_convertible_to(src_type_kind: str, dst_type_kind: str) -> bool:
-    return (
-        src_type_kind == dst_type_kind
-        or (src_type_kind, dst_type_kind) in _CONVERTIBLE_KINDS
-    )
+    Returns:
+        An encoder from a Python value to an engine value.
+    """
+    type_info = analyze_type_info(type_annotation)
+
+    if isinstance(type_info.variant, AnalyzedUnknownType):
+        raise ValueError(f"Type annotation `{type_info.core_type}` is unsupported")
+
+    def encode_value(value: Any, in_struct: bool = False) -> Any:
+        return _encode_engine_value_core(
+            value, in_struct=in_struct, type_variant=type_info
+        )
+
+    return lambda value: encode_value(value, in_struct=False)
+
+
+def encode_engine_value(
+    value: Any, in_struct: bool = False, type_hint: Type[Any] | str | None = None
+) -> Any:
+    """
+    Encode a Python value to an engine value.
+
+    Args:
+        value: The Python value to encode
+        in_struct: Whether this value is being encoded within a struct context
+        type_hint: Type annotation for the value. When provided, enables optimized
+                  type-aware encoding. For top-level calls, this should always be provided.
+
+    Returns:
+        The encoded engine value
+    """
+    if type_hint is not None:
+        encoder = make_engine_value_encoder(type_hint)
+        return encoder(value)
+
+    return _encode_engine_value_core(value, in_struct=in_struct)
 
 
 def make_engine_value_decoder(

python/cocoindex/flow.py

@@ -9,28 +9,20 @@
 import functools
 import inspect
 import re
-
-from .validation import (
-    validate_flow_name,
-    NamingError,
-    validate_full_flow_name,
-    validate_target_name,
-)
-
 from dataclasses import dataclass
 from enum import Enum
 from threading import Lock
 from typing import (
     Any,
     Callable,
     Generic,
+    Iterable,
     NamedTuple,
     Sequence,
     TypeVar,
     cast,
     get_args,
     get_origin,
-    Iterable,
 )
 
 from rich.text import Text
@@ -45,6 +37,11 @@
 from .runtime import execution_context
 from .setup import SetupChangeBundle
 from .typing import encode_enriched_type
+from .validation import (
+    validate_flow_name,
+    validate_full_flow_name,
+    validate_target_name,
+)
 
 
 class _NameBuilder:
@@ -1083,10 +1080,13 @@ async def eval_async(self, *args: Any, **kwargs: Any) -> T:
         flow_info = await self._flow_info_async()
         params = []
         for i, arg in enumerate(self._param_names):
+            param_type = (
+                self._flow_arg_types[i] if i < len(self._flow_arg_types) else None
+            )
             if i < len(args):
-                params.append(encode_engine_value(args[i]))
+                params.append(encode_engine_value(args[i], type_hint=param_type))
             elif arg in kwargs:
-                params.append(encode_engine_value(kwargs[arg]))
+                params.append(encode_engine_value(kwargs[arg], type_hint=param_type))
             else:
                 raise ValueError(f"Parameter {arg} is not provided")
         engine_result = await flow_info.engine_flow.evaluate_async(params)

python/cocoindex/op.py

@@ -6,7 +6,7 @@
 import dataclasses
 import inspect
 from enum import Enum
-from typing import Any, Awaitable, Callable, Protocol, dataclass_transform, Annotated
+from typing import Annotated, Any, Awaitable, Callable, Protocol, dataclass_transform
 
 from . import _engine  # type: ignore
 from .convert import encode_engine_value, make_engine_value_decoder
@@ -277,7 +277,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> Any:
                     output = await self._acall(*decoded_args, **decoded_kwargs)
             else:
                 output = await self._acall(*decoded_args, **decoded_kwargs)
-            return encode_engine_value(output)
+            return encode_engine_value(output, type_hint=expected_return)
 
     _WrappedClass.__name__ = executor_cls.__name__
     _WrappedClass.__doc__ = executor_cls.__doc__