feat(query-handler): expose query handler API in Python SDK (#1011)

Author: georgeh0

python/cocoindex/__init__.py

@@ -24,6 +24,7 @@
 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 @@
     "Range",
     "Vector",
     "Json",
+    # Query handler
+    "QueryHandlerResultFields",
+    "QueryInfo",
+    "QueryOutput",
 ]

python/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

python/cocoindex/flow.py

@@ -38,9 +38,10 @@
     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,
@@ -697,13 +698,15 @@ class Flow:
     _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
         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:
         """
@@ -809,6 +812,9 @@ def _internal_flow(self) -> _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
 
@@ -855,6 +861,43 @@ def close(self) -> None:
         with _flows_lock:
             del _flows[self.name]
 
+    def add_query_handler(
+        self,
+        name: str,
+        handler: Callable[[str], Any],
+        /,
+        *,
+        result_fields: QueryHandlerResultFields | None = None,
+    ) -> None:
+        async_handler = to_async_call(handler)
+
+        async def _handler(query: str) -> dict[str, Any]:
+            handler_result = await async_handler(query)
+            return {
+                "results": dump_engine_object(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]]:
+        def _inner(handler: Callable[[str], Any]) -> Callable[[str], Any]:
+            self.add_query_handler(
+                name or handler.__name__, handler, result_fields=result_fields
+            )
+            return handler
+
+        return _inner
+
 
 def _create_lazy_flow(
     name: str | None, fl_def: Callable[[FlowBuilder, DataScope], None]

python/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 @@
     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 @@ async def prepare(self) -> None:
             """
             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 @@ def __init__(self, spec_cls: type, connector_cls: type):
         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(

python/cocoindex/query.py

@@ -0,0 +1,48 @@
+import dataclasses
+import numpy as np
+from numpy import typing as npt
+from typing import Generic, TypeVar
+
+
+@dataclasses.dataclass
+class QueryHandlerResultFields:
+    """
+    Specify field names in the returned query handler.
+    """
+
+    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
+
+
+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 | None = None

python/cocoindex/query_handler.py

@@ -5,7 +5,8 @@
 
 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 @@ def run(self, coro: Coroutine[Any, Any, T]) -> T:
 
 
 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))

python/cocoindex/runtime.py

@@ -434,7 +434,12 @@ impl Flow {
         SetupChangeBundle(Arc::new(bundle))
     }
 
-    pub fn add_query_handler(&self, name: String, handler: Py<PyAny>) -> PyResult<()> {
+    pub fn add_query_handler(
+        &self,
+        name: String,
+        handler: Py<PyAny>,
+        handler_info: Pythonized<Option<QueryHandlerInfo>>,
+    ) -> PyResult<()> {
         struct PyQueryHandler {
             handler: Py<PyAny>,
         }
@@ -483,7 +488,7 @@ impl Flow {
         handlers.insert(
             name,
             QueryHandlerContext {
-                info: Arc::new(QueryHandlerInfo {}),
+                info: Arc::new(handler_info.into_inner().unwrap_or_default()),
                 handler: Arc::new(PyQueryHandler { handler }),
             },
         );

src/py/mod.rs

@@ -1,7 +1,16 @@
 use crate::prelude::*;
 
-#[derive(Serialize)]
-pub struct QueryHandlerInfo {}
+#[derive(Serialize, Deserialize, Default)]
+pub struct QueryHandlerResultFields {
+    embedding: Vec<String>,
+    score: Option<String>,
+}
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct QueryHandlerInfo {
+    #[serde(default)]
+    result_fields: QueryHandlerResultFields,
+}
 
 #[derive(Serialize, Deserialize)]
 pub struct QueryInput {
@@ -15,7 +24,7 @@ pub struct QueryInfo {
 
 #[derive(Serialize, Deserialize)]
 pub struct QueryOutput {
-    pub results: Vec<IndexMap<String, serde_json::Value>>,
+    pub results: Vec<HashMap<String, serde_json::Value>>,
     pub query_info: QueryInfo,
 }