cocoindex 0.2.10__cp311-abi3-win_amd64.whl → 0.2.11__cp311-abi3-win_amd64.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=KQcBWfVM6cEl93NcnSCeaPZx3pQSmkQr030zQZ2vvAw,13215
-cocoindex-0.2.10.dist-info/WHEEL,sha256=8hEf8NzM1FnmM77AjVt5h8nDuYkN3UqZ79LoIAHXeRE,95
-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=5QfLfEQ3NpWzhLCbOCLFcvNd-tvCgp_d9_MZjksSFqE,2362
-cocoindex/_engine.pyd,sha256=PuMNCJ6Oax64wTR7hDa4bP_3RxgsRSxQrdSDinbyaxk,72942592
-cocoindex/auth_registry.py,sha256=Qq1IVZb-7K4luRrQSDlOPbISnGEZ4kIDsrCU8H2ARw0,1529
+cocoindex-0.2.11.dist-info/METADATA,sha256=azjBskKzo4W0_RuO8XyolrdX9B7AAv30ryyfVn2kZ4c,13215
+cocoindex-0.2.11.dist-info/WHEEL,sha256=8hEf8NzM1FnmM77AjVt5h8nDuYkN3UqZ79LoIAHXeRE,95
+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=7FXm4Q_I_aLgVxY2nd5PTWQ6Z92mTjJ1B5WxTHpj6p8,2531
+cocoindex/_engine.pyd,sha256=lVSYJ3dZJ5ouqBsel_g3qXP5gzcBHSLD5kWujb5AlFs,73029632
+cocoindex/auth_registry.py,sha256=vfVen6nm_sVMoUdXKf4ucrO-wOudgculrxJ2Ob94Lwo,1185
 cocoindex/cli.py,sha256=Wt8wUjWqHEBoR3gDnHqIplLaL_tbStMCHAOuYhYq84A,22338
-cocoindex/convert.py,sha256=iKVufQXjsrvS8psKKZZYxuLohJvBIVctgKaRLLt8jK0,22711
-cocoindex/flow.py,sha256=w6FimPK4Nll9nYccbX-vjDMXzYb2lV66GmgikAMR5UQ,38802
+cocoindex/convert.py,sha256=AhF33zfRFeqnRb8fiEzmhD7S4EWJHyeXW0RrB6_KNPw,22775
+cocoindex/flow.py,sha256=TD-9w2uB-hvOVrs8CTE2lqKKzewKRT2qbJ3SAYhiAxE,41039
 cocoindex/functions.py,sha256=CtiwTVW6g4BtO5_EvVcij7Si4Bx-axnM1hsdU43aM4g,12617
 cocoindex/index.py,sha256=GrqTm1rLwICQ8hadtNvJAxVg7GWMvtMmFcbiNtNzmP0,569
 cocoindex/lib.py,sha256=-Ys4a7BdDOW_ctz9vEBA5fSOqwsY_IKLHxaDwxvoXRg,2359
 cocoindex/llm.py,sha256=JM5EPup7FZojynCI0rg4gnMBYmdPZpaHJbVBz3f7BAI,897
-cocoindex/op.py,sha256=yO8QdfauUEe-3n1xZmhHhyr0EHqYEokcJRxtd3OtMo8,22255
+cocoindex/op.py,sha256=6y7UHk0mThJVfYzPfyefjxKTAXIWoUp6pvYqNUo6HLI,22038
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/runtime.py,sha256=6mE-jR1Kh5c4GETDvBgwjXZq69TK5rh1qNpaseRDZnw,1117
+cocoindex/query_handler.py,sha256=iC4_j0yiQrW1izHDJUgaQyNrTpyfTvYIOwoshYNclKk,1243
+cocoindex/runtime.py,sha256=ZjyPmohIWGFDRaOdS0MgeVuf4Ra2Ggbp6xraxwvHFBo,1386
 cocoindex/setting.py,sha256=oohs7H52_h1AulDcp96cS9SVAAm_uFaQyJYog9EvCj0,5131
 cocoindex/setup.py,sha256=KbJvmeFu0NbeoH-5iDmHZP86f26HIId8kHmGUNZAePI,3160
 cocoindex/sources.py,sha256=Ij8LyYbM49zjPg-pJHMbas8sdLquEcgvmuYfv7wCdjI,3290
@@ -30,4 +31,4 @@ cocoindex/typing.py,sha256=bkKcp4G_xQODwQn92sfeNi-kXSWnbtlwTpFbrMGCnj4,14673
 cocoindex/user_app_loader.py,sha256=ZkvUG9aJNNECAjwTY0ZYtNpFd9dNBPVoPKGTtB7dSZg,1926
 cocoindex/utils.py,sha256=U3W39zD2uZpXX8v84tJD7sRmbC5ar3z_ljAP1cJrYXI,618
 cocoindex/validation.py,sha256=4ZjsW-SZT8X_TEEhEE6QG6D-8Oq_TkPAhTqP0mdFYSE,3194
-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>