cocoindex 0.2.10__cp311-abi3-macosx_11_0_arm64.whl → 0.2.12__cp311-abi3-macosx_11_0_arm64.whl

cocoindex/op.py

@@ -2,7 +2,6 @@
 Facilities for defining cocoindex operations.
 """
 
-import asyncio
 import dataclasses
 import inspect
 from enum import Enum
@@ -19,6 +18,8 @@ from typing import (
 from . import _engine  # type: ignore
 from .subprocess_exec import executor_stub
 from .convert import (
+    dump_engine_object,
+    load_engine_object,
     make_engine_value_encoder,
     make_engine_value_decoder,
     make_engine_key_decoder,
@@ -31,7 +32,12 @@ from .typing import (
     analyze_type_info,
     AnalyzedAnyType,
     AnalyzedDictType,
+    EnrichedValueType,
+    decode_engine_field_schemas,
+    FieldSchema,
 )
+from .runtime import to_async_call
+from .index import IndexOptions
 
 
 class OpCategory(Enum):
@@ -86,15 +92,6 @@ class Executor(Protocol):
     op_category: OpCategory
 
 
-def _load_spec_from_engine(
-    spec_loader: Callable[..., Any], spec: dict[str, Any]
-) -> Any:
-    """
-    Load a spec from the engine.
-    """
-    return spec_loader(**spec)
-
-
 def _get_required_method(cls: type, name: str) -> Callable[..., Any]:
     method = getattr(cls, name, None)
     if method is None:
@@ -105,7 +102,7 @@ def _get_required_method(cls: type, name: str) -> Callable[..., Any]:
 
 
 class _EngineFunctionExecutorFactory:
-    _spec_loader: Callable[..., Any]
+    _spec_loader: Callable[[Any], Any]
     _executor_cls: type
 
     def __init__(self, spec_loader: Callable[..., Any], executor_cls: type):
@@ -113,9 +110,9 @@ class _EngineFunctionExecutorFactory:
         self._executor_cls = executor_cls
 
     def __call__(
-        self, spec: dict[str, Any], *args: Any, **kwargs: Any
+        self, raw_spec: dict[str, Any], *args: Any, **kwargs: Any
     ) -> tuple[dict[str, Any], Executor]:
-        spec = _load_spec_from_engine(self._spec_loader, spec)
+        spec = self._spec_loader(raw_spec)
         executor = self._executor_cls(spec)
         result_type = executor.analyze_schema(*args, **kwargs)
         return (result_type, executor)
@@ -150,12 +147,6 @@ class OpArgs:
     arg_relationship: tuple[ArgRelationship, str] | None = None
 
 
-def _to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
-    if inspect.iscoroutinefunction(call):
-        return call
-    return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))
-
-
 @dataclasses.dataclass
 class _ArgInfo:
     decoder: Callable[[Any], Any]
@@ -218,8 +209,9 @@ def _register_op_factory(
                             TypeAttr(related_attr.value, actual_arg.analyzed_value)
                         )
                 type_info = analyze_type_info(arg_param.annotation)
+                enriched = EnrichedValueType.decode(actual_arg.value_type)
                 decoder = make_engine_value_decoder(
-                    [arg_name], actual_arg.value_type["type"], type_info
+                    [arg_name], enriched.type, type_info
                 )
                 is_required = not type_info.nullable
                 if is_required and actual_arg.value_type.get("nullable", False):
@@ -319,8 +311,8 @@ def _register_op_factory(
             """
             prepare_method = getattr(self._executor, "prepare", None)
             if prepare_method is not None:
-                await _to_async_call(prepare_method)()
-            self._acall = _to_async_call(self._executor.__call__)
+                await to_async_call(prepare_method)()
+            self._acall = to_async_call(self._executor.__call__)
 
         async def __call__(self, *args: Any, **kwargs: Any) -> Any:
             decoded_args = []
@@ -379,7 +371,7 @@ def executor_class(**args: Any) -> Callable[[type], type]:
             expected_args=list(sig.parameters.items())[1:],  # First argument is `self`
             expected_return=sig.return_annotation,
             executor_factory=cls,
-            spec_loader=spec_cls,
+            spec_loader=lambda v: load_engine_object(spec_cls, v),
             op_kind=spec_cls.__name__,
             op_args=op_args,
         )
@@ -415,7 +407,7 @@ def function(**args: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
             expected_args=list(sig.parameters.items()),
             expected_return=sig.return_annotation,
             executor_factory=_SimpleFunctionExecutor,
-            spec_loader=lambda: fn,
+            spec_loader=lambda _: fn,
             op_kind=op_kind,
             op_args=op_args,
         )
@@ -435,8 +427,44 @@ class _TargetConnectorContext:
     target_name: str
     spec: Any
     prepared_spec: Any
+    key_fields_schema: list[FieldSchema]
     key_decoder: Callable[[Any], Any]
+    value_fields_schema: list[FieldSchema]
     value_decoder: Callable[[Any], Any]
+    index_options: IndexOptions
+    setup_state: Any
+
+
+def _build_args(
+    method: Callable[..., Any], num_required_args: int, **kwargs: Any
+) -> list[Any]:
+    signature = inspect.signature(method)
+    for param in signature.parameters.values():
+        if param.kind not in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        ):
+            raise ValueError(
+                f"Method {method.__name__} should only have positional arguments, got {param.kind.name}"
+            )
+    if len(signature.parameters) < num_required_args:
+        raise ValueError(
+            f"Method {method.__name__} must have at least {num_required_args} required arguments: "
+            f"{', '.join(list(kwargs.keys())[:num_required_args])}"
+        )
+    if len(kwargs) > len(kwargs):
+        raise ValueError(
+            f"Method {method.__name__} can only have at most {num_required_args} arguments: {', '.join(kwargs.keys())}"
+        )
+    return [v for _, v in zip(signature.parameters, kwargs.values())]
+
+
+class TargetStateCompatibility(Enum):
+    """The compatibility of the target state."""
+
+    COMPATIBLE = "Compatible"
+    PARTIALLY_COMPATIBLE = "PartialCompatible"
+    NOT_COMPATIBLE = "NotCompatible"
 
 
 class _TargetConnector:
@@ -444,8 +472,10 @@ class _TargetConnector:
     The connector class passed to the engine.
     """
 
-    _spec_cls: type
-    _connector_cls: type
+    _spec_cls: type[Any]
+    _persistent_key_type: Any
+    _setup_state_cls: type[Any]
+    _connector_cls: type[Any]
 
     _get_persistent_key_fn: Callable[[_TargetConnectorContext, str], Any]
     _apply_setup_change_async_fn: Callable[
@@ -454,19 +484,27 @@ class _TargetConnector:
     _mutate_async_fn: Callable[..., Awaitable[None]]
     _mutatation_type: AnalyzedDictType | None
 
-    def __init__(self, spec_cls: type, connector_cls: type):
+    def __init__(
+        self,
+        spec_cls: type[Any],
+        persistent_key_type: Any,
+        setup_state_cls: type[Any],
+        connector_cls: type[Any],
+    ):
         self._spec_cls = spec_cls
+        self._persistent_key_type = persistent_key_type
+        self._setup_state_cls = setup_state_cls
         self._connector_cls = connector_cls
 
         self._get_persistent_key_fn = _get_required_method(
             connector_cls, "get_persistent_key"
         )
-        self._apply_setup_change_async_fn = _to_async_call(
+        self._apply_setup_change_async_fn = to_async_call(
             _get_required_method(connector_cls, "apply_setup_change")
         )
 
         mutate_fn = _get_required_method(connector_cls, "mutate")
-        self._mutate_async_fn = _to_async_call(mutate_fn)
+        self._mutate_async_fn = to_async_call(mutate_fn)
 
         # Store the type annotation for later use
         self._mutatation_type = self._analyze_mutate_mutation_type(
@@ -513,15 +551,16 @@ class _TargetConnector:
         raise ValueError(
             f"Method {connector_cls.__name__}.mutate(*args) parameter must be a tuple with "
             f"2 elements (tuple[SpecType, dict[str, ValueStruct]], spec and mutation in dict), "
-            "got {args_type}"
+            f"got {analyzed_args_type.core_type}"
         )
 
     def create_export_context(
         self,
         name: str,
-        spec: dict[str, Any],
-        key_fields_schema: list[Any],
-        value_fields_schema: list[Any],
+        raw_spec: dict[str, Any],
+        raw_key_fields_schema: list[Any],
+        raw_value_fields_schema: list[Any],
+        raw_index_options: dict[str, Any],
     ) -> _TargetConnectorContext:
         key_annotation, value_annotation = (
             (
@@ -532,34 +571,103 @@ class _TargetConnector:
             else (Any, Any)
         )
 
+        key_fields_schema = decode_engine_field_schemas(raw_key_fields_schema)
         key_decoder = make_engine_key_decoder(
-            ["(key)"], key_fields_schema, analyze_type_info(key_annotation)
+            ["<key>"], key_fields_schema, analyze_type_info(key_annotation)
         )
+        value_fields_schema = decode_engine_field_schemas(raw_value_fields_schema)
         value_decoder = make_engine_struct_decoder(
-            ["(value)"], value_fields_schema, analyze_type_info(value_annotation)
+            ["<value>"], value_fields_schema, analyze_type_info(value_annotation)
         )
 
-        loaded_spec = _load_spec_from_engine(self._spec_cls, spec)
-        prepare_method = getattr(self._connector_cls, "prepare", None)
-        if prepare_method is None:
-            prepared_spec = loaded_spec
-        else:
-            prepared_spec = prepare_method(loaded_spec)
-
+        spec = load_engine_object(self._spec_cls, raw_spec)
+        index_options = load_engine_object(IndexOptions, raw_index_options)
         return _TargetConnectorContext(
             target_name=name,
-            spec=loaded_spec,
-            prepared_spec=prepared_spec,
+            spec=spec,
+            prepared_spec=None,
+            key_fields_schema=key_fields_schema,
             key_decoder=key_decoder,
+            value_fields_schema=value_fields_schema,
             value_decoder=value_decoder,
+            index_options=index_options,
+            setup_state=None,
         )
 
     def get_persistent_key(self, export_context: _TargetConnectorContext) -> Any:
-        return self._get_persistent_key_fn(
-            export_context.spec, export_context.target_name
+        args = _build_args(
+            self._get_persistent_key_fn,
+            1,
+            spec=export_context.spec,
+            target_name=export_context.target_name,
+        )
+        return dump_engine_object(self._get_persistent_key_fn(*args))
+
+    def get_setup_state(self, export_context: _TargetConnectorContext) -> Any:
+        get_setup_state_fn = getattr(self._connector_cls, "get_setup_state", None)
+        if get_setup_state_fn is None:
+            state = export_context.spec
+            if not isinstance(state, self._setup_state_cls):
+                raise ValueError(
+                    f"Expect a get_setup_state() method for {self._connector_cls} that returns an instance of {self._setup_state_cls}"
+                )
+        else:
+            args = _build_args(
+                get_setup_state_fn,
+                1,
+                spec=export_context.spec,
+                key_fields_schema=export_context.key_fields_schema,
+                value_fields_schema=export_context.value_fields_schema,
+                index_options=export_context.index_options,
+            )
+            state = get_setup_state_fn(*args)
+            if not isinstance(state, self._setup_state_cls):
+                raise ValueError(
+                    f"Method {get_setup_state_fn.__name__} must return an instance of {self._setup_state_cls}, got {type(state)}"
+                )
+        export_context.setup_state = state
+        return dump_engine_object(state)
+
+    def check_state_compatibility(
+        self, raw_desired_state: Any, raw_existing_state: Any
+    ) -> Any:
+        check_state_compatibility_fn = getattr(
+            self._connector_cls, "check_state_compatibility", None
         )
+        if check_state_compatibility_fn is not None:
+            compatibility = check_state_compatibility_fn(
+                load_engine_object(self._setup_state_cls, raw_desired_state),
+                load_engine_object(self._setup_state_cls, raw_existing_state),
+            )
+        else:
+            compatibility = (
+                TargetStateCompatibility.COMPATIBLE
+                if raw_desired_state == raw_existing_state
+                else TargetStateCompatibility.PARTIALLY_COMPATIBLE
+            )
+        return dump_engine_object(compatibility)
 
-    def describe_resource(self, key: Any) -> str:
+    async def prepare_async(
+        self,
+        export_context: _TargetConnectorContext,
+    ) -> None:
+        prepare_fn = getattr(self._connector_cls, "prepare", None)
+        if prepare_fn is None:
+            export_context.prepared_spec = export_context.spec
+            return
+        args = _build_args(
+            prepare_fn,
+            1,
+            spec=export_context.spec,
+            setup_state=export_context.setup_state,
+            key_fields_schema=export_context.key_fields_schema,
+            value_fields_schema=export_context.value_fields_schema,
+        )
+        async_prepare_fn = to_async_call(prepare_fn)
+        export_context.prepared_spec = await async_prepare_fn(*args)
+
+    def describe_resource(self, raw_key: Any) -> str:
+        key = load_engine_object(self._persistent_key_type, raw_key)
         describe_fn = getattr(self._connector_cls, "describe", None)
         if describe_fn is None:
             return str(key)
@@ -569,15 +677,16 @@ class _TargetConnector:
         self,
         changes: list[tuple[Any, list[dict[str, Any] | None], dict[str, Any] | None]],
     ) -> None:
-        for key, previous, current in changes:
+        for raw_key, previous, current in changes:
+            key = load_engine_object(self._persistent_key_type, raw_key)
             prev_specs = [
-                _load_spec_from_engine(self._spec_cls, spec)
+                load_engine_object(self._setup_state_cls, spec)
                 if spec is not None
                 else None
                 for spec in previous
             ]
             curr_spec = (
-                _load_spec_from_engine(self._spec_cls, current)
+                load_engine_object(self._setup_state_cls, current)
                 if current is not None
                 else None
             )
@@ -610,7 +719,12 @@ class _TargetConnector:
         )
 
 
-def target_connector(spec_cls: type) -> Callable[[type], type]:
+def target_connector(
+    *,
+    spec_cls: type[Any],
+    persistent_key_type: Any = Any,
+    setup_state_cls: type[Any] | None = None,
+) -> Callable[[type], type]:
     """
     Decorate a class to provide a target connector for an op.
     """
@@ -621,7 +735,9 @@ def target_connector(spec_cls: type) -> Callable[[type], type]:
 
     # Register the target connector.
     def _inner(connector_cls: type) -> type:
-        connector = _TargetConnector(spec_cls, connector_cls)
+        connector = _TargetConnector(
+            spec_cls, persistent_key_type, setup_state_cls or spec_cls, connector_cls
+        )
         _engine.register_target_connector(spec_cls.__name__, connector)
         return connector_cls
 

cocoindex/query_handler.py

@@ -0,0 +1,51 @@
+import dataclasses
+import numpy as np
+from numpy import typing as npt
+from typing import Generic, TypeVar
+from .index import VectorSimilarityMetric
+
+
+@dataclasses.dataclass
+class QueryHandlerResultFields:
+    """
+    Specify field names in query results returned by the query handler.
+    This provides metadata for tools like CocoInsight to recognize structure of the query results.
+    """
+
+    embedding: list[str] = dataclasses.field(default_factory=list)
+    score: str | None = None
+
+
+@dataclasses.dataclass
+class QueryHandlerInfo:
+    """
+    Info to configure a query handler.
+    """
+
+    result_fields: QueryHandlerResultFields | None = None
+
+
+@dataclasses.dataclass
+class QueryInfo:
+    """
+    Info about the query.
+    """
+
+    embedding: list[float] | npt.NDArray[np.float32] | None = None
+    similarity_metric: VectorSimilarityMetric | None = None
+
+
+R = TypeVar("R")
+
+
+@dataclasses.dataclass
+class QueryOutput(Generic[R]):
+    """
+    Output of a query handler.
+
+    results: list of results. Each result can be a dict or a dataclass.
+    query_info: Info about the query.
+    """
+
+    results: list[R]
+    query_info: QueryInfo = dataclasses.field(default_factory=QueryInfo)

cocoindex/runtime.py

@@ -5,7 +5,8 @@ manner.
 
 import threading
 import asyncio
-from typing import Any, Coroutine, TypeVar
+import inspect
+from typing import Any, Callable, Coroutine, TypeVar, Awaitable
 
 
 T = TypeVar("T")
@@ -35,3 +36,9 @@ class _ExecutionContext:
 
 
 execution_context = _ExecutionContext()
+
+
+def to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
+    if inspect.iscoroutinefunction(call):
+        return call
+    return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))

cocoindex/targets/__init__.py

@@ -0,0 +1,5 @@
+"""
+Targets supported by CocoIndex.
+"""
+
+from ._engine_builtin_specs import *

cocoindex/{targets.py → targets/_engine_builtin_specs.py}

@@ -3,10 +3,10 @@
 from dataclasses import dataclass
 from typing import Sequence
 
-from . import op
-from . import index
-from .auth_registry import AuthEntryReference
-from .setting import DatabaseConnectionSpec
+from .. import op
+from .. import index
+from ..auth_registry import AuthEntryReference
+from ..setting import DatabaseConnectionSpec
 
 
 class Postgres(op.TargetSpec):