cocoindex 0.1.53__cp313-cp313-macosx_11_0_arm64.whl → 0.1.54__cp313-cp313-macosx_11_0_arm64.whl

cocoindex/cli.py

@@ -1,19 +1,19 @@
-import click
+import atexit
 import datetime
-import sys
 import importlib.util
 import os
-import atexit
+import sys
 import types
+from typing import Any
 
-from dotenv import load_dotenv, find_dotenv
+import click
+from dotenv import find_dotenv, load_dotenv
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
-from typing import Any
 
 from . import flow, lib, setting
-from .setup import sync_setup, drop_setup, flow_names_with_setup, apply_setup_changes
+from .setup import apply_setup_changes, drop_setup, flow_names_with_setup, sync_setup
 
 # Create ServerSettings lazily upon first call, as environment variables may be loaded from files, etc.
 COCOINDEX_HOST = "https://cocoindex.io"

cocoindex/convert.py

@@ -14,6 +14,7 @@ import numpy as np
 from .typing import (
     KEY_FIELD_NAME,
     TABLE_TYPES,
+    AnalyzedTypeInfo,
     DtypeRegistry,
     analyze_type_info,
     encode_enriched_type,
@@ -46,6 +47,19 @@ def encode_engine_value(value: Any) -> Any:
     return value
 
 
+_CONVERTIBLE_KINDS = {
+    ("Float32", "Float64"),
+    ("LocalDateTime", "OffsetDateTime"),
+}
+
+
+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 make_engine_value_decoder(
     field_path: list[str],
     src_type: dict[str, Any],
@@ -65,11 +79,23 @@ def make_engine_value_decoder(
 
     src_type_kind = src_type["kind"]
 
+    dst_type_info: AnalyzedTypeInfo | None = None
     if (
-        dst_annotation is None
-        or dst_annotation is inspect.Parameter.empty
-        or dst_annotation is Any
+        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 src_type_kind == "Uuid":
+        return lambda value: uuid.UUID(bytes=value)
+
+    if dst_type_info is None:
         if src_type_kind == "Struct" or src_type_kind in TABLE_TYPES:
             raise ValueError(
                 f"Missing type annotation for `{''.join(field_path)}`."
@@ -77,32 +103,66 @@ def make_engine_value_decoder(
             )
         return lambda value: value
 
-    dst_type_info = analyze_type_info(dst_annotation)
+    if dst_type_info.kind in ("Float32", "Float64", "Int64"):
+        dst_core_type = dst_type_info.core_type
 
-    if 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})"
-        )
+        def decode_scalar(value: Any) -> Any | None:
+            if value is None:
+                if dst_type_info.nullable:
+                    return None
+                raise ValueError(
+                    f"Received null for non-nullable scalar `{''.join(field_path)}`"
+                )
+            return dst_core_type(value)
 
-    if dst_type_info.struct_type is not None:
-        return _make_engine_struct_value_decoder(
-            field_path, src_type["fields"], dst_type_info.struct_type
+        return decode_scalar
+
+    if src_type_kind == "Vector":
+        field_path_str = "".join(field_path)
+        expected_dim = (
+            dst_type_info.vector_info.dim if dst_type_info.vector_info else None
         )
 
-    if dst_type_info.np_number_type is not None and src_type_kind != "Vector":
-        numpy_type = dst_type_info.np_number_type
+        elem_decoder = None
+        scalar_dtype = None
+        if dst_type_info.np_number_type is None:  # for Non-NDArray vector
+            elem_decoder = make_engine_value_decoder(
+                field_path + ["[*]"],
+                src_type["element_type"],
+                dst_type_info.elem_type,
+            )
+        else:  # for NDArray vector
+            scalar_dtype = extract_ndarray_scalar_dtype(dst_type_info.np_number_type)
+            _ = DtypeRegistry.validate_dtype_and_get_kind(scalar_dtype)
 
-        def decode_numpy_scalar(value: Any) -> Any | None:
+        def decode_vector(value: Any) -> Any | None:
             if value is None:
                 if dst_type_info.nullable:
                     return None
                 raise ValueError(
-                    f"Received null for non-nullable scalar `{''.join(field_path)}`"
+                    f"Received null for non-nullable vector `{field_path_str}`"
                 )
-            return numpy_type(value)
+            if not isinstance(value, (np.ndarray, list)):
+                raise TypeError(
+                    f"Expected NDArray or list for vector `{field_path_str}`, got {type(value)}"
+                )
+            if expected_dim is not None and len(value) != expected_dim:
+                raise ValueError(
+                    f"Vector dimension mismatch for `{field_path_str}`: "
+                    f"expected {expected_dim}, got {len(value)}"
+                )
+
+            if elem_decoder is not None:  # for Non-NDArray vector
+                return [elem_decoder(v) for v in value]
+            else:  # for NDArray vector
+                return np.array(value, dtype=scalar_dtype)
 
-        return decode_numpy_scalar
+        return decode_vector
+
+    if dst_type_info.struct_type is not None:
+        return _make_engine_struct_value_decoder(
+            field_path, src_type["fields"], dst_type_info.struct_type
+        )
 
     if src_type_kind in TABLE_TYPES:
         field_path.append("[*]")
@@ -141,48 +201,8 @@ def make_engine_value_decoder(
         field_path.pop()
         return decode
 
-    if src_type_kind == "Uuid":
-        return lambda value: uuid.UUID(bytes=value)
-
-    if src_type_kind == "Vector":
-
-        def decode_vector(value: Any) -> Any | None:
-            field_path_str = "".join(field_path)
-            expected_dim = (
-                dst_type_info.vector_info.dim if dst_type_info.vector_info else None
-            )
-
-            if value is None:
-                if dst_type_info.nullable:
-                    return None
-                raise ValueError(
-                    f"Received null for non-nullable vector `{field_path_str}`"
-                )
-            if not isinstance(value, (np.ndarray, list)):
-                raise TypeError(
-                    f"Expected NDArray or list for vector `{field_path_str}`, got {type(value)}"
-                )
-            if expected_dim is not None and len(value) != expected_dim:
-                raise ValueError(
-                    f"Vector dimension mismatch for `{field_path_str}`: "
-                    f"expected {expected_dim}, got {len(value)}"
-                )
-
-            if dst_type_info.np_number_type is None:  # for Non-NDArray vector
-                elem_decoder = make_engine_value_decoder(
-                    field_path + ["[*]"],
-                    src_type["element_type"],
-                    dst_type_info.elem_type,
-                )
-                return [elem_decoder(v) for v in value]
-            else:  # for NDArray vector
-                scalar_dtype = extract_ndarray_scalar_dtype(
-                    dst_type_info.np_number_type
-                )
-                _ = DtypeRegistry.validate_dtype_and_get_kind(scalar_dtype)
-                return np.array(value, dtype=scalar_dtype)
-
-        return decode_vector
+    if src_type_kind == "Union":
+        return lambda value: value[1]
 
     return lambda value: value
 

cocoindex/flow.py

@@ -92,6 +92,7 @@ def _spec_kind(spec: Any) -> str:
 
 
 T = TypeVar("T")
+S = TypeVar("S")
 
 
 class _DataSliceState:
@@ -185,7 +186,7 @@ class DataSlice(Generic[T]):
 
     def transform(
         self, fn_spec: op.FunctionSpec, *args: Any, **kwargs: Any
-    ) -> DataSlice[T]:
+    ) -> DataSlice[Any]:
         """
         Apply a function to the data slice.
         """
@@ -216,7 +217,7 @@ class DataSlice(Generic[T]):
             ),
         )
 
-    def call(self, func: Callable[[DataSlice[T]], T], *args: Any, **kwargs: Any) -> T:
+    def call(self, func: Callable[..., S], *args: Any, **kwargs: Any) -> S:
         """
         Call a function with the data slice.
         """

cocoindex/functions.py

@@ -32,6 +32,16 @@ class SplitRecursively(op.FunctionSpec):
     custom_languages: list[CustomLanguageSpec] = dataclasses.field(default_factory=list)
 
 
+class EmbedText(op.FunctionSpec):
+    """Embed a text into a vector space."""
+
+    api_type: llm.LlmApiType
+    model: str
+    address: str | None = None
+    output_dimension: int | None = None
+    task_type: str | None = None
+
+
 class ExtractByLlm(op.FunctionSpec):
     """Extract information from a text using a LLM."""
 

cocoindex/llm.py

@@ -11,6 +11,7 @@ class LlmApiType(Enum):
     ANTHROPIC = "Anthropic"
     LITE_LLM = "LiteLlm"
     OPEN_ROUTER = "OpenRouter"
+    VOYAGE = "Voyage"
 
 
 @dataclass

cocoindex/tests/__init__.py

@@ -1 +0,0 @@
-

cocoindex/tests/test_convert.py

@@ -91,23 +91,26 @@ def validate_full_roundtrip(
     """
     from cocoindex import _engine  # type: ignore
 
+    def eq(a: Any, b: Any) -> bool:
+        if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
+            return np.array_equal(a, b)
+        return type(a) == type(b) and not not (a == b)
+
     encoded_value = encode_engine_value(value)
     value_type = value_type or type(value)
     encoded_output_type = encode_enriched_type(value_type)["type"]
     value_from_engine = _engine.testutil.seder_roundtrip(
         encoded_value, encoded_output_type
     )
-    decoded_value = build_engine_value_decoder(value_type, value_type)(
-        value_from_engine
-    )
-    np.testing.assert_array_equal(decoded_value, value)
+    decoder = make_engine_value_decoder([], encoded_output_type, value_type)
+    decoded_value = decoder(value_from_engine)
+    assert eq(decoded_value, value)
 
     if other_decoded_values is not None:
         for other_value, other_type in other_decoded_values:
-            other_decoded_value = build_engine_value_decoder(other_type, other_type)(
-                value_from_engine
-            )
-            np.testing.assert_array_equal(other_decoded_value, other_value)
+            decoder = make_engine_value_decoder([], encoded_output_type, other_type)
+            other_decoded_value = decoder(value_from_engine)
+            assert eq(other_decoded_value, other_value)
 
 
 def test_encode_engine_value_basic_types() -> None:
@@ -215,19 +218,38 @@ def test_encode_engine_value_none() -> None:
 
 
 def test_roundtrip_basic_types() -> None:
-    validate_full_roundtrip(42, int)
+    validate_full_roundtrip(42, int, (42, None))
     validate_full_roundtrip(3.25, float, (3.25, Float64))
-    validate_full_roundtrip(3.25, Float64, (3.25, float))
-    validate_full_roundtrip(3.25, Float32)
-    validate_full_roundtrip("hello", str)
-    validate_full_roundtrip(True, bool)
-    validate_full_roundtrip(False, bool)
-    validate_full_roundtrip(datetime.date(2025, 1, 1), datetime.date)
-    validate_full_roundtrip(datetime.datetime.now(), cocoindex.LocalDateTime)
     validate_full_roundtrip(
-        datetime.datetime.now(datetime.UTC), cocoindex.OffsetDateTime
+        3.25, Float64, (3.25, float), (np.float64(3.25), np.float64)
+    )
+    validate_full_roundtrip(
+        3.25, Float32, (3.25, float), (np.float32(3.25), np.float32)
+    )
+    validate_full_roundtrip("hello", str, ("hello", None))
+    validate_full_roundtrip(True, bool, (True, None))
+    validate_full_roundtrip(False, bool, (False, None))
+    validate_full_roundtrip(
+        datetime.date(2025, 1, 1), datetime.date, (datetime.date(2025, 1, 1), None)
+    )
+
+    validate_full_roundtrip(
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456),
+        cocoindex.LocalDateTime,
+        (datetime.datetime(2025, 1, 2, 3, 4, 5, 123456), datetime.datetime),
+    )
+    validate_full_roundtrip(
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, datetime.UTC),
+        cocoindex.OffsetDateTime,
+        (
+            datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, datetime.UTC),
+            datetime.datetime,
+        ),
     )
 
+    uuid_value = uuid.uuid4()
+    validate_full_roundtrip(uuid_value, uuid.UUID, (uuid_value, None))
+
 
 def test_decode_scalar_numpy_values() -> None:
     test_cases = [
@@ -549,6 +571,48 @@ def test_field_position_cases(
     assert decoder(engine_val) == PythonOrder(**expected_dict)
 
 
+def test_roundtrip_union_simple() -> None:
+    t = int | str | float
+    value = 10.4
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_with_active_uuid() -> None:
+    t = str | uuid.UUID | int
+    value = uuid.uuid4().bytes
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_with_inactive_uuid() -> None:
+    t = str | uuid.UUID | int
+    value = "5a9f8f6a-318f-4f1f-929d-566d7444a62d"  # it's a string
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_offset_datetime() -> None:
+    t = str | uuid.UUID | float | int | datetime.datetime
+    value = datetime.datetime.now(datetime.UTC)
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_date() -> None:
+    t = str | uuid.UUID | float | int | datetime.date
+    value = datetime.date.today()
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_time() -> None:
+    t = str | uuid.UUID | float | int | datetime.time
+    value = datetime.time()
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_timedelta() -> None:
+    t = str | uuid.UUID | float | int | datetime.timedelta
+    value = datetime.timedelta(hours=39, minutes=10, seconds=1)
+    validate_full_roundtrip(value, t)
+
+
 def test_roundtrip_ltable() -> None:
     t = list[Order]
     value = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]
@@ -807,37 +871,72 @@ def test_dump_vector_type_annotation_no_dim() -> None:
 
 def test_full_roundtrip_vector_numeric_types() -> None:
     """Test full roundtrip for numeric vector types using NDArray."""
-    value_f32: Vector[np.float32, Literal[3]] = np.array(
-        [1.0, 2.0, 3.0], dtype=np.float32
+    value_f32 = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    validate_full_roundtrip(
+        value_f32,
+        Vector[np.float32, Literal[3]],
+        ([np.float32(1.0), np.float32(2.0), np.float32(3.0)], list[np.float32]),
+        ([1.0, 2.0, 3.0], list[cocoindex.Float32]),
+        ([1.0, 2.0, 3.0], list[float]),
     )
-    validate_full_roundtrip(value_f32, Vector[np.float32, Literal[3]])
-    value_f64: Vector[np.float64, Literal[3]] = np.array(
-        [1.0, 2.0, 3.0], dtype=np.float64
+    validate_full_roundtrip(
+        value_f32,
+        np.typing.NDArray[np.float32],
+        ([np.float32(1.0), np.float32(2.0), np.float32(3.0)], list[np.float32]),
+        ([1.0, 2.0, 3.0], list[cocoindex.Float32]),
+        ([1.0, 2.0, 3.0], list[float]),
     )
-    validate_full_roundtrip(value_f64, Vector[np.float64, Literal[3]])
-    value_i64: Vector[np.int64, Literal[3]] = np.array([1, 2, 3], dtype=np.int64)
-    validate_full_roundtrip(value_i64, Vector[np.int64, Literal[3]])
-    value_i32: Vector[np.int32, Literal[3]] = np.array([1, 2, 3], dtype=np.int32)
+    validate_full_roundtrip(
+        value_f32.tolist(),
+        list[np.float32],
+        (value_f32, Vector[np.float32, Literal[3]]),
+        ([1.0, 2.0, 3.0], list[cocoindex.Float32]),
+        ([1.0, 2.0, 3.0], list[float]),
+    )
+
+    value_f64 = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    validate_full_roundtrip(
+        value_f64,
+        Vector[np.float64, Literal[3]],
+        ([np.float64(1.0), np.float64(2.0), np.float64(3.0)], list[np.float64]),
+        ([1.0, 2.0, 3.0], list[cocoindex.Float64]),
+        ([1.0, 2.0, 3.0], list[float]),
+    )
+
+    value_i64 = np.array([1, 2, 3], dtype=np.int64)
+    validate_full_roundtrip(
+        value_i64,
+        Vector[np.int64, Literal[3]],
+        ([np.int64(1), np.int64(2), np.int64(3)], list[np.int64]),
+        ([1, 2, 3], list[int]),
+    )
+
+    value_i32 = np.array([1, 2, 3], dtype=np.int32)
     with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_i32, Vector[np.int32, Literal[3]])
-    value_u8: Vector[np.uint8, Literal[3]] = np.array([1, 2, 3], dtype=np.uint8)
+    value_u8 = np.array([1, 2, 3], dtype=np.uint8)
     with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u8, Vector[np.uint8, Literal[3]])
-    value_u16: Vector[np.uint16, Literal[3]] = np.array([1, 2, 3], dtype=np.uint16)
+    value_u16 = np.array([1, 2, 3], dtype=np.uint16)
     with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u16, Vector[np.uint16, Literal[3]])
-    value_u32: Vector[np.uint32, Literal[3]] = np.array([1, 2, 3], dtype=np.uint32)
+    value_u32 = np.array([1, 2, 3], dtype=np.uint32)
     with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u32, Vector[np.uint32, Literal[3]])
-    value_u64: Vector[np.uint64, Literal[3]] = np.array([1, 2, 3], dtype=np.uint64)
+    value_u64 = np.array([1, 2, 3], dtype=np.uint64)
     with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u64, Vector[np.uint64, Literal[3]])
 
 
 def test_roundtrip_vector_no_dimension() -> None:
     """Test full roundtrip for vector types without dimension annotation."""
-    value_f64: Vector[np.float64] = np.array([1.0, 2.0, 3.0], dtype=np.float64)
-    validate_full_roundtrip(value_f64, Vector[np.float64])
+    value_f64 = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    validate_full_roundtrip(
+        value_f64,
+        Vector[np.float64],
+        ([1.0, 2.0, 3.0], list[float]),
+        (np.array([1.0, 2.0, 3.0], dtype=np.float64), np.typing.NDArray[np.float64]),
+    )
 
 
 def test_roundtrip_string_vector() -> None:
@@ -862,9 +961,9 @@ def test_roundtrip_dimension_mismatch() -> None:
 def test_full_roundtrip_scalar_numeric_types() -> None:
     """Test full roundtrip for scalar NumPy numeric types."""
     # Test supported scalar types
-    validate_full_roundtrip(np.int64(42), np.int64)
-    validate_full_roundtrip(np.float32(3.14), np.float32)
-    validate_full_roundtrip(np.float64(2.718), np.float64)
+    validate_full_roundtrip(np.int64(42), np.int64, (42, int))
+    validate_full_roundtrip(np.float32(3.25), np.float32, (3.25, cocoindex.Float32))
+    validate_full_roundtrip(np.float64(3.25), np.float64, (3.25, cocoindex.Float64))
 
     # Test unsupported scalar types
     for unsupported_type in [np.int32, np.uint8, np.uint16, np.uint32, np.uint64]:

cocoindex/tests/test_typing.py

@@ -2,7 +2,7 @@ import dataclasses
 import datetime
 import uuid
 from collections.abc import Mapping, Sequence
-from typing import Annotated, Any, Dict, List, Literal, NamedTuple, get_args, get_origin
+from typing import Annotated, Any, Literal, NamedTuple, get_args, get_origin
 
 import numpy as np
 import pytest
@@ -162,10 +162,11 @@ def test_ndarray_any_dtype() -> None:
 
 
 def test_list_of_primitives() -> None:
-    typ = List[str]
+    typ = list[str]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=list[str],
         vector_info=VectorInfo(dim=None),
         elem_type=str,
         key_type=None,
@@ -177,10 +178,11 @@ def test_list_of_primitives() -> None:
 
 
 def test_list_of_structs() -> None:
-    typ = List[SimpleDataclass]
+    typ = list[SimpleDataclass]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="LTable",
+        core_type=list[SimpleDataclass],
         vector_info=None,
         elem_type=SimpleDataclass,
         key_type=None,
@@ -196,6 +198,7 @@ def test_sequence_of_int() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=Sequence[int],
         vector_info=VectorInfo(dim=None),
         elem_type=int,
         key_type=None,
@@ -207,10 +210,11 @@ def test_sequence_of_int() -> None:
 
 
 def test_list_with_vector_info() -> None:
-    typ = Annotated[List[int], VectorInfo(dim=5)]
+    typ = Annotated[list[int], VectorInfo(dim=5)]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=list[int],
         vector_info=VectorInfo(dim=5),
         elem_type=int,
         key_type=None,
@@ -222,10 +226,11 @@ def test_list_with_vector_info() -> None:
 
 
 def test_dict_str_int() -> None:
-    typ = Dict[str, int]
+    typ = dict[str, int]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="KTable",
+        core_type=dict[str, int],
         vector_info=None,
         elem_type=(str, int),
         key_type=None,
@@ -241,6 +246,7 @@ def test_mapping_str_dataclass() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="KTable",
+        core_type=Mapping[str, SimpleDataclass],
         vector_info=None,
         elem_type=(str, SimpleDataclass),
         key_type=None,
@@ -256,6 +262,7 @@ def test_dataclass() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Struct",
+        core_type=SimpleDataclass,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -271,6 +278,7 @@ def test_named_tuple() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Struct",
+        core_type=SimpleNamedTuple,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -286,6 +294,7 @@ def test_tuple_key_value() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Int64",
+        core_type=int,
         vector_info=None,
         elem_type=None,
         key_type=str,
@@ -301,6 +310,7 @@ def test_str() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Str",
+        core_type=str,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -316,6 +326,7 @@ def test_bool() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Bool",
+        core_type=bool,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -331,6 +342,7 @@ def test_bytes() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Bytes",
+        core_type=bytes,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -346,6 +358,7 @@ def test_uuid() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Uuid",
+        core_type=uuid.UUID,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -361,6 +374,7 @@ def test_date() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Date",
+        core_type=datetime.date,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -376,6 +390,7 @@ def test_time() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Time",
+        core_type=datetime.time,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -391,6 +406,7 @@ def test_timedelta() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="TimeDelta",
+        core_type=datetime.timedelta,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -406,6 +422,7 @@ def test_float() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Float64",
+        core_type=float,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -421,6 +438,7 @@ def test_int() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Int64",
+        core_type=int,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -436,6 +454,7 @@ def test_type_with_attributes() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Str",
+        core_type=str,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -472,7 +491,7 @@ def test_encode_enriched_type_vector() -> None:
 
 
 def test_encode_enriched_type_ltable() -> None:
-    typ = List[SimpleDataclass]
+    typ = list[SimpleDataclass]
     result = encode_enriched_type(typ)
     assert result["type"]["kind"] == "LTable"
     assert result["type"]["row"]["kind"] == "Struct"
@@ -513,7 +532,7 @@ def test_invalid_struct_kind() -> None:
 
 
 def test_invalid_list_kind() -> None:
-    typ = Annotated[List[int], TypeKind("Struct")]
+    typ = Annotated[list[int], TypeKind("Struct")]
     with pytest.raises(ValueError, match="Unexpected type kind for list: Struct"):
         analyze_type_info(typ)
 

cocoindex/typing.py

@@ -150,6 +150,7 @@ class AnalyzedTypeInfo:
     """
 
     kind: str
+    core_type: Any
     vector_info: VectorInfo | None  # For Vector
     elem_type: ElementType | None  # For Vector and Table
 
@@ -161,6 +162,7 @@ class AnalyzedTypeInfo:
 
     attrs: dict[str, Any] | None
     nullable: bool = False
+    union_variant_types: typing.List[ElementType] | None = None  # For Union
 
 
 def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
@@ -181,18 +183,6 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         if base_type is Annotated:
             annotations = t.__metadata__
             t = t.__origin__
-        elif base_type is types.UnionType:
-            possible_types = typing.get_args(t)
-            non_none_types = [
-                arg for arg in possible_types if arg not in (None, types.NoneType)
-            ]
-            if len(non_none_types) != 1:
-                raise ValueError(
-                    f"Expect exactly one non-None choice for Union type, but got {len(non_none_types)}: {t}"
-                )
-            t = non_none_types[0]
-            if len(possible_types) > 1:
-                nullable = True
         else:
             break
 
@@ -211,6 +201,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
     struct_type: type | None = None
     elem_type: ElementType | None = None
+    union_variant_types: typing.List[ElementType] | None = None
     key_type: type | None = None
     np_number_type: type | None = None
     if _is_struct_type(t):
@@ -251,6 +242,24 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         args = typing.get_args(t)
         elem_type = (args[0], args[1])
         kind = "KTable"
+    elif base_type is types.UnionType:
+        possible_types = typing.get_args(t)
+        non_none_types = [
+            arg for arg in possible_types if arg not in (None, types.NoneType)
+        ]
+
+        if len(non_none_types) == 0:
+            return analyze_type_info(None)
+
+        nullable = len(non_none_types) < len(possible_types)
+
+        if len(non_none_types) == 1:
+            result = analyze_type_info(non_none_types[0])
+            result.nullable = nullable
+            return result
+
+        kind = "Union"
+        union_variant_types = non_none_types
     elif kind is None:
         if t is bytes:
             kind = "Bytes"
@@ -277,8 +286,10 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
     return AnalyzedTypeInfo(
         kind=kind,
+        core_type=t,
         vector_info=vector_info,
         elem_type=elem_type,
+        union_variant_types=union_variant_types,
         key_type=key_type,
         struct_type=struct_type,
         np_number_type=np_number_type,
@@ -338,6 +349,14 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         encoded_type["element_type"] = _encode_type(elem_type_info)
         encoded_type["dimension"] = type_info.vector_info.dim
 
+    elif type_info.kind == "Union":
+        if type_info.union_variant_types is None:
+            raise ValueError("Union type must have a variant type list")
+        encoded_type["types"] = [
+            _encode_type(analyze_type_info(typ))
+            for typ in type_info.union_variant_types
+        ]
+
     elif type_info.kind in TABLE_TYPES:
         if type_info.elem_type is None:
             raise ValueError(f"{type_info.kind} type must have an element type")

{cocoindex-0.1.53.dist-info → cocoindex-0.1.54.dist-info}/METADATA

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.53
+Version: 0.1.54
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: pytest ; extra == 'test'
 Requires-Dist: ruff ; extra == 'dev'
+Requires-Dist: pre-commit ; extra == 'dev'
 Provides-Extra: test
 Provides-Extra: dev
 License-File: LICENSE
@@ -51,10 +52,10 @@ Unlike a workflow orchestration framework where data is usually opaque, in CocoI
 
 ```python
 # import
-data['content'] = flow_builder.add_source(...) 
+data['content'] = flow_builder.add_source(...)
 
 # transform
-data['out'] = data['content'] 
+data['out'] = data['content']
     .transform(...)
     .transform(...)
 
@@ -75,17 +76,17 @@ As a data framework, CocoIndex takes it to the next level on data freshness. **I
 The frameworks takes care of
 - Change data capture.
 - Figure out what exactly needs to be updated, and only updating that without having to recompute everything.
-  
+
 This makes it fast to reflect any source updates to the target store. If you have concerns with surfacing stale data to AI agents and are spending lots of efforts working on infra piece to optimize the latency, the framework actually handles it for you.
 
 
 ## Quick Start:
-If you're new to CocoIndex, we recommend checking out 
+If you're new to CocoIndex, we recommend checking out
 - 📖 [Documentation](https://cocoindex.io/docs)
 - ⚡  [Quick Start Guide](https://cocoindex.io/docs/getting_started/quickstart)
-- 🎬 [Quick Start Video Tutorial](https://youtu.be/gv5R8nOXsWU?si=9ioeKYkMEnYevTXT) 
+- 🎬 [Quick Start Video Tutorial](https://youtu.be/gv5R8nOXsWU?si=9ioeKYkMEnYevTXT)
 
-### Setup 
+### Setup
 
 1. Install CocoIndex Python library
 
@@ -155,8 +156,8 @@ It defines an index flow like this:
 | [Google Drive Text Embedding](examples/gdrive_text_embedding) | Index text documents from Google Drive |
 | [Docs to Knowledge Graph](examples/docs_to_knowledge_graph) | Extract relationships from Markdown documents and build a knowledge graph |
 | [Embeddings to Qdrant](examples/text_embedding_qdrant) | Index documents in a Qdrant collection for semantic search |
-| [FastAPI Server with Docker](examples/fastapi_server_docker) | Run the semantic search server in a Dockerized FastAPI setup | 
-| [Product Recommendation](examples/product_recommendation) | Build real-time product recommendations with LLM and graph database| 
+| [FastAPI Server with Docker](examples/fastapi_server_docker) | Run the semantic search server in a Dockerized FastAPI setup |
+| [Product Recommendation](examples/product_recommendation) | Build real-time product recommendations with LLM and graph database|
 | [Image Search with Vision API](examples/image_search) | Generates detailed captions for images using a vision model, embeds them, enables live-updating semantic search via FastAPI and served on a React frontend|
 
 More coming and stay tuned 👀!
@@ -178,7 +179,7 @@ Join our community here:
 - 📜 [Read our blog posts](https://cocoindex.io/blogs/)
 
 ## Support us:
-We are constantly improving, and more features and examples are coming soon. If you love this project, please drop us a star ⭐ at GitHub repo [![GitHub](https://img.shields.io/github/stars/cocoindex-io/cocoindex?color=5B5BD6)](https://github.com/cocoindex-io/cocoindex) to stay tuned and help us grow. 
+We are constantly improving, and more features and examples are coming soon. If you love this project, please drop us a star ⭐ at GitHub repo [![GitHub](https://img.shields.io/github/stars/cocoindex-io/cocoindex?color=5B5BD6)](https://github.com/cocoindex-io/cocoindex) to stay tuned and help us grow.
 
 ## License
 CocoIndex is Apache 2.0 licensed.

cocoindex-0.1.54.dist-info/RECORD

@@ -0,0 +1,28 @@
+cocoindex-0.1.54.dist-info/METADATA,sha256=YhrfcjAPuh3OkyFCVB_tkaT_dY9NV7HMhOHIQVBqnCk,9907
+cocoindex-0.1.54.dist-info/WHEEL,sha256=YLLGMgiebwxzWB63F0KVTBu1H1ad3TdbsHYfFKfk6fc,104
+cocoindex-0.1.54.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.54.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex/__init__.py,sha256=3potlGouaGAznMTEqU5QKuB934qpD1xR6ZH6EL3YMWw,1782
+cocoindex/_engine.cpython-313-darwin.so,sha256=VwEk8aK1CW5YXpFr2tv46tLcD1DRCYamkYLrrBseNK0,57458240
+cocoindex/auth_registry.py,sha256=1XqO7ibjmBBd8i11XSJTvTgdz8p1ptW-ZpuSgo_5zzk,716
+cocoindex/cli.py,sha256=yXlUZqJuB9baofnpHFOu3S-0x0xGglmJ88nd9_T-fEs,18223
+cocoindex/convert.py,sha256=4-a0uWLOfQ_e2OOwHjIYAsBMZIDRnGbK01Ye1E2ybIQ,10377
+cocoindex/flow.py,sha256=djpGsQlSZV10OCS_81GY9k8q_0l-6MeismLnoF9XteE,30023
+cocoindex/functions.py,sha256=Lh1aSovd6HaikAWvDi6dX0poyma1MqBX5S9FAzIR8c4,2622
+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/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cocoindex/runtime.py,sha256=bAdHYaXFWiiUWyAgzmKTeaAaRR0D_AmaqVCIdPO-v00,1056
+cocoindex/setting.py,sha256=Zl8K86r8RVvG9c3pCsH0Ot8BHhDAQAQCjoBp7TnXMLQ,3590
+cocoindex/setup.py,sha256=u5dYZFKfz4yZLiGHD0guNaR0s4zY9JAoZWrWHpAHw_0,773
+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=RxbCmW3cSZQVJtAWZBBH9YNAViWVSi7u0JO8JGJl7M8,36174
+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.54.dist-info/RECORD,,

cocoindex-0.1.53.dist-info/RECORD

@@ -1,28 +0,0 @@
-cocoindex-0.1.53.dist-info/METADATA,sha256=HQKpE-cc_Dgi2zpIZOAX8aVCl5Zt8hd__egLbAdlOSU,9874
-cocoindex-0.1.53.dist-info/WHEEL,sha256=YLLGMgiebwxzWB63F0KVTBu1H1ad3TdbsHYfFKfk6fc,104
-cocoindex-0.1.53.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.53.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-cocoindex/__init__.py,sha256=3potlGouaGAznMTEqU5QKuB934qpD1xR6ZH6EL3YMWw,1782
-cocoindex/_engine.cpython-313-darwin.so,sha256=RNI4THwgDUk7SquPLUz81q0MH73WqAYDkESTApebEFU,57345008
-cocoindex/auth_registry.py,sha256=1XqO7ibjmBBd8i11XSJTvTgdz8p1ptW-ZpuSgo_5zzk,716
-cocoindex/cli.py,sha256=joSucqLibjZHs27fwNPzuMPGLeVJJZqaUbeSM-OnEyA,18223
-cocoindex/convert.py,sha256=qP_LRTg_B8YPEeHprnX4kl-QV4ks0oqQzH7A1totZRA,9787
-cocoindex/flow.py,sha256=uQKvIyWSysSdhNFI87QD6f2kbF-2R6zFMLZWBEk4rZU,30015
-cocoindex/functions.py,sha256=9A61Jj5a3vQoI2MIAhjXvJrDxSzDhe6VncQWbiVtwcg,2393
-cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
-cocoindex/lib.py,sha256=BeRUn3RqE_wSsVtsgCzbFFKe1LXgRyRmMOcmwWBuEXo,2940
-cocoindex/llm.py,sha256=gbjAsD8Y9QZyICCQfSDZYQU2pIA6GqG4jKbRpbnQniA,408
-cocoindex/op.py,sha256=Z7V9Fdz4qeTeozzKmp1Dk1lPUP0GBgVgD_vBEhTJS5Y,11810
-cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/runtime.py,sha256=bAdHYaXFWiiUWyAgzmKTeaAaRR0D_AmaqVCIdPO-v00,1056
-cocoindex/setting.py,sha256=Zl8K86r8RVvG9c3pCsH0Ot8BHhDAQAQCjoBp7TnXMLQ,3590
-cocoindex/setup.py,sha256=u5dYZFKfz4yZLiGHD0guNaR0s4zY9JAoZWrWHpAHw_0,773
-cocoindex/sources.py,sha256=JCnOhv1w4o28e03i7yvo4ESicWYAhckkBg5bQlxNH4U,1330
-cocoindex/targets.py,sha256=Nfh_tpFd1goTnS_cxBjIs4j9zl3Z4Z1JomAQ1dl3Sic,2796
-cocoindex/tests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-cocoindex/tests/test_convert.py,sha256=iDA4j-X_r1KeJgaa3osU7SyY8rBc57R0YVW8G4Z83YI,33235
-cocoindex/tests/test_optional_database.py,sha256=snAmkNa6wtOSaxoZE1HgjvL5v_ylitt3Jt_9df4Cgdc,8506
-cocoindex/tests/test_typing.py,sha256=jeZFUyzUTC6hdyo1uERhlKaKAiBKeKCrZ20OLGBkluo,14184
-cocoindex/typing.py,sha256=yLtXdXEF0sYe3cMxTR__g39Gi286CnOKGI0T-hjaCrw,11999
-cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
-cocoindex-0.1.53.dist-info/RECORD,,