cocoindex 0.1.26__cp312-cp312-macosx_11_0_arm64.whl → 0.1.27__cp312-cp312-macosx_11_0_arm64.whl

cocoindex/__init__.py

@@ -9,4 +9,5 @@ from .llm import LlmSpec, LlmApiType
 from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
 from .auth_registry import AuthEntryReference, add_auth_entry, ref_auth_entry
 from .lib import *
-from ._engine import OpArgSchema
+from ._engine import OpArgSchema
+from .typing import Float32, Float64, LocalDateTime, OffsetDateTime, Range, Vector, Json

cocoindex/convert.py

@@ -8,25 +8,28 @@ import uuid
 
 from enum import Enum
 from typing import Any, Callable, get_origin
-from .typing import analyze_type_info, encode_enriched_type, COLLECTION_TYPES
+from .typing import analyze_type_info, encode_enriched_type, TABLE_TYPES, KEY_FIELD_NAME
 
-def to_engine_value(value: Any) -> Any:
-    """Convert a Python value to an engine value."""
+
+def encode_engine_value(value: Any) -> Any:
+    """Encode a Python value to an engine value."""
     if dataclasses.is_dataclass(value):
-        return [to_engine_value(getattr(value, f.name)) for f in dataclasses.fields(value)]
+        return [encode_engine_value(getattr(value, f.name)) for f in dataclasses.fields(value)]
     if isinstance(value, (list, tuple)):
-        return [to_engine_value(v) for v in value]
+        return [encode_engine_value(v) for v in value]
+    if isinstance(value, dict):
+        return [[encode_engine_value(k)] + encode_engine_value(v) for k, v in value.items()]
     if isinstance(value, uuid.UUID):
         return value.bytes
     return value
 
-def make_engine_value_converter(
+def make_engine_value_decoder(
         field_path: list[str],
         src_type: dict[str, Any],
         dst_annotation,
     ) -> Callable[[Any], Any]:
     """
-    Make a converter from an engine value to a Python value.
+    Make a decoder from an engine value to a Python value.
 
     Args:
         field_path: The path to the field in the engine value. For error messages.
@@ -34,13 +37,13 @@ def make_engine_value_converter(
         dst_annotation: The type annotation of the Python value.
 
     Returns:
-        A converter from an engine value to a Python value.
+        A decoder from an engine value to a Python value.
     """
 
     src_type_kind = src_type['kind']
 
     if dst_annotation is inspect.Parameter.empty:
-        if src_type_kind == 'Struct' or src_type_kind in COLLECTION_TYPES:
+        if src_type_kind == 'Struct' or src_type_kind in TABLE_TYPES:
             raise ValueError(f"Missing type annotation for `{''.join(field_path)}`."
                              f"It's required for {src_type_kind} type.")
         return lambda value: value
@@ -53,41 +56,59 @@ def make_engine_value_converter(
             f"passed in {src_type_kind}, declared {dst_annotation} ({dst_type_info.kind})")
 
     if dst_type_info.dataclass_type is not None:
-        return _make_engine_struct_value_converter(
+        return _make_engine_struct_value_decoder(
             field_path, src_type['fields'], dst_type_info.dataclass_type)
 
-    if src_type_kind in COLLECTION_TYPES:
+    if src_type_kind in TABLE_TYPES:
         field_path.append('[*]')
         elem_type_info = analyze_type_info(dst_type_info.elem_type)
         if elem_type_info.dataclass_type is None:
             raise ValueError(f"Type mismatch for `{''.join(field_path)}`: "
-                             f"declared `{dst_type_info.kind}`, a dataclass type expected")
-        elem_converter = _make_engine_struct_value_converter(
-            field_path, src_type['row']['fields'], elem_type_info.dataclass_type)
+                            f"declared `{dst_type_info.kind}`, a dataclass type expected")
+        engine_fields_schema = src_type['row']['fields']
+        if elem_type_info.key_type is not None:
+            key_field_schema = engine_fields_schema[0]
+            field_path.append(f".{key_field_schema.get('name', KEY_FIELD_NAME)}")
+            key_decoder = make_engine_value_decoder(
+                field_path, key_field_schema['type'], elem_type_info.key_type)
+            field_path.pop()
+            value_decoder = _make_engine_struct_value_decoder(
+                field_path, engine_fields_schema[1:], elem_type_info.dataclass_type)
+            def decode(value):
+                if value is None:
+                    return None
+                return {key_decoder(v[0]): value_decoder(v[1:]) for v in value}
+        else:
+            elem_decoder = _make_engine_struct_value_decoder(
+                field_path, engine_fields_schema, elem_type_info.dataclass_type)
+            def decode(value):
+                if value is None:
+                    return None
+                return [elem_decoder(v) for v in value]
         field_path.pop()
-        return lambda value: [elem_converter(v) for v in value] if value is not None else None
+        return decode
 
     if src_type_kind == 'Uuid':
         return lambda value: uuid.UUID(bytes=value)
 
     return lambda value: value
 
-def _make_engine_struct_value_converter(
+def _make_engine_struct_value_decoder(
         field_path: list[str],
         src_fields: list[dict[str, Any]],
         dst_dataclass_type: type,
     ) -> Callable[[list], Any]:
-    """Make a converter from an engine field values to a Python value."""
+    """Make a decoder from an engine field values to a Python value."""
 
     src_name_to_idx = {f['name']: i for i, f in enumerate(src_fields)}
     def make_closure_for_value(name: str, param: inspect.Parameter) -> Callable[[list], Any]:
         src_idx = src_name_to_idx.get(name)
         if src_idx is not None:
             field_path.append(f'.{name}')
-            field_converter = make_engine_value_converter(
+            field_decoder = make_engine_value_decoder(
                 field_path, src_fields[src_idx]['type'], param.annotation)
             field_path.pop()
-            return lambda values: field_converter(values[src_idx])
+            return lambda values: field_decoder(values[src_idx])
 
         default_value = param.default
         if default_value is inspect.Parameter.empty:
@@ -96,12 +117,12 @@ def _make_engine_struct_value_converter(
 
         return lambda _: default_value
 
-    field_value_converters = [
+    field_value_decoder = [
         make_closure_for_value(name, param)
         for (name, param) in inspect.signature(dst_dataclass_type).parameters.items()]
 
     return lambda values: dst_dataclass_type(
-        *(converter(values) for converter in field_value_converters))
+        *(decoder(values) for decoder in field_value_decoder))
 
 def dump_engine_object(v: Any) -> Any:
     """Recursively dump an object for engine. Engine side uses `Pythonized` to catch."""

cocoindex/flow.py

@@ -142,9 +142,9 @@ class DataSlice:
 
     def row(self) -> DataScope:
         """
-        Return a scope representing each entry of the collection.
+        Return a scope representing each row of the table.
         """
-        row_scope = self._state.engine_data_slice.collection_entry_scope()
+        row_scope = self._state.engine_data_slice.table_row_scope()
         return DataScope(self._state.flow_builder_state, row_scope)
 
     def for_each(self, f: Callable[[DataScope], None]) -> None:

cocoindex/functions.py

@@ -41,7 +41,7 @@ class SentenceTransformerEmbedExecutor:
         args = self.spec.args or {}
         self._model = sentence_transformers.SentenceTransformer(self.spec.model, **args)
         dim = self._model.get_sentence_embedding_dimension()
-        return Annotated[list[Float32], Vector(dim=dim), TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value)]
+        return Annotated[Vector[Float32, dim], TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value)]
 
     def __call__(self, text: str) -> list[Float32]:
         return self._model.encode(text).tolist()

cocoindex/lib.py

@@ -1,7 +1,6 @@
 """
 Library level functions and states.
 """
-import asyncio
 import os
 import sys
 import functools
@@ -12,6 +11,7 @@ from dataclasses import dataclass
 
 from . import _engine
 from . import flow, query, cli
+from .convert import dump_engine_object
 
 
 def _load_field(target: dict[str, str], name: str, env_name: str, required: bool = False):
@@ -22,24 +22,32 @@ def _load_field(target: dict[str, str], name: str, env_name: str, required: bool
     else:
         target[name] = value
 
+@dataclass
+class DatabaseConnectionSpec:
+    url: str
+    user: str | None = None
+    password: str | None = None
+
 @dataclass
 class Settings:
     """Settings for the cocoindex library."""
-    database_url: str
+    database: DatabaseConnectionSpec
 
     @classmethod
     def from_env(cls) -> Self:
         """Load settings from environment variables."""
 
-        kwargs: dict[str, str] = dict()
-        _load_field(kwargs, "database_url", "COCOINDEX_DATABASE_URL", required=True)
-
-        return cls(**kwargs)
+        db_kwargs: dict[str, str] = dict()
+        _load_field(db_kwargs, "url", "COCOINDEX_DATABASE_URL", required=True)
+        _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
+        _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+        database = DatabaseConnectionSpec(**db_kwargs)
+        return cls(database=database)
 
 
 def init(settings: Settings):
     """Initialize the cocoindex library."""
-    _engine.init(settings.__dict__)
+    _engine.init(dump_engine_object(settings))
 
 @dataclass
 class ServerSettings:

cocoindex/llm.py

@@ -6,6 +6,7 @@ class LlmApiType(Enum):
     OPENAI = "OpenAi"
     OLLAMA = "Ollama"
     GEMINI = "Gemini"
+    ANTHROPIC = "Anthropic"
 
 @dataclass
 class LlmSpec:

cocoindex/op.py

@@ -9,7 +9,7 @@ from typing import get_type_hints, Protocol, Any, Callable, Awaitable, dataclass
 from enum import Enum
 
 from .typing import encode_enriched_type
-from .convert import to_engine_value, make_engine_value_converter
+from .convert import encode_engine_value, make_engine_value_decoder
 from . import _engine
 
 class OpCategory(Enum):
@@ -129,7 +129,7 @@ def _register_op_factory(
                     raise ValueError(
                         f"Too many positional arguments passed in: {len(args)} > {next_param_idx}")
                 self._args_converters.append(
-                    make_engine_value_converter(
+                    make_engine_value_decoder(
                         [arg_name], arg.value_type['type'], arg_param.annotation))
                 if arg_param.kind != inspect.Parameter.VAR_POSITIONAL:
                     next_param_idx += 1
@@ -146,7 +146,7 @@ def _register_op_factory(
                 if expected_arg is None:
                     raise ValueError(f"Unexpected keyword argument passed in: {kwarg_name}")
                 arg_param = expected_arg[1]
-                self._kwargs_converters[kwarg_name] = make_engine_value_converter(
+                self._kwargs_converters[kwarg_name] = make_engine_value_decoder(
                     [kwarg_name], kwarg.value_type['type'], arg_param.annotation)
 
             missing_args = [name for (name, arg) in expected_kwargs
@@ -188,7 +188,7 @@ def _register_op_factory(
                     output = await self._acall(*converted_args, **converted_kwargs)
             else:
                 output = await self._acall(*converted_args, **converted_kwargs)
-            return to_engine_value(output)
+            return encode_engine_value(output)
 
     _WrappedClass.__name__ = executor_cls.__name__
     _WrappedClass.__doc__ = executor_cls.__doc__

cocoindex/query.py

@@ -66,15 +66,16 @@ class SimpleSemanticsQueryHandler:
         return self._lazy_query_handler()
 
     def search(self, query: str, limit: int, vector_field_name: str | None = None,
-               similarity_matric: index.VectorSimilarityMetric | None = None) -> tuple[list[QueryResult], SimpleSemanticsQueryInfo]:
+               similarity_metric: index.VectorSimilarityMetric | None = None
+               ) -> tuple[list[QueryResult], SimpleSemanticsQueryInfo]:
         """
         Search the index with the given query, limit, vector field name, and similarity metric.
         """
         internal_results, internal_info = self.internal_handler().search(
             query, limit, vector_field_name,
-            similarity_matric.value if similarity_matric is not None else None)
-        fields = [field['name'] for field in internal_results['fields']]
-        results = [QueryResult(data=dict(zip(fields, result['data'])),  score=result['score']) for result in internal_results['results']]
+            similarity_metric.value if similarity_metric is not None else None)
+        results = [QueryResult(data=result['data'],  score=result['score'])
+                   for result in internal_results]
         info = SimpleSemanticsQueryInfo(
             similarity_metric=index.VectorSimilarityMetric(internal_info['similarity_metric']),
             query_vector=internal_info['query_vector'],

cocoindex/storages.py

@@ -9,7 +9,7 @@ from .auth_registry import AuthEntryReference
 class Postgres(op.StorageSpec):
     """Storage powered by Postgres and pgvector."""
 
-    database_url: str | None = None
+    database: AuthEntryReference | None = None
     table_name: str | None = None
 
 @dataclass

cocoindex/tests/test_convert.py

@@ -1,12 +1,11 @@
-import dataclasses
 import uuid
 import datetime
 from dataclasses import dataclass, make_dataclass
 import pytest
+import cocoindex
 from cocoindex.typing import encode_enriched_type
-from cocoindex.convert import to_engine_value
-from cocoindex.convert import make_engine_value_converter
-
+from cocoindex.convert import encode_engine_value, make_engine_value_decoder
+from typing import Literal
 @dataclass
 class Order:
     order_id: str
@@ -26,7 +25,7 @@ class Basket:
 class Customer:
     name: str
     order: Order
-    tags: list[Tag] = None
+    tags: list[Tag] | None = None
 
 @dataclass
 class NestedStruct:
@@ -34,63 +33,63 @@ class NestedStruct:
     orders: list[Order]
     count: int = 0
 
-def build_engine_value_converter(engine_type_in_py, python_type=None):
+def build_engine_value_decoder(engine_type_in_py, python_type=None):
     """
     Helper to build a converter for the given engine-side type (as represented in Python).
     If python_type is not specified, uses engine_type_in_py as the target.
     """
     engine_type = encode_enriched_type(engine_type_in_py)["type"]
-    return make_engine_value_converter([], engine_type, python_type or engine_type_in_py)
+    return make_engine_value_decoder([], engine_type, python_type or engine_type_in_py)
 
-def test_to_engine_value_basic_types():
-    assert to_engine_value(123) == 123
-    assert to_engine_value(3.14) == 3.14
-    assert to_engine_value("hello") == "hello"
-    assert to_engine_value(True) is True
+def test_encode_engine_value_basic_types():
+    assert encode_engine_value(123) == 123
+    assert encode_engine_value(3.14) == 3.14
+    assert encode_engine_value("hello") == "hello"
+    assert encode_engine_value(True) is True
 
-def test_to_engine_value_uuid():
+def test_encode_engine_value_uuid():
     u = uuid.uuid4()
-    assert to_engine_value(u) == u.bytes
+    assert encode_engine_value(u) == u.bytes
 
-def test_to_engine_value_date_time_types():
+def test_encode_engine_value_date_time_types():
     d = datetime.date(2024, 1, 1)
-    assert to_engine_value(d) == d
+    assert encode_engine_value(d) == d
     t = datetime.time(12, 30)
-    assert to_engine_value(t) == t
+    assert encode_engine_value(t) == t
     dt = datetime.datetime(2024, 1, 1, 12, 30)
-    assert to_engine_value(dt) == dt
+    assert encode_engine_value(dt) == dt
 
-def test_to_engine_value_struct():
+def test_encode_engine_value_struct():
     order = Order(order_id="O123", name="mixed nuts", price=25.0)
-    assert to_engine_value(order) == ["O123", "mixed nuts", 25.0, "default_extra"]
+    assert encode_engine_value(order) == ["O123", "mixed nuts", 25.0, "default_extra"]
 
-def test_to_engine_value_list_of_structs():
+def test_encode_engine_value_list_of_structs():
     orders = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]
-    assert to_engine_value(orders) == [["O1", "item1", 10.0, "default_extra"], ["O2", "item2", 20.0, "default_extra"]]
+    assert encode_engine_value(orders) == [["O1", "item1", 10.0, "default_extra"], ["O2", "item2", 20.0, "default_extra"]]
 
-def test_to_engine_value_struct_with_list():
+def test_encode_engine_value_struct_with_list():
     basket = Basket(items=["apple", "banana"])
-    assert to_engine_value(basket) == [["apple", "banana"]]
+    assert encode_engine_value(basket) == [["apple", "banana"]]
 
-def test_to_engine_value_nested_struct():
+def test_encode_engine_value_nested_struct():
     customer = Customer(name="Alice", order=Order("O1", "item1", 10.0))
-    assert to_engine_value(customer) == ["Alice", ["O1", "item1", 10.0, "default_extra"], None]
+    assert encode_engine_value(customer) == ["Alice", ["O1", "item1", 10.0, "default_extra"], None]
 
-def test_to_engine_value_empty_list():
-    assert to_engine_value([]) == []
-    assert to_engine_value([[]]) == [[]]
+def test_encode_engine_value_empty_list():
+    assert encode_engine_value([]) == []
+    assert encode_engine_value([[]]) == [[]]
 
-def test_to_engine_value_tuple():
-    assert to_engine_value(()) == []
-    assert to_engine_value((1, 2, 3)) == [1, 2, 3]
-    assert to_engine_value(((1, 2), (3, 4))) == [[1, 2], [3, 4]]
-    assert to_engine_value(([],)) == [[]]
-    assert to_engine_value(((),)) == [[]]
+def test_encode_engine_value_tuple():
+    assert encode_engine_value(()) == []
+    assert encode_engine_value((1, 2, 3)) == [1, 2, 3]
+    assert encode_engine_value(((1, 2), (3, 4))) == [[1, 2], [3, 4]]
+    assert encode_engine_value(([],)) == [[]]
+    assert encode_engine_value(((),)) == [[]]
 
-def test_to_engine_value_none():
-    assert to_engine_value(None) is None
+def test_encode_engine_value_none():
+    assert encode_engine_value(None) is None
 
-def test_make_engine_value_converter_basic_types():
+def test_make_engine_value_decoder_basic_types():
     for engine_type_in_py, value in [
         (int, 42),
         (float, 3.14),
@@ -98,11 +97,11 @@ def test_make_engine_value_converter_basic_types():
         (bool, True),
         # (type(None), None),  # Removed unsupported NoneType
     ]:
-        converter = build_engine_value_converter(engine_type_in_py)
-        assert converter(value) == value
+        decoder = build_engine_value_decoder(engine_type_in_py)
+        assert decoder(value) == value
 
 @pytest.mark.parametrize(
-    "converter_type, engine_val, expected",
+    "data_type, engine_val, expected",
     [
         # All fields match
         (Order, ["O123", "mixed nuts", 25.0, "default_extra"], Order("O123", "mixed nuts", 25.0, "default_extra")),
@@ -120,30 +119,30 @@ def test_make_engine_value_converter_basic_types():
         (Customer, ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"]], "extra"], Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip")])),
     ]
 )
-def test_struct_conversion_cases(converter_type, engine_val, expected):
-    converter = build_engine_value_converter(converter_type)
-    assert converter(engine_val) == expected
+def test_struct_decoder_cases(data_type, engine_val, expected):
+    decoder = build_engine_value_decoder(data_type)
+    assert decoder(engine_val) == expected
 
-def test_make_engine_value_converter_collections():
+def test_make_engine_value_decoder_collections():
     # List of structs
-    converter = build_engine_value_converter(list[Order])
+    decoder = build_engine_value_decoder(list[Order])
     engine_val = [
         ["O1", "item1", 10.0, "default_extra"],
         ["O2", "item2", 20.0, "default_extra"]
     ]
-    assert converter(engine_val) == [Order("O1", "item1", 10.0, "default_extra"), Order("O2", "item2", 20.0, "default_extra")]
+    assert decoder(engine_val) == [Order("O1", "item1", 10.0, "default_extra"), Order("O2", "item2", 20.0, "default_extra")]
     # Struct with list field
-    converter = build_engine_value_converter(Customer)
+    decoder = build_engine_value_decoder(Customer)
     engine_val = ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"], ["premium"]]]
-    assert converter(engine_val) == Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip"), Tag("premium")])
+    assert decoder(engine_val) == Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip"), Tag("premium")])
     # Struct with struct field
-    converter = build_engine_value_converter(NestedStruct)
+    decoder = build_engine_value_decoder(NestedStruct)
     engine_val = [
         ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"]]],
         [["O1", "item1", 10.0, "default_extra"], ["O2", "item2", 20.0, "default_extra"]],
         2
     ]
-    assert converter(engine_val) == NestedStruct(
+    assert decoder(engine_val) == NestedStruct(
         Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip")]),
         [Order("O1", "item1", 10.0, "default_extra"), Order("O2", "item2", 20.0, "default_extra")],
         2
@@ -227,8 +226,55 @@ def make_python_order(fields, defaults=None):
 def test_field_position_cases(engine_fields, python_fields, python_defaults, engine_val, expected_python_val):
     EngineOrder = make_engine_order(engine_fields)
     PythonOrder = make_python_order(python_fields, python_defaults)
-    converter = build_engine_value_converter(EngineOrder, PythonOrder)
+    decoder = build_engine_value_decoder(EngineOrder, PythonOrder)
     # Map field names to expected values
     expected_dict = dict(zip([f[0] for f in python_fields], expected_python_val))
     # Instantiate using keyword arguments (order doesn't matter)
-    assert converter(engine_val) == PythonOrder(**expected_dict)
+    assert decoder(engine_val) == PythonOrder(**expected_dict)
+
+def test_roundtrip_ltable():
+    t = list[Order]
+    value = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]
+    encoded = encode_engine_value(value)
+    assert encoded == [["O1", "item1", 10.0, "default_extra"], ["O2", "item2", 20.0, "default_extra"]]
+    decoded = build_engine_value_decoder(t)(encoded)
+    assert decoded == value
+
+def test_roundtrip_ktable_str_key():
+    t = dict[str, Order]
+    value = {"K1": Order("O1", "item1", 10.0), "K2": Order("O2", "item2", 20.0)}
+    encoded = encode_engine_value(value)
+    assert encoded == [["K1", "O1", "item1", 10.0, "default_extra"], ["K2", "O2", "item2", 20.0, "default_extra"]]
+    decoded = build_engine_value_decoder(t)(encoded)
+    assert decoded == value
+
+def test_roundtrip_ktable_struct_key():
+    @dataclass(frozen=True)
+    class OrderKey:
+        shop_id: str
+        version: int
+
+    t = dict[OrderKey, Order]
+    value = {OrderKey("A", 3): Order("O1", "item1", 10.0), OrderKey("B", 4): Order("O2", "item2", 20.0)}
+    encoded = encode_engine_value(value)
+    assert encoded == [[["A", 3], "O1", "item1", 10.0, "default_extra"],
+                       [["B", 4], "O2", "item2", 20.0, "default_extra"]]
+    decoded = build_engine_value_decoder(t)(encoded)
+    assert decoded == value
+
+IntVectorType = cocoindex.Vector[int, Literal[5]]
+def test_vector_as_vector() -> None:
+    value: IntVectorType = [1, 2, 3, 4, 5]
+    encoded = encode_engine_value(value)
+    assert encoded == [1, 2, 3, 4, 5]
+    decoded = build_engine_value_decoder(IntVectorType)(encoded)
+    assert decoded == value
+
+ListIntType = list[int]
+def test_vector_as_list() -> None:
+    value: ListIntType = [1, 2, 3, 4, 5]
+    encoded = encode_engine_value(value)
+    assert encoded == [1, 2, 3, 4, 5]
+    decoded = build_engine_value_decoder(ListIntType)(encoded)
+    assert decoded == value
+

cocoindex/typing.py

@@ -5,9 +5,9 @@ import datetime
 import types
 import inspect
 import uuid
-from typing import Annotated, NamedTuple, Any, TypeVar, TYPE_CHECKING, overload
+from typing import Annotated, NamedTuple, Any, TypeVar, TYPE_CHECKING, overload, Sequence, Protocol, Generic, Literal
 
-class Vector(NamedTuple):
+class VectorInfo(NamedTuple):
     dim: int | None
 
 class TypeKind(NamedTuple):
@@ -21,7 +21,7 @@ class TypeAttr:
         self.key = key
         self.value = value
 
-Annotation = Vector | TypeKind | TypeAttr
+Annotation = TypeKind | TypeAttr | VectorInfo
 
 Float32 = Annotated[float, TypeKind('Float32')]
 Float64 = Annotated[float, TypeKind('Float64')]
@@ -30,29 +30,34 @@ Json = Annotated[Any, TypeKind('Json')]
 LocalDateTime = Annotated[datetime.datetime, TypeKind('LocalDateTime')]
 OffsetDateTime = Annotated[datetime.datetime, TypeKind('OffsetDateTime')]
 
-COLLECTION_TYPES = ('Table', 'List')
-
-R = TypeVar("R")
-
 if TYPE_CHECKING:
-    Table = Annotated[list[R], TypeKind('Table')]
-    List = Annotated[list[R], TypeKind('List')]
+    T_co = TypeVar('T_co', covariant=True)
+    Dim_co = TypeVar('Dim_co', bound=int, covariant=True)
+
+    class Vector(Sequence[T_co], Generic[T_co, Dim_co], Protocol):
+        """Vector[T, Dim] is a special typing alias for a list[T] with optional dimension info"""
 else:
-    # pylint: disable=too-few-public-methods
-    class Table:  # type: ignore[unreachable]
-        """
-        A Table type, which has a list of rows. The first field of each row is the key.
-        """
-        def __class_getitem__(cls, item: type[R]):
-            return Annotated[list[item], TypeKind('Table')]
-
-    # pylint: disable=too-few-public-methods
-    class List:  # type: ignore[unreachable]
-        """
-        A List type, which has a list of ordered rows.
-        """
-        def __class_getitem__(cls, item: type[R]):
-            return Annotated[list[item], TypeKind('List')]
+    class Vector:  # type: ignore[unreachable]
+        """ A special typing alias for a list[T] with optional dimension info """
+        def __class_getitem__(self, params):
+            if not isinstance(params, tuple):
+                # Only element type provided
+                elem_type = params
+                return Annotated[list[elem_type], VectorInfo(dim=None)]
+            else:
+                # Element type and dimension provided
+                elem_type, dim = params
+                if typing.get_origin(dim) is Literal:
+                    dim = typing.get_args(dim)[0]  # Extract the literal value
+                return Annotated[list[elem_type], VectorInfo(dim=dim)]
+
+TABLE_TYPES = ('KTable', 'LTable')
+KEY_FIELD_NAME = '_key'
+
+ElementType = type | tuple[type, type]
+
+def _is_struct_type(t) -> bool:
+    return isinstance(t, type) and dataclasses.is_dataclass(t)
 
 @dataclasses.dataclass
 class AnalyzedTypeInfo:
@@ -60,9 +65,12 @@ class AnalyzedTypeInfo:
     Analyzed info of a Python type.
     """
     kind: str
-    vector_info: Vector | None
-    elem_type: type | None
-    dataclass_type: type | None
+    vector_info: VectorInfo | None  # For Vector
+    elem_type: ElementType | None   # For Vector and Table
+
+    key_type: type | None           # For element of KTable
+    dataclass_type: type | None     # For Struct
+
     attrs: dict[str, Any] | None
     nullable: bool = False
 
@@ -70,6 +78,12 @@ def analyze_type_info(t) -> AnalyzedTypeInfo:
     """
     Analyze a Python type and return the analyzed info.
     """
+    if isinstance(t, tuple) and len(t) == 2:
+        key_type, value_type = t
+        result = analyze_type_info(value_type)
+        result.key_type = key_type
+        return result
+
     annotations: tuple[Annotation, ...] = ()
     base_type = None
     nullable = False
@@ -98,33 +112,41 @@ def analyze_type_info(t) -> AnalyzedTypeInfo:
             if attrs is None:
                 attrs = dict()
             attrs[attr.key] = attr.value
-        elif isinstance(attr, Vector):
+        elif isinstance(attr, VectorInfo):
             vector_info = attr
         elif isinstance(attr, TypeKind):
             kind = attr.kind
 
     dataclass_type = None
     elem_type = None
-    if isinstance(t, type) and dataclasses.is_dataclass(t):
+    key_type = None
+    if _is_struct_type(t):
         if kind is None:
             kind = 'Struct'
         elif kind != 'Struct':
             raise ValueError(f"Unexpected type kind for struct: {kind}")
         dataclass_type = t
     elif base_type is collections.abc.Sequence or base_type is list:
+        args = typing.get_args(t)
+        elem_type = args[0]
+
         if kind is None:
-            kind = 'Vector' if vector_info is not None else 'List'
-        elif not (kind == 'Vector' or kind in COLLECTION_TYPES):
+            if _is_struct_type(elem_type):
+                kind = 'LTable'
+                if vector_info is not None:
+                    raise ValueError("Vector element must be a simple type, not a struct")
+            else:
+                kind = 'Vector'
+                if vector_info is None:
+                    vector_info = VectorInfo(dim=None)
+        elif not (kind == 'Vector' or kind in TABLE_TYPES):
             raise ValueError(f"Unexpected type kind for list: {kind}")
-
+    elif base_type is collections.abc.Mapping or base_type is dict:
         args = typing.get_args(t)
-        if len(args) != 1:
-            raise ValueError(f"{kind} must have exactly one type argument")
-        elem_type = args[0]
+        elem_type = (args[0], args[1])
+        kind = 'KTable'
     elif kind is None:
-        if base_type is collections.abc.Sequence or base_type is list:
-            kind = 'Vector' if vector_info is not None else 'List'
-        elif t is bytes:
+        if t is bytes:
             kind = 'Bytes'
         elif t is str:
             kind = 'Str'
@@ -145,20 +167,26 @@ def analyze_type_info(t) -> AnalyzedTypeInfo:
         else:
             raise ValueError(f"type unsupported yet: {t}")
 
-    return AnalyzedTypeInfo(kind=kind, vector_info=vector_info, elem_type=elem_type,
-                            dataclass_type=dataclass_type, attrs=attrs, nullable=nullable)
+    return AnalyzedTypeInfo(kind=kind, vector_info=vector_info,
+                            elem_type=elem_type, key_type=key_type, dataclass_type=dataclass_type,
+                            attrs=attrs, nullable=nullable)
 
-def _encode_fields_schema(dataclass_type: type) -> list[dict[str, Any]]:
+def _encode_fields_schema(dataclass_type: type, key_type: type | None = None) -> list[dict[str, Any]]:
     result = []
-    for field in dataclasses.fields(dataclass_type):
+    def add_field(name: str, t) -> None:
         try:
-            type_info = encode_enriched_type_info(analyze_type_info(field.type))
+            type_info = encode_enriched_type_info(analyze_type_info(t))
         except ValueError as e:
             e.add_note(f"Failed to encode annotation for field - "
-                       f"{dataclass_type.__name__}.{field.name}: {field.type}")
+                       f"{dataclass_type.__name__}.{name}: {t}")
             raise
-        type_info['name'] = field.name
+        type_info['name'] = name
         result.append(type_info)
+
+    if key_type is not None:
+        add_field(KEY_FIELD_NAME, key_type)
+    for field in dataclasses.fields(dataclass_type):
+        add_field(field.name, field.type)
     return result
 
 def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
@@ -167,7 +195,7 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
     if type_info.kind == 'Struct':
         if type_info.dataclass_type is None:
             raise ValueError("Struct type must have a dataclass type")
-        encoded_type['fields'] = _encode_fields_schema(type_info.dataclass_type)
+        encoded_type['fields'] = _encode_fields_schema(type_info.dataclass_type, type_info.key_type)
         if doc := inspect.getdoc(type_info.dataclass_type):
             encoded_type['description'] = doc
 
@@ -179,7 +207,7 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         encoded_type['element_type'] = _encode_type(analyze_type_info(type_info.elem_type))
         encoded_type['dimension'] = type_info.vector_info.dim
 
-    elif type_info.kind in COLLECTION_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")
         row_type_info = analyze_type_info(type_info.elem_type)

{cocoindex-0.1.26.dist-info → cocoindex-0.1.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.26
+Version: 0.1.27
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: pytest ; extra == 'test'

cocoindex-0.1.27.dist-info/RECORD

@@ -0,0 +1,24 @@
+cocoindex-0.1.27.dist-info/METADATA,sha256=upAIi8VNoLnqN1xUoKPpw1Hqspc3Rpb4cxKB3ahbFTs,8079
+cocoindex-0.1.27.dist-info/WHEEL,sha256=e_pkmRfvNwfhA3ReKBun4dlOSPa_df-kVPas-U1KpNY,104
+cocoindex-0.1.27.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex/functions.py,sha256=clnpHCYSsjUnc8Spbc1-5sQedG-60fmibodv9LpHgqo,1647
+cocoindex/query.py,sha256=8_3Lb_EVjZtl2ZyJNZGX16LoKXEd-PL8OjY-zs9GQeA,3205
+cocoindex/index.py,sha256=LssEOuZi6AqhwKtZM3QFeQpa9T-0ELi8G5DsrYKECvc,534
+cocoindex/lib.py,sha256=c6D7NuuTJj20WgVhnp0QGyK18lKMUvoDCiFr3PFs71s,3871
+cocoindex/auth_registry.py,sha256=lZ2rD5_9aC_UpGk7t4TmSYal_rjN7eHgO4_sU7FR0Zw,620
+cocoindex/convert.py,sha256=mBUTa_Ag39_ut-yE_jc1wqS3zLjtOm6QKet-bqJ-RWc,5947
+cocoindex/tests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+cocoindex/tests/test_convert.py,sha256=WPRKp0jv_uSEM81RGWEAmsax-J-FtXt90mZ0yEnvGLs,11236
+cocoindex/__init__.py,sha256=8atBT1HjclUOeiXd7TSzZWaqOR4x_qr5epvCKB7Z7oY,661
+cocoindex/flow.py,sha256=zuwvByhQxg0fMYshWCcq0YYe6X2TvAJJie4xetnNIPE,21054
+cocoindex/llm.py,sha256=_3rtahuKcqcEHPkFSwhXOSrekZyGxVApPoYtlU_chcA,348
+cocoindex/runtime.py,sha256=jqRnWkkIlAhE04gi4y0Y5bzuq9FX4j0aVNU-nengLJk,980
+cocoindex/op.py,sha256=ICCKZw6peCFu-CtMeIEaz6vlBxrf5dZwgUs9R4ALYNU,10604
+cocoindex/sources.py,sha256=wZFU8lwSXjyofJR-syySH9fTyPnBlAPJ6-1hQNX8fGA,936
+cocoindex/setup.py,sha256=W1HshwYk_K2aeLOVn_e62ZOXBO9yWsoUboRiH4SjF48,496
+cocoindex/cli.py,sha256=MvEUbQVrJy-sYbGQNsqIaMJvcQXQn1OQVNues22Hph0,7061
+cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cocoindex/typing.py,sha256=p3FEUnoQc6zmiG8YwO4T155sgZtyc_1AufiJe3bNol8,8458
+cocoindex/storages.py,sha256=bWsV4rRs4PJp6GemOb_PwoRu5XRgOsG-4sRMziztliU,1957
+cocoindex/_engine.cpython-312-darwin.so,sha256=JBHj9qaowOY_yuYszvJ7-DNf7zK88xRFRcz2Dt4jRds,59514816
+cocoindex-0.1.27.dist-info/RECORD,,

cocoindex-0.1.26.dist-info/RECORD

@@ -1,24 +0,0 @@
-cocoindex-0.1.26.dist-info/METADATA,sha256=BwpB45xAUpYFZugav3l9u1uJ_0s4Kdrnp-Y2FnSJZ0A,8079
-cocoindex-0.1.26.dist-info/WHEEL,sha256=e_pkmRfvNwfhA3ReKBun4dlOSPa_df-kVPas-U1KpNY,104
-cocoindex-0.1.26.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-cocoindex/functions.py,sha256=xcAeRQTy9JObfxpjyMn-dPY2y7XhVWjB7759xVyup6o,1657
-cocoindex/query.py,sha256=XsVY5cBJJ3a70qazkcCHjWZLE1zBqzMQ4HVSulicGMA,3273
-cocoindex/index.py,sha256=LssEOuZi6AqhwKtZM3QFeQpa9T-0ELi8G5DsrYKECvc,534
-cocoindex/lib.py,sha256=48nfWSg5IMzTSkVxdrWF8d9Hi-Bw8in_2rs7rQRwAs8,3505
-cocoindex/auth_registry.py,sha256=lZ2rD5_9aC_UpGk7t4TmSYal_rjN7eHgO4_sU7FR0Zw,620
-cocoindex/convert.py,sha256=tzlHadc-SaZCRBWxZEp08T4clJQPab_eXt6mUub0iQQ,5017
-cocoindex/tests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-cocoindex/tests/test_convert.py,sha256=mT_7HhXRu1PzGFbdmPmrq3O5-bPyDDG7mbr-mHaa2i0,9339
-cocoindex/__init__.py,sha256=f4LTPg4db7Wm3QO9HirvhsT11OVykiFxGbt1JK6taFA,572
-cocoindex/flow.py,sha256=1Mx-rYBzPIlFLNsiNVGhPtJKy2u6stZT1xjyPzERbFI,21068
-cocoindex/llm.py,sha256=4b20wpSHcgfDM7tdxRm1KIo_7C30nT7h0gCsWvs686I,320
-cocoindex/runtime.py,sha256=jqRnWkkIlAhE04gi4y0Y5bzuq9FX4j0aVNU-nengLJk,980
-cocoindex/op.py,sha256=zOQzgnVDvET8LtUt7TW8PfHMhaA0eAXoZ3c_EsL6jtU,10602
-cocoindex/sources.py,sha256=wZFU8lwSXjyofJR-syySH9fTyPnBlAPJ6-1hQNX8fGA,936
-cocoindex/setup.py,sha256=W1HshwYk_K2aeLOVn_e62ZOXBO9yWsoUboRiH4SjF48,496
-cocoindex/cli.py,sha256=MvEUbQVrJy-sYbGQNsqIaMJvcQXQn1OQVNues22Hph0,7061
-cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/typing.py,sha256=4mP9VXS75s3VMfF1LDc1LXsBng7uGdR5aD73N8iaeSM,7282
-cocoindex/storages.py,sha256=GRHkmwuSAU7neF3H0pjAPfeEkmtXv-DJc7CKYjcATvE,1946
-cocoindex/_engine.cpython-312-darwin.so,sha256=nVnQv2Ok4mUXYFOvKEdd-YlWdzF-YGwDU-DACxkjASQ,59217776
-cocoindex-0.1.26.dist-info/RECORD,,