cocoindex 0.1.38__cp311-cp311-win_amd64.whl → 0.1.39__cp311-cp311-win_amd64.whl

cocoindex/__init__.py

@@ -2,13 +2,14 @@
 Cocoindex is a framework for building and running indexing pipelines.
 """
 from . import functions, query, sources, storages, cli
-from .flow import FlowBuilder, DataScope, DataSlice, Flow, flow_def
+from .flow import FlowBuilder, DataScope, DataSlice, Flow, flow_def, transform_flow
 from .flow import EvaluateAndDumpOptions, GeneratedField
 from .flow import update_all_flows_async, FlowLiveUpdater, FlowLiveUpdaterOptions
 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 .setting import *
+from .setting import DatabaseConnectionSpec, Settings, ServerSettings
+from .setting import get_app_namespace
 from ._engine import OpArgSchema
 from .typing import Float32, Float64, LocalDateTime, OffsetDateTime, Range, Vector, Json

cocoindex/convert.py

@@ -44,7 +44,7 @@ def make_engine_value_decoder(
 
     src_type_kind = src_type['kind']
 
-    if dst_annotation is inspect.Parameter.empty:
+    if dst_annotation is None or dst_annotation is inspect.Parameter.empty or dst_annotation is Any:
         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.")

cocoindex/flow.py

@@ -8,8 +8,9 @@ import asyncio
 import re
 import inspect
 import datetime
+import functools
 
-from typing import Any, Callable, Sequence, TypeVar
+from typing import Any, Callable, Sequence, TypeVar, Generic, get_args, get_origin, Type, NamedTuple
 from threading import Lock
 from enum import Enum
 from dataclasses import dataclass
@@ -20,7 +21,7 @@ from . import _engine
 from . import index
 from . import op
 from . import setting
-from .convert import dump_engine_object
+from .convert import dump_engine_object, encode_engine_value, make_engine_value_decoder
 from .typing import encode_enriched_type
 from .runtime import execution_context
 
@@ -123,7 +124,7 @@ class _DataSliceState:
         # TODO: We'll support this by an identity transformer or "aliasing" in the future.
         raise ValueError("DataSlice is already attached to a field")
 
-class DataSlice:
+class DataSlice(Generic[T]):
     """A data slice represents a slice of data in a flow. It's readonly."""
 
     _state: _DataSliceState
@@ -183,11 +184,11 @@ class DataSlice:
                         name, prefix=_to_snake_case(_spec_kind(fn_spec))+'_'),
                 ))
 
-    def call(self, func: Callable[[DataSlice], T]) -> T:
+    def call(self, func: Callable[[DataSlice], T], *args, **kwargs) -> T:
         """
         Call a function with the data slice.
         """
-        return func(self)
+        return func(self, *args, **kwargs)
 
 def _data_slice_state(data_slice: DataSlice) -> _DataSliceState:
     return data_slice._state  # pylint: disable=protected-access
@@ -642,27 +643,67 @@ async def update_all_flows_async(options: FlowLiveUpdaterOptions) -> dict[str, _
     all_stats = await asyncio.gather(*(_update_flow(name, fl) for (name, fl) in fls.items()))
     return dict(all_stats)
 
-_transient_flow_name_builder = _NameBuilder()
-class TransientFlow:
+def _get_data_slice_annotation_type(data_slice_type: Type[DataSlice[T]]) -> Type[T] | None:
+    type_args = get_args(data_slice_type)
+    if data_slice_type is DataSlice:
+        return None
+    if get_origin(data_slice_type) != DataSlice or len(type_args) != 1:
+        raise ValueError(f"Expect a DataSlice[T] type, but got {data_slice_type}")
+    return type_args[0]
+
+_transform_flow_name_builder = _NameBuilder()
+
+class TransformFlowInfo(NamedTuple):
+    engine_flow: _engine.TransientFlow
+    result_decoder: Callable[[Any], T]
+
+class TransformFlow(Generic[T]):
     """
     A transient transformation flow that transforms in-memory data.
     """
-    _engine_flow: _engine.TransientFlow
+    _flow_fn: Callable[..., DataSlice[T]]
+    _flow_name: str
+    _flow_arg_types: list[Any]
+    _param_names: list[str]
+
+    _lazy_lock: asyncio.Lock
+    _lazy_flow_info: TransformFlowInfo | None = None
 
     def __init__(
-            self, flow_fn: Callable[..., DataSlice],
+            self, flow_fn: Callable[..., DataSlice[T]],
             flow_arg_types: Sequence[Any], /, name: str | None = None):
+        self._flow_fn = flow_fn
+        self._flow_name = _transform_flow_name_builder.build_name(name, prefix="_transform_flow_")
+        self._flow_arg_types = list(flow_arg_types)
+        self._lazy_lock = asyncio.Lock()
+
+    def __call__(self, *args, **kwargs) -> DataSlice[T]:
+        return self._flow_fn(*args, **kwargs)
 
-        flow_builder_state = _FlowBuilderState(
-            name=_transient_flow_name_builder.build_name(name, prefix="_transient_flow_"))
-        sig = inspect.signature(flow_fn)
-        if len(sig.parameters) != len(flow_arg_types):
+    @property
+    def _flow_info(self) -> TransformFlowInfo:
+        if self._lazy_flow_info is not None:
+            return self._lazy_flow_info
+        return execution_context.run(self._flow_info_async())
+
+    async def _flow_info_async(self) -> TransformFlowInfo:
+        if self._lazy_flow_info is not None:
+            return self._lazy_flow_info
+        async with self._lazy_lock:
+            if self._lazy_flow_info is None:
+                self._lazy_flow_info = await self._build_flow_info_async()
+            return self._lazy_flow_info
+
+    async def _build_flow_info_async(self) -> TransformFlowInfo:
+        flow_builder_state = _FlowBuilderState(name=self._flow_name)
+        sig = inspect.signature(self._flow_fn)
+        if len(sig.parameters) != len(self._flow_arg_types):
             raise ValueError(
                 f"Number of parameters in the flow function ({len(sig.parameters)}) "
-                "does not match the number of argument types ({len(flow_arg_types)})")
+                f"does not match the number of argument types ({len(self._flow_arg_types)})")
 
         kwargs: dict[str, DataSlice] = {}
-        for (param_name, param), param_type in zip(sig.parameters.items(), flow_arg_types):
+        for (param_name, param), param_type in zip(sig.parameters.items(), self._flow_arg_types):
             if param.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
                                   inspect.Parameter.KEYWORD_ONLY):
                 raise ValueError(f"Parameter {param_name} is not a parameter can be passed by name")
@@ -670,20 +711,68 @@ class TransientFlow:
                 param_name, encode_enriched_type(param_type))
             kwargs[param_name] = DataSlice(_DataSliceState(flow_builder_state, engine_ds))
 
-        output = flow_fn(**kwargs)
+        output = self._flow_fn(**kwargs)
         flow_builder_state.engine_flow_builder.set_direct_output(
             _data_slice_state(output).engine_data_slice)
-        self._engine_flow = flow_builder_state.engine_flow_builder.build_transient_flow(
-            execution_context.event_loop)
+        engine_flow = await flow_builder_state.engine_flow_builder.build_transient_flow_async(execution_context.event_loop)
+        self._param_names = list(sig.parameters.keys())
+
+        engine_return_type = _data_slice_state(output).engine_data_slice.data_type().schema()
+        python_return_type = _get_data_slice_annotation_type(sig.return_annotation)
+        result_decoder = make_engine_value_decoder([], engine_return_type['type'], python_return_type)
+
+        return TransformFlowInfo(engine_flow, result_decoder)
 
     def __str__(self):
-        return str(self._engine_flow)
+        return str(self._flow_info.engine_flow)
 
     def __repr__(self):
-        return repr(self._engine_flow)
+        return repr(self._flow_info.engine_flow)
 
     def internal_flow(self) -> _engine.TransientFlow:
         """
         Get the internal flow.
         """
-        return self._engine_flow
+        return self._flow_info.engine_flow
+
+    def eval(self, *args, **kwargs) -> T:
+        """
+        Evaluate the transform flow.
+        """
+        return execution_context.run(self.eval_async(*args, **kwargs))
+
+    async def eval_async(self, *args, **kwargs) -> T:
+        """
+        Evaluate the transform flow.
+        """
+        flow_info = await self._flow_info_async()
+        params = []
+        for i, arg in enumerate(self._param_names):
+            if i < len(args):
+                params.append(encode_engine_value(args[i]))
+            elif arg in kwargs:
+                params.append(encode_engine_value(kwargs[arg]))
+            else:
+                raise ValueError(f"Parameter {arg} is not provided")
+        engine_result = await flow_info.engine_flow.evaluate_async(params)
+        return flow_info.result_decoder(engine_result)
+
+
+def transform_flow() -> Callable[[Callable[..., DataSlice[T]]], TransformFlow[T]]:
+    """
+    A decorator to wrap the transform function.
+    """
+    def _transform_flow_wrapper(fn: Callable[..., DataSlice[T]]):
+        sig = inspect.signature(fn)
+        arg_types = []
+        for (param_name, param) in sig.parameters.items():
+            if param.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                                  inspect.Parameter.KEYWORD_ONLY):
+                raise ValueError(f"Parameter {param_name} is not a parameter can be passed by name")
+            arg_types.append(_get_data_slice_annotation_type(param.annotation))
+
+        _transform_flow = TransformFlow(fn, arg_types)
+        functools.update_wrapper(_transform_flow, fn)
+        return _transform_flow
+
+    return _transform_flow_wrapper

cocoindex/op.py

@@ -100,8 +100,8 @@ def _register_op_factory(
             return op_args.behavior_version
 
     class _WrappedClass(executor_cls, _Fallback):
-        _args_converters: list[Callable[[Any], Any]]
-        _kwargs_converters: dict[str, Callable[[str, Any], Any]]
+        _args_decoders: list[Callable[[Any], Any]]
+        _kwargs_decoders: dict[str, Callable[[str, Any], Any]]
         _acall: Callable
 
         def __init__(self, spec):
@@ -109,17 +109,17 @@ def _register_op_factory(
             self.spec = spec
             self._acall = _to_async_call(super().__call__)
 
-        def analyze(self, *args, **kwargs):
+        def analyze(self, *args: _engine.OpArgSchema, **kwargs: _engine.OpArgSchema):
             """
             Analyze the spec and arguments. In this phase, argument types should be validated.
             It should return the expected result type for the current op.
             """
-            self._args_converters = []
-            self._kwargs_converters = {}
+            self._args_decoders = []
+            self._kwargs_decoders = {}
 
             # Match arguments with parameters.
             next_param_idx = 0
-            for arg in  args:
+            for arg in args:
                 if next_param_idx >= len(expected_args):
                     raise ValueError(
                         f"Too many arguments passed in: {len(args)} > {len(expected_args)}")
@@ -128,7 +128,7 @@ def _register_op_factory(
                     inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.VAR_KEYWORD):
                     raise ValueError(
                         f"Too many positional arguments passed in: {len(args)} > {next_param_idx}")
-                self._args_converters.append(
+                self._args_decoders.append(
                     make_engine_value_decoder(
                         [arg_name], arg.value_type['type'], arg_param.annotation))
                 if arg_param.kind != inspect.Parameter.VAR_POSITIONAL:
@@ -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_decoder(
+                self._kwargs_decoders[kwarg_name] = make_engine_value_decoder(
                     [kwarg_name], kwarg.value_type['type'], arg_param.annotation)
 
             missing_args = [name for (name, arg) in expected_kwargs
@@ -174,8 +174,8 @@ def _register_op_factory(
                 await _to_async_call(setup_method)()
 
         async def __call__(self, *args, **kwargs):
-            converted_args = (converter(arg) for converter, arg in zip(self._args_converters, args))
-            converted_kwargs = {arg_name: self._kwargs_converters[arg_name](arg)
+            decoded_args = (decoder(arg) for decoder, arg in zip(self._args_decoders, args))
+            decoded_kwargs = {arg_name: self._kwargs_decoders[arg_name](arg)
                                 for arg_name, arg in kwargs.items()}
 
             if op_args.gpu:
@@ -185,9 +185,9 @@ def _register_op_factory(
                 # For now, we use a lock to ensure only one task is executed at a time.
                 # TODO: Implement multi-processing dispatching.
                 async with _gpu_dispatch_lock:
-                    output = await self._acall(*converted_args, **converted_kwargs)
+                    output = await self._acall(*decoded_args, **decoded_kwargs)
             else:
-                output = await self._acall(*converted_args, **converted_kwargs)
+                output = await self._acall(*decoded_args, **decoded_kwargs)
             return encode_engine_value(output)
 
     _WrappedClass.__name__ = executor_cls.__name__

cocoindex/query.py

@@ -50,7 +50,7 @@ class SimpleSemanticsQueryHandler:
                     if engine_handler is None:
                         engine_handler = _engine.SimpleSemanticsQueryHandler(
                             flow.internal_flow(), target_name,
-                            fl.TransientFlow(query_transform_flow, [str]).internal_flow(),
+                            fl.TransformFlow(query_transform_flow, [str]).internal_flow(),
                             default_similarity_metric.value)
                         engine_handler.register_query_handler(name)
             return engine_handler

cocoindex/setting.py

@@ -49,7 +49,7 @@ def _load_field(target: dict[str, Any], name: str, env_name: str, required: bool
 class Settings:
     """Settings for the cocoindex library."""
     database: DatabaseConnectionSpec
-    app_namespace: str
+    app_namespace: str = ""
 
     @classmethod
     def from_env(cls) -> Self:

{cocoindex-0.1.38.dist-info → cocoindex-0.1.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.38
+Version: 0.1.39
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0

{cocoindex-0.1.38.dist-info → cocoindex-0.1.39.dist-info}/RECORD

@@ -1,25 +1,25 @@
-cocoindex-0.1.38.dist-info/METADATA,sha256=F04y39dSBt_UKZkInX4seXLMjI-bPLW3WYVD7mZSsnQ,9958
-cocoindex-0.1.38.dist-info/WHEEL,sha256=Kw8y023UaufEXFM028WbBCk7rJUUAhAX1Zw-54pv0m4,96
-cocoindex-0.1.38.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-cocoindex/__init__.py,sha256=euvHrtzQEXtAEut0EFlgcc4VtogA6zltqWRdt70Zgrg,703
-cocoindex/_engine.cp311-win_amd64.pyd,sha256=OKJEyUBeM0IofIMrcK4Nza4Emfg6_KhJ47EM7TkI4fQ,60772352
+cocoindex-0.1.39.dist-info/METADATA,sha256=pnwutYA2ul-s26Xw4BkP6nqpTCoye74FvOqRPw1AOEk,9958
+cocoindex-0.1.39.dist-info/WHEEL,sha256=Kw8y023UaufEXFM028WbBCk7rJUUAhAX1Zw-54pv0m4,96
+cocoindex-0.1.39.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+cocoindex/__init__.py,sha256=pHcVzNeWewxtx9gKiPnTB8jLjuFlX9Bw9hX1cEGOWvI,806
+cocoindex/_engine.cp311-win_amd64.pyd,sha256=DM4eY3Dq3MlYvk6doOpMFTvqMzsVlsWCiK9G9XCwSco,60931584
 cocoindex/auth_registry.py,sha256=-EhwmyIwJLPZthPOvEK21TfMbxQtndRbogk8q7y4XuU,716
 cocoindex/cli.py,sha256=iJUQ2GSI8iHW3wfKRxHH5nixug92imPVmTalOjUVa3w,9237
-cocoindex/convert.py,sha256=XT747PZcerRGlsaJ69mZoG-7Z4ruwDNX1pO0VlGKlwY,7034
-cocoindex/flow.py,sha256=CjcMmx_2HcDkgy7qEs0fXRpNZFi3-4ajOAegR0veamE,24503
+cocoindex/convert.py,sha256=Co8MhwbnGWONu2mkrQYzFM1gw3stOIdFHueT1mszdV4,7085
+cocoindex/flow.py,sha256=xMJM0vzTrx4PHgQ5OxktCaBW1q54B_E11MpP-a_SD-s,28396
 cocoindex/functions.py,sha256=m7R8gNVK5RcavPY37d2lkRXwLBdMblZexCJ0JYcruss,1881
 cocoindex/index.py,sha256=32iiQI60VKhRRHle17rpoEVa_tsAHnyXXXnrLaX68uQ,557
 cocoindex/lib.py,sha256=Yv3_qKESgrjxJRYQn8a89aZcEoLZ5-xewgfEN0l9K1w,2461
 cocoindex/llm.py,sha256=POMdB-huMvPkRDpvcaBeOgXfbm0YZa0swNOdD2s4TRc,364
-cocoindex/op.py,sha256=PeoBddukU5vOBvUrFPC6fYH35L59THy9ZXVm_Vc2ng0,10998
+cocoindex/op.py,sha256=B1P7vCmrcuBuxnX8tHqGvfFgUMRBu4Om2VFnzk00MpI,11007
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/query.py,sha256=GU3V8AoSK--xOU2yX0Ao-Rz1t3mlrsJ4JIhzsLVwFH8,3297
+cocoindex/query.py,sha256=ljsylpRRxEpBXQakE2_4a89wWLYXTCE9HjacyjhxsQw,3297
 cocoindex/runtime.py,sha256=WwyDSpJhvaHjnxH4q6r2MVftJB-QkNVH_TxZLxuDtFw,1009
-cocoindex/setting.py,sha256=SODjbKgIM9_E9VEwMWcwR7BpuwEVl3Iypg1_eFdZfho,3359
+cocoindex/setting.py,sha256=dTtUD_L2b4Jf2w1-Dt8M1eFMpYQV9sMS95fJ4wB_9-I,3364
 cocoindex/setup.py,sha256=YAH9ssZB27MXYvsiqtBA1IDjAfeUci--d3cjFE76i0Y,767
 cocoindex/sources.py,sha256=W5wgiNwhTjPyNxt3DLKBbOnGwrHHfi36b9Gfr1GYhe0,1369
 cocoindex/storages.py,sha256=oJcY_zZx1wXjoJ_gitSvaFTUqhL6lUJihZQivvMMZ5U,2292
 cocoindex/tests/__init__.py,sha256=frcCV1k9oG9oKj3dpUqdJg1PxRT2RSN_XKdLCPjaYaY,2
 cocoindex/tests/test_convert.py,sha256=XbjzUoKazBJyBRCPaNElbRGoDD4NJZ9J-i4el39O7w8,15863
 cocoindex/typing.py,sha256=ZZ-QqkfNNSRkzqWahx-8w3X4upjy4kyIFr7OAPWYz5Q,9232
-cocoindex-0.1.38.dist-info/RECORD,,
+cocoindex-0.1.39.dist-info/RECORD,,