cocoindex 0.1.52__cp313-cp313-macosx_10_12_x86_64.whl → 0.1.54__cp313-cp313-macosx_10_12_x86_64.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

@@ -6,17 +6,20 @@ import dataclasses
 import datetime
 import inspect
 import uuid
+from enum import Enum
+from typing import Any, Callable, Mapping, get_origin
+
 import numpy as np
 
-from enum import Enum
-from typing import Any, Callable, get_origin, Mapping
 from .typing import (
+    KEY_FIELD_NAME,
+    TABLE_TYPES,
+    AnalyzedTypeInfo,
+    DtypeRegistry,
     analyze_type_info,
     encode_enriched_type,
+    extract_ndarray_scalar_dtype,
     is_namedtuple_type,
-    TABLE_TYPES,
-    KEY_FIELD_NAME,
-    DtypeRegistry,
 )
 
 
@@ -29,6 +32,8 @@ def encode_engine_value(value: Any) -> Any:
         ]
     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)):
@@ -42,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],
@@ -61,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)}`."
@@ -73,14 +103,62 @@ 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
+
+        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)
+
+        return decode_scalar
 
-    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})"
+    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
         )
 
+        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_vector(value: Any) -> Any | 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 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_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
@@ -123,39 +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":
-        dtype_info = DtypeRegistry.get_by_dtype(dst_type_info.np_number_type)
-
-        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 vector `{''.join(field_path)}`"
-                )
-
-            if not isinstance(value, (np.ndarray, list)):
-                raise TypeError(
-                    f"Expected NDArray or list for vector `{''.join(field_path)}`, got {type(value)}"
-                )
-            expected_dim = (
-                dst_type_info.vector_info.dim if dst_type_info.vector_info else None
-            )
-            if expected_dim is not None and len(value) != expected_dim:
-                raise ValueError(
-                    f"Vector dimension mismatch for `{''.join(field_path)}`: "
-                    f"expected {expected_dim}, got {len(value)}"
-                )
-
-            # Use NDArray for supported numeric dtypes, else return list
-            if dtype_info is not None:
-                return np.array(value, dtype=dtype_info.numpy_dtype)
-            return value
-
-        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

@@ -9,6 +9,9 @@ class LlmApiType(Enum):
     OLLAMA = "Ollama"
     GEMINI = "Gemini"
     ANTHROPIC = "Anthropic"
+    LITE_LLM = "LiteLlm"
+    OPEN_ROUTER = "OpenRouter"
+    VOYAGE = "Voyage"
 
 
 @dataclass

cocoindex/tests/__init__.py

@@ -1 +0,0 @@
-