cocoindex 0.1.57__cp311-cp311-manylinux_2_28_x86_64.whl → 0.1.59__cp311-cp311-manylinux_2_28_x86_64.whl

cocoindex/convert.py

@@ -5,7 +5,6 @@ Utilities to convert between Python and engine values.
 import dataclasses
 import datetime
 import inspect
-import uuid
 from enum import Enum
 from typing import Any, Callable, Mapping, get_origin
 
@@ -14,7 +13,6 @@ import numpy as np
 from .typing import (
     KEY_FIELD_NAME,
     TABLE_TYPES,
-    AnalyzedTypeInfo,
     DtypeRegistry,
     analyze_type_info,
     encode_enriched_type,
@@ -74,23 +72,16 @@ def make_engine_value_decoder(
     Returns:
         A decoder from an engine value to a Python value.
     """
-
     src_type_kind = src_type["kind"]
 
-    dst_type_info: AnalyzedTypeInfo | None = None
-    if (
-        dst_annotation is not None
-        and dst_annotation is not inspect.Parameter.empty
-        and dst_annotation is not Any
-    ):
-        dst_type_info = analyze_type_info(dst_annotation)
-        if not _is_type_kind_convertible_to(src_type_kind, dst_type_info.kind):
-            raise ValueError(
-                f"Type mismatch for `{''.join(field_path)}`: "
-                f"passed in {src_type_kind}, declared {dst_annotation} ({dst_type_info.kind})"
-            )
-
-    if dst_type_info is None:
+    dst_is_any = (
+        dst_annotation is None
+        or dst_annotation is inspect.Parameter.empty
+        or dst_annotation is Any
+    )
+    if dst_is_any:
+        if src_type_kind == "Union":
+            return lambda value: value[1]
         if src_type_kind == "Struct" or src_type_kind in TABLE_TYPES:
             raise ValueError(
                 f"Missing type annotation for `{''.join(field_path)}`."
@@ -98,6 +89,41 @@ def make_engine_value_decoder(
             )
         return lambda value: value
 
+    dst_type_info = analyze_type_info(dst_annotation)
+
+    if src_type_kind == "Union":
+        dst_type_variants = (
+            dst_type_info.union_variant_types
+            if dst_type_info.union_variant_types is not None
+            else [dst_annotation]
+        )
+        src_type_variants = src_type["types"]
+        decoders = []
+        for i, src_type_variant in enumerate(src_type_variants):
+            src_field_path = field_path + [f"[{i}]"]
+            decoder = None
+            for dst_type_variant in dst_type_variants:
+                try:
+                    decoder = make_engine_value_decoder(
+                        src_field_path, src_type_variant, dst_type_variant
+                    )
+                    break
+                except ValueError:
+                    pass
+            if decoder is None:
+                raise ValueError(
+                    f"Type mismatch for `{''.join(field_path)}`: "
+                    f"cannot find matched target type for source type variant {src_type_variant}"
+                )
+            decoders.append(decoder)
+        return lambda value: decoders[value[0]](value[1])
+
+    if not _is_type_kind_convertible_to(src_type_kind, dst_type_info.kind):
+        raise ValueError(
+            f"Type mismatch for `{''.join(field_path)}`: "
+            f"passed in {src_type_kind}, declared {dst_annotation} ({dst_type_info.kind})"
+        )
+
     if dst_type_info.kind in ("Float32", "Float64", "Int64"):
         dst_core_type = dst_type_info.core_type
 
@@ -196,9 +222,6 @@ def make_engine_value_decoder(
         field_path.pop()
         return decode
 
-    if src_type_kind == "Union":
-        return lambda value: value[1]
-
     return lambda value: value
 
 

cocoindex/flow.py

@@ -5,25 +5,27 @@ Flow is the main interface for building and running flows.
 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,
     Iterable,
 )
-from threading import Lock
-from enum import Enum
-from dataclasses import dataclass
+
 from rich.text import Text
 from rich.tree import Tree
 
@@ -32,9 +34,10 @@ from . import index
 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 .setup import SetupChangeBundle
+from .typing import encode_enriched_type
 
 
 class _NameBuilder:
@@ -92,6 +95,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")
 
@@ -191,31 +218,19 @@ class DataSlice(Generic[T]):
         """
         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:
@@ -446,6 +461,24 @@ class FlowBuilder:
             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.

cocoindex/llm.py

@@ -12,6 +12,7 @@ class LlmApiType(Enum):
     LITE_LLM = "LiteLlm"
     OPEN_ROUTER = "OpenRouter"
     VOYAGE = "Voyage"
+    VLLM = "Vllm"
 
 
 @dataclass

cocoindex/op.py

@@ -5,13 +5,12 @@ Facilities for defining cocoindex operations.
 import asyncio
 import dataclasses
 import inspect
-
-from typing import Protocol, Any, Callable, Awaitable, dataclass_transform
 from enum import Enum
+from typing import Any, Awaitable, Callable, Protocol, dataclass_transform
 
-from .typing import encode_enriched_type, resolve_forward_ref
-from .convert import encode_engine_value, make_engine_value_decoder
 from . import _engine  # type: ignore
+from .convert import encode_engine_value, make_engine_value_decoder
+from .typing import encode_enriched_type, resolve_forward_ref
 
 
 class OpCategory(Enum):

cocoindex/tests/test_convert.py

@@ -104,7 +104,7 @@ def validate_full_roundtrip(
     )
     decoder = make_engine_value_decoder([], encoded_output_type, value_type)
     decoded_value = decoder(value_from_engine)
-    assert eq(decoded_value, value)
+    assert eq(decoded_value, value), f"{decoded_value} != {value}"
 
     if other_decoded_values is not None:
         for other_value, other_type in other_decoded_values:
@@ -613,6 +613,18 @@ def test_roundtrip_union_timedelta() -> None:
     validate_full_roundtrip(value, t)
 
 
+def test_roundtrip_vector_of_union() -> None:
+    t = list[str | int]
+    value = ["a", 1]
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_with_vector() -> None:
+    t = NDArray[np.float32] | str
+    value = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    validate_full_roundtrip(value, t, ([1.0, 2.0, 3.0], list[float] | str))
+
+
 def test_roundtrip_ltable() -> None:
     t = list[Order]
     value = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]

{cocoindex-0.1.57.dist-info → cocoindex-0.1.59.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.57
+Version: 0.1.59
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0

{cocoindex-0.1.57.dist-info → cocoindex-0.1.59.dist-info}/RECORD

@@ -1,18 +1,18 @@
-cocoindex-0.1.57.dist-info/METADATA,sha256=B506wZaMqr_x4PwQrv-tQZqvkukTH4nkUkreAFAZkIk,10020
-cocoindex-0.1.57.dist-info/WHEEL,sha256=1ugA4CjcVVPH4LDWfawu_6l5KZf3P6xl5-C8vT5fNPY,108
-cocoindex-0.1.57.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.57.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex-0.1.59.dist-info/METADATA,sha256=XyuEgs1E6LhwzhvSwCwjcRmzFZkEHaw-GFEgpp5qVn8,10020
+cocoindex-0.1.59.dist-info/WHEEL,sha256=1ugA4CjcVVPH4LDWfawu_6l5KZf3P6xl5-C8vT5fNPY,108
+cocoindex-0.1.59.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.59.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 cocoindex/__init__.py,sha256=MFm-QJzrr0ODJCsAAsPUzJXh8KH1WdZod8F60B1iUYw,1877
-cocoindex/_engine.cpython-311-x86_64-linux-gnu.so,sha256=aWjA6EYfzpbkXfwKNSVxnAYhflM2KmlB0n_sNO5OdKU,63198000
+cocoindex/_engine.cpython-311-x86_64-linux-gnu.so,sha256=Z584GErs4MYSUSRtODhQ7dFHZqDTBn-GFKipJyZwce0,63206328
 cocoindex/auth_registry.py,sha256=1XqO7ibjmBBd8i11XSJTvTgdz8p1ptW-ZpuSgo_5zzk,716
 cocoindex/cli.py,sha256=8bDL-Qmd9NYtn1DsDfvUMk45xfAqNf9YTyM7H9KRuNU,21345
-cocoindex/convert.py,sha256=XeSr0ykudBrB-RWRmcbdbt3WCihlLDSBIYUcDtPbTdA,10228
-cocoindex/flow.py,sha256=YUDXhCEolKIMIW_0ba2bwu27-kJIo1_t3ePy_sSkYHw,32191
+cocoindex/convert.py,sha256=FsKb2Pfbm7e1VQDOs_AsoiW9PbIUuyHQuqUlrENXmUY,11199
+cocoindex/flow.py,sha256=UQviW2O6nzCKt3jf7N1xAXULvK7YRCZDAj5hy5VLkjM,33041
 cocoindex/functions.py,sha256=IBwvdPpGR-S5mk53HvHpT2GVs15MI9wQznxgOdxA0ac,3202
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
 cocoindex/lib.py,sha256=BeRUn3RqE_wSsVtsgCzbFFKe1LXgRyRmMOcmwWBuEXo,2940
-cocoindex/llm.py,sha256=hpwvHjWiLkoV018TkC191rztq62NQeewwF3kDsHUols,430
-cocoindex/op.py,sha256=Z7V9Fdz4qeTeozzKmp1Dk1lPUP0GBgVgD_vBEhTJS5Y,11810
+cocoindex/llm.py,sha256=0ri8ZRg9_Zf2gyC5xuQ1Kq6kdZUO8r-A5WLnxit5S_4,448
+cocoindex/op.py,sha256=r_Usx7Jqh49Cck3tsYLx2vLRNUZArkQP_g7bIID6LPU,11809
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/runtime.py,sha256=bAdHYaXFWiiUWyAgzmKTeaAaRR0D_AmaqVCIdPO-v00,1056
 cocoindex/setting.py,sha256=Zl8K86r8RVvG9c3pCsH0Ot8BHhDAQAQCjoBp7TnXMLQ,3590
@@ -20,9 +20,9 @@ cocoindex/setup.py,sha256=7uIHKN4FOCuoidPXcKyGTrkqpkl9luL49-6UcnMxYzw,3068
 cocoindex/sources.py,sha256=JCnOhv1w4o28e03i7yvo4ESicWYAhckkBg5bQlxNH4U,1330
 cocoindex/targets.py,sha256=Nfh_tpFd1goTnS_cxBjIs4j9zl3Z4Z1JomAQ1dl3Sic,2796
 cocoindex/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/tests/test_convert.py,sha256=1IoJzfPtM5zoy0YKhTsZTX-fF4xA3w_4usyS4UN44vE,36150
+cocoindex/tests/test_convert.py,sha256=fY3p1gRHZpuWUgk5v28UGPYzyaMI4Qakv4I4qHSNGu8,36528
 cocoindex/tests/test_optional_database.py,sha256=snAmkNa6wtOSaxoZE1HgjvL5v_ylitt3Jt_9df4Cgdc,8506
 cocoindex/tests/test_typing.py,sha256=t6UCYShcfonTfjBlGRWPiFGMZ8DGFfABXo6idekPoJE,14757
 cocoindex/typing.py,sha256=kPMFVKs2i4SCLzW1Tn5NP_Ev9DAc-2qW6eJ68gpLexU,12580
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
-cocoindex-0.1.57.dist-info/RECORD,,
+cocoindex-0.1.59.dist-info/RECORD,,