cocoindex 0.2.10__cp311-abi3-macosx_10_12_x86_64.whl → 0.2.11__cp311-abi3-macosx_10_12_x86_64.whl

cocoindex/__init__.py

@@ -24,6 +24,7 @@ from .llm import LlmSpec, LlmApiType
 from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
 from .setting import DatabaseConnectionSpec, Settings, ServerSettings
 from .setting import get_app_namespace
+from .query_handler import QueryHandlerResultFields, QueryInfo, QueryOutput
 from .typing import (
     Int64,
     Float32,
@@ -95,4 +96,8 @@ __all__ = [
     "Range",
     "Vector",
     "Json",
+    # Query handler
+    "QueryHandlerResultFields",
+    "QueryInfo",
+    "QueryOutput",
 ]

cocoindex/auth_registry.py

@@ -11,18 +11,6 @@ from .convert import dump_engine_object
 
 T = TypeVar("T")
 
-# Global atomic counter for generating unique auth entry keys
-_counter_lock = threading.Lock()
-_auth_key_counter = 0
-
-
-def _generate_auth_key() -> str:
-    """Generate a unique auth entry key using a global atomic counter."""
-    global _auth_key_counter  # pylint: disable=global-statement
-    with _counter_lock:
-        _auth_key_counter += 1
-        return f"__auth_{_auth_key_counter}"
-
 
 @dataclass
 class TransientAuthEntryReference(Generic[T]):
@@ -37,7 +25,8 @@ class AuthEntryReference(TransientAuthEntryReference[T]):
 
 def add_transient_auth_entry(value: T) -> TransientAuthEntryReference[T]:
     """Add an auth entry to the registry. Returns its reference."""
-    return add_auth_entry(_generate_auth_key(), value)
+    key = _engine.add_transient_auth_entry(dump_engine_object(value))
+    return TransientAuthEntryReference(key)
 
 
 def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]:

cocoindex/convert.py

@@ -616,6 +616,8 @@ def dump_engine_object(v: Any) -> Any:
         return s
     elif isinstance(v, (list, tuple)):
         return [dump_engine_object(item) for item in v]
+    elif isinstance(v, np.ndarray):
+        return v.tolist()
     elif isinstance(v, dict):
         return {k: dump_engine_object(v) for k, v in v.items()}
     return v

cocoindex/flow.py

@@ -38,9 +38,10 @@ from .convert import (
     make_engine_value_encoder,
 )
 from .op import FunctionSpec
-from .runtime import execution_context
+from .runtime import execution_context, to_async_call
 from .setup import SetupChangeBundle
 from .typing import analyze_type_info, encode_enriched_type
+from .query_handler import QueryHandlerInfo, QueryHandlerResultFields
 from .validation import (
     validate_flow_name,
     validate_full_flow_name,
@@ -694,23 +695,18 @@ class Flow:
     """
 
     _name: str
-    _lazy_engine_flow: Callable[[], _engine.Flow] | None
+    _engine_flow_creator: Callable[[], _engine.Flow]
+
+    _lazy_flow_lock: Lock
+    _lazy_query_handler_args: list[tuple[Any, ...]]
+    _lazy_engine_flow: _engine.Flow | None = None
 
     def __init__(self, name: str, engine_flow_creator: Callable[[], _engine.Flow]):
         validate_flow_name(name)
         self._name = name
-        engine_flow = None
-        lock = Lock()
-
-        def _lazy_engine_flow() -> _engine.Flow:
-            nonlocal engine_flow, lock
-            if engine_flow is None:
-                with lock:
-                    if engine_flow is None:
-                        engine_flow = engine_flow_creator()
-            return engine_flow
-
-        self._lazy_engine_flow = _lazy_engine_flow
+        self._engine_flow_creator = engine_flow_creator
+        self._lazy_flow_lock = Lock()
+        self._lazy_query_handler_args = []
 
     def _render_spec(self, verbose: bool = False) -> Tree:
         """
@@ -794,15 +790,33 @@ class Flow:
         """
         Get the engine flow.
         """
-        if self._lazy_engine_flow is None:
-            raise RuntimeError(f"Flow {self.full_name} is already removed")
-        return self._lazy_engine_flow()
+        if self._lazy_engine_flow is not None:
+            return self._lazy_engine_flow
+        return self._internal_flow()
 
     async def internal_flow_async(self) -> _engine.Flow:
         """
         Get the engine flow. The async version.
         """
-        return await asyncio.to_thread(self.internal_flow)
+        if self._lazy_engine_flow is not None:
+            return self._lazy_engine_flow
+        return await asyncio.to_thread(self._internal_flow)
+
+    def _internal_flow(self) -> _engine.Flow:
+        """
+        Get the engine flow. The async version.
+        """
+        with self._lazy_flow_lock:
+            if self._lazy_engine_flow is not None:
+                return self._lazy_engine_flow
+
+            engine_flow = self._engine_flow_creator()
+            self._lazy_engine_flow = engine_flow
+            for args in self._lazy_query_handler_args:
+                engine_flow.add_query_handler(*args)
+            self._lazy_query_handler_args = []
+
+            return engine_flow
 
     def setup(self, report_to_stdout: bool = False) -> None:
         """
@@ -847,6 +861,53 @@ class Flow:
         with _flows_lock:
             del _flows[self.name]
 
+    def add_query_handler(
+        self,
+        name: str,
+        handler: Callable[[str], Any],
+        /,
+        *,
+        result_fields: QueryHandlerResultFields | None = None,
+    ) -> None:
+        """
+        Add a query handler to the flow.
+        """
+        async_handler = to_async_call(handler)
+
+        async def _handler(query: str) -> dict[str, Any]:
+            handler_result = await async_handler(query)
+            return {
+                "results": [
+                    [(k, dump_engine_object(v)) for (k, v) in result.items()]
+                    for result in handler_result.results
+                ],
+                "query_info": dump_engine_object(handler_result.query_info),
+            }
+
+        handler_info = dump_engine_object(QueryHandlerInfo(result_fields=result_fields))
+        with self._lazy_flow_lock:
+            if self._lazy_engine_flow is not None:
+                self._lazy_engine_flow.add_query_handler(name, _handler, handler_info)
+            else:
+                self._lazy_query_handler_args.append((name, _handler, handler_info))
+
+    def query_handler(
+        self,
+        name: str | None = None,
+        result_fields: QueryHandlerResultFields | None = None,
+    ) -> Callable[[Callable[[str], Any]], Callable[[str], Any]]:
+        """
+        A decorator to declare a query handler.
+        """
+
+        def _inner(handler: Callable[[str], Any]) -> Callable[[str], Any]:
+            self.add_query_handler(
+                name or handler.__qualname__, handler, result_fields=result_fields
+            )
+            return handler
+
+        return _inner
+
 
 def _create_lazy_flow(
     name: str | None, fl_def: Callable[[FlowBuilder, DataScope], None]

cocoindex/op.py

@@ -2,7 +2,6 @@
 Facilities for defining cocoindex operations.
 """
 
-import asyncio
 import dataclasses
 import inspect
 from enum import Enum
@@ -32,6 +31,7 @@ from .typing import (
     AnalyzedAnyType,
     AnalyzedDictType,
 )
+from .runtime import to_async_call
 
 
 class OpCategory(Enum):
@@ -150,12 +150,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]
@@ -319,8 +313,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 = []
@@ -461,12 +455,12 @@ class _TargetConnector:
         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(

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-0.2.10.dist-info → cocoindex-0.2.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.2.10
+Version: 0.2.11
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent

{cocoindex-0.2.10.dist-info → cocoindex-0.2.11.dist-info}/RECORD

@@ -1,20 +1,21 @@
-cocoindex-0.2.10.dist-info/METADATA,sha256=W1EtSw4PWmq2rT-jS5OA_3x85UIXrJmgtyJLRo7Zhgc,12999
-cocoindex-0.2.10.dist-info/WHEEL,sha256=N8W3-0eDM6igWj-H12r7VkxoMaJIqJLxUyWCFstEaGg,105
-cocoindex-0.2.10.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.2.10.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=s4PFFZPP5i9n3Gw4VuPYKYYU5pc2bK9W9ipjRkQoygw,717804
-cocoindex/__init__.py,sha256=r0WRAMjHW5gq85fgvxD_EcWW1um4irffpY-ifMrnhKw,2264
-cocoindex/_engine.abi3.so,sha256=MXkLdSPva8dt-wMRatsdelrGHLsCFyFqp0IyttZwuUw,69966180
-cocoindex/auth_registry.py,sha256=PE1-kVkcyC1G2C_V7b1kvYzeq73OFQehWKQP7ln7fJ8,1478
+cocoindex-0.2.11.dist-info/METADATA,sha256=u-dE4rSfVStoZHzhM01B2D0N-8AqLGiV3NAYs1zs19U,12999
+cocoindex-0.2.11.dist-info/WHEEL,sha256=N8W3-0eDM6igWj-H12r7VkxoMaJIqJLxUyWCFstEaGg,105
+cocoindex-0.2.11.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.2.11.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=o8g8gIN-zaK02_7ZMPlqompannwA2SwEB0e30CKRN1E,717804
+cocoindex/__init__.py,sha256=AsoNLBgjJ-1AEKtwTnENJM9dEilHTq6rJ4qWUzGKSdc,2428
+cocoindex/_engine.abi3.so,sha256=H-x9Yqz5RWpB0vkeGPpjb3QuOX1h8-ocQlqrPaXCnqw,70039588
+cocoindex/auth_registry.py,sha256=9GgeuyE48vYhsJQJgI-uRjUviZkrQE63MUenP3OREko,1145
 cocoindex/cli.py,sha256=69X30bFTFdM7c0_6lgIHR19CeQ7UEkobEQYihy8IdOQ,21599
-cocoindex/convert.py,sha256=Eh9Co37BtW_PK3Oi-pFEiFt8cc_6g7XLcurV-NeH6GU,22090
-cocoindex/flow.py,sha256=GGk2wtCGfP9cfgdk23SmwR2yBLAuoCDvBGmzmp-XEzM,37591
+cocoindex/convert.py,sha256=Apk5PU4XGnSvEE2HMRN2W9H4ApBIHnkupegv_-TaVfA,22152
+cocoindex/flow.py,sha256=u_ta8wW5oGTDINs2dY4d_1hybySp3wuX79VJ4BDXZDc,39767
 cocoindex/functions.py,sha256=09erNt3WbzY9l1KER-akBF2O5-6xEahV2ORBECaL6yk,12260
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
 cocoindex/lib.py,sha256=0XheDF7fiFdqExpdqzU-VKun_Zll6DwZ5JfTm7u42aY,2284
 cocoindex/llm.py,sha256=Pv_cdnRngTLtuLU9AUmS8izIHhcKVnuBNolC33f9BDI,851
-cocoindex/op.py,sha256=3jTEpmzHPjg-euZhXjQqgF0i8V2LHUdkq7bp4AAIHRM,21627
+cocoindex/op.py,sha256=SRtI5ApgRQJHNsljd9Z2qNdS68_pYYQj9oVg3Bqc8W4,21416
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/runtime.py,sha256=povilB3HH3y1JF-yxKwU-pD8n2WnAqyQxIgvXXHNc60,1080
+cocoindex/query_handler.py,sha256=o9XFnpXWa00pAFImJ9oq-_kjI7w69dgCalkq63H9mA8,1192
+cocoindex/runtime.py,sha256=4NxcltaDZvA3RR3Pnt6gH_f99jcWSyMH_1Xi5BjbtwY,1342
 cocoindex/setting.py,sha256=1Dx8ktjwf-8BiXrbsmfn5Mzudb2SQYqFdRnSNGVKaLk,4960
 cocoindex/setup.py,sha256=7uIHKN4FOCuoidPXcKyGTrkqpkl9luL49-6UcnMxYzw,3068
 cocoindex/sources.py,sha256=FYz7cWYasLGDaYoIEQ1dF2uprgUETHWsTIrIS7n6pQE,3188
@@ -30,4 +31,4 @@ cocoindex/typing.py,sha256=lEQYIzAGVKQp6RnhyeopY9Q7xEED7yQj3ZMxvTPblV8,14200
 cocoindex/user_app_loader.py,sha256=bc3Af-gYRxJ9GpObtpjegZY855oQBCv5FGkrkWV2yGY,1873
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
 cocoindex/validation.py,sha256=PZnJoby4sLbsmPv9fOjOQXuefjfZ7gmtsiTGU8SH-tc,3090
-cocoindex-0.2.10.dist-info/RECORD,,
+cocoindex-0.2.11.dist-info/RECORD,,

{cocoindex-0.2.10.dist-info → cocoindex-0.2.11.dist-info}/licenses/THIRD_PARTY_NOTICES.html

@@ -2428,7 +2428,7 @@ Software.
                 <h3 id="Apache-2.0">Apache License 2.0</h3>
                 <h4>Used by:</h4>
                 <ul class="license-used-by">
-                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.2.10</a></li>
+                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.2.11</a></li>
                     <li><a href=" https://github.com/awesomized/crc-fast-rust ">crc-fast 1.3.0</a></li>
                     <li><a href=" https://github.com/qdrant/rust-client ">qdrant-client 1.15.0</a></li>
                 </ul>