cocoindex 0.1.43__cp311-cp311-macosx_11_0_arm64.whl → 0.1.45__cp311-cp311-macosx_11_0_arm64.whl

cocoindex/functions.py

@@ -1,4 +1,5 @@
 """All builtin functions."""
+
 from typing import Annotated, Any, TYPE_CHECKING
 
 from .typing import Float32, Vector, TypeAttr
@@ -8,12 +9,15 @@ from . import op, llm
 if TYPE_CHECKING:
     import sentence_transformers
 
+
 class ParseJson(op.FunctionSpec):
     """Parse a text into a JSON object."""
 
+
 class SplitRecursively(op.FunctionSpec):
     """Split a document (in string) recursively."""
 
+
 class ExtractByLlm(op.FunctionSpec):
     """Extract information from a text using a LLM."""
 
@@ -21,6 +25,7 @@ class ExtractByLlm(op.FunctionSpec):
     output_type: type
     instruction: str | None = None
 
+
 class SentenceTransformerEmbed(op.FunctionSpec):
     """
     `SentenceTransformerEmbed` embeds a text into a vector space using the [SentenceTransformer](https://huggingface.co/sentence-transformers) library.
@@ -30,9 +35,11 @@ class SentenceTransformerEmbed(op.FunctionSpec):
         model: The name of the SentenceTransformer model to use.
         args: Additional arguments to pass to the SentenceTransformer constructor. e.g. {"trust_remote_code": True}
     """
+
     model: str
     args: dict[str, Any] | None = None
 
+
 @op.executor_class(gpu=True, cache=True, behavior_version=1)
 class SentenceTransformerEmbedExecutor:
     """Executor for SentenceTransformerEmbed."""
@@ -40,12 +47,18 @@ class SentenceTransformerEmbedExecutor:
     spec: SentenceTransformerEmbed
     _model: "sentence_transformers.SentenceTransformer"
 
-    def analyze(self, text):
-        import sentence_transformers # pylint: disable=import-outside-toplevel
+    def analyze(self, text: Any) -> type:
+        import sentence_transformers  # pylint: disable=import-outside-toplevel
+
         args = self.spec.args or {}
         self._model = sentence_transformers.SentenceTransformer(self.spec.model, **args)
         dim = self._model.get_sentence_embedding_dimension()
-        return Annotated[Vector[Float32, dim], TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value)]
+        result: type = Annotated[
+            Vector[Float32, dim],  # type: ignore
+            TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value),
+        ]
+        return result
 
     def __call__(self, text: str) -> list[Float32]:
-        return self._model.encode(text).tolist()
+        result: list[Float32] = self._model.encode(text).tolist()
+        return result

cocoindex/index.py

@@ -1,23 +1,29 @@
 from enum import Enum
 from dataclasses import dataclass
 from typing import Sequence
+
+
 class VectorSimilarityMetric(Enum):
     COSINE_SIMILARITY = "CosineSimilarity"
     L2_DISTANCE = "L2Distance"
     INNER_PRODUCT = "InnerProduct"
 
+
 @dataclass
 class VectorIndexDef:
     """
     Define a vector index on a field.
     """
+
     field_name: str
     metric: VectorSimilarityMetric
 
+
 @dataclass
 class IndexOptions:
     """
     Options for an index.
     """
+
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[VectorIndexDef] = ()

cocoindex/lib.py

@@ -1,14 +1,16 @@
 """
 Library level functions and states.
 """
+
 import warnings
 from typing import Callable, Any
 
-from . import _engine, flow, query, setting
+from . import _engine  # type: ignore
+from . import flow, query, setting
 from .convert import dump_engine_object
 
 
-def init(settings: setting.Settings | None = None):
+def init(settings: setting.Settings | None = None) -> None:
     """
     Initialize the cocoindex library.
 
@@ -19,20 +21,22 @@ def init(settings: setting.Settings | None = None):
     setting.set_app_namespace(settings.app_namespace)
 
 
-def start_server(settings: setting.ServerSettings):
+def start_server(settings: setting.ServerSettings) -> None:
     """Start the cocoindex server."""
     flow.ensure_all_flows_built()
     query.ensure_all_handlers_built()
     _engine.start_server(settings.__dict__)
 
-def stop():
+
+def stop() -> None:
     """Stop the cocoindex library."""
     _engine.stop()
 
+
 def main_fn(
-        settings: Any | None = None,
-        cocoindex_cmd: str | None = None,
-        ) -> Callable[[Callable], Callable]:
+    settings: Any | None = None,
+    cocoindex_cmd: str | None = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """
     DEPRECATED: The @cocoindex.main_fn() decorator is obsolete and has no effect.
     It will be removed in a future version, which will cause an AttributeError.
@@ -63,9 +67,10 @@ def main_fn(
         "See cocoindex <command> --help for more details.\n"
         "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n\n",
         DeprecationWarning,
-        stacklevel=2
+        stacklevel=2,
     )
 
-    def _main_wrapper(fn: Callable) -> Callable:
+    def _main_wrapper(fn: Callable[..., Any]) -> Callable[..., Any]:
         return fn
+
     return _main_wrapper

cocoindex/llm.py

@@ -1,16 +1,20 @@
 from dataclasses import dataclass
 from enum import Enum
 
+
 class LlmApiType(Enum):
     """The type of LLM API to use."""
+
     OPENAI = "OpenAi"
     OLLAMA = "Ollama"
     GEMINI = "Gemini"
     ANTHROPIC = "Anthropic"
 
+
 @dataclass
 class LlmSpec:
     """A specification for a LLM."""
+
     api_type: LlmApiType
     model: str
     address: str | None = None

cocoindex/op.py

@@ -1,6 +1,7 @@
 """
 Facilities for defining cocoindex operations.
 """
+
 import asyncio
 import dataclasses
 import inspect
@@ -10,40 +11,58 @@ from enum import Enum
 
 from .typing import encode_enriched_type, resolve_forward_ref
 from .convert import encode_engine_value, make_engine_value_decoder
-from . import _engine
+from . import _engine  # type: ignore
+
 
 class OpCategory(Enum):
     """The category of the operation."""
+
     FUNCTION = "function"
     SOURCE = "source"
     STORAGE = "storage"
     DECLARATION = "declaration"
+
+
 @dataclass_transform()
 class SpecMeta(type):
     """Meta class for spec classes."""
-    def __new__(mcs, name, bases, attrs, category: OpCategory | None = None):
+
+    def __new__(
+        mcs,
+        name: str,
+        bases: tuple[type, ...],
+        attrs: dict[str, Any],
+        category: OpCategory | None = None,
+    ) -> type:
         cls: type = super().__new__(mcs, name, bases, attrs)
         if category is not None:
             # It's the base class.
-            setattr(cls, '_op_category', category)
+            setattr(cls, "_op_category", category)
         else:
             # It's the specific class providing specific fields.
             cls = dataclasses.dataclass(cls)
         return cls
 
-class SourceSpec(metaclass=SpecMeta, category=OpCategory.SOURCE): # pylint: disable=too-few-public-methods
+
+class SourceSpec(metaclass=SpecMeta, category=OpCategory.SOURCE):  # pylint: disable=too-few-public-methods
     """A source spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
-class FunctionSpec(metaclass=SpecMeta, category=OpCategory.FUNCTION): # pylint: disable=too-few-public-methods
+
+class FunctionSpec(metaclass=SpecMeta, category=OpCategory.FUNCTION):  # pylint: disable=too-few-public-methods
     """A function spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
-class StorageSpec(metaclass=SpecMeta, category=OpCategory.STORAGE): # pylint: disable=too-few-public-methods
+
+class StorageSpec(metaclass=SpecMeta, category=OpCategory.STORAGE):  # pylint: disable=too-few-public-methods
     """A storage spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
 
-class DeclarationSpec(metaclass=SpecMeta, category=OpCategory.DECLARATION): # pylint: disable=too-few-public-methods
+
+class DeclarationSpec(metaclass=SpecMeta, category=OpCategory.DECLARATION):  # pylint: disable=too-few-public-methods
     """A declaration spec. All its subclass can be instantiated similar to a dataclass, i.e. ClassName(field1=value1, field2=value2, ...)"""
+
+
 class Executor(Protocol):
     """An executor for an operation."""
+
     op_category: OpCategory
 
 
@@ -55,7 +74,9 @@ class _FunctionExecutorFactory:
         self._spec_cls = spec_cls
         self._executor_cls = executor_cls
 
-    def __call__(self, spec: dict[str, Any], *args, **kwargs):
+    def __call__(
+        self, spec: dict[str, Any], *args: Any, **kwargs: Any
+    ) -> tuple[dict[str, Any], Executor]:
         spec = self._spec_cls(**spec)
         executor = self._executor_cls(spec)
         result_type = executor.analyze(*args, **kwargs)
@@ -64,6 +85,7 @@ class _FunctionExecutorFactory:
 
 _gpu_dispatch_lock = asyncio.Lock()
 
+
 @dataclasses.dataclass
 class OpArgs:
     """
@@ -72,44 +94,50 @@ class OpArgs:
     - behavior_version: The behavior version of the executor. Cache will be invalidated if it
       changes. Must be provided if `cache` is True.
     """
+
     gpu: bool = False
     cache: bool = False
     behavior_version: int | None = None
 
-def _to_async_call(call: Callable) -> Callable[..., Awaitable[Any]]:
+
+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))
 
+
 def _register_op_factory(
-        category: OpCategory,
-        expected_args: list[tuple[str, inspect.Parameter]],
-        expected_return,
-        executor_cls: type,
-        spec_cls: type,
-        op_args: OpArgs,
-    ):
+    category: OpCategory,
+    expected_args: list[tuple[str, inspect.Parameter]],
+    expected_return: Any,
+    executor_cls: type,
+    spec_cls: type,
+    op_args: OpArgs,
+) -> type:
     """
     Register an op factory.
     """
+
     class _Fallback:
-        def enable_cache(self):
+        def enable_cache(self) -> bool:
             return op_args.cache
 
-        def behavior_version(self):
+        def behavior_version(self) -> int | None:
             return op_args.behavior_version
 
-    class _WrappedClass(executor_cls, _Fallback):
+    class _WrappedClass(executor_cls, _Fallback):  # type: ignore[misc]
         _args_decoders: list[Callable[[Any], Any]]
-        _kwargs_decoders: dict[str, Callable[[str, Any], Any]]
-        _acall: Callable
+        _kwargs_decoders: dict[str, Callable[[Any], Any]]
+        _acall: Callable[..., Awaitable[Any]]
 
-        def __init__(self, spec):
+        def __init__(self, spec: Any) -> None:
             super().__init__()
             self.spec = spec
             self._acall = _to_async_call(super().__call__)
 
-        def analyze(self, *args: _engine.OpArgSchema, **kwargs: _engine.OpArgSchema):
+        def analyze(
+            self, *args: _engine.OpArgSchema, **kwargs: _engine.OpArgSchema
+        ) -> Any:
             """
             Analyze the spec and arguments. In this phase, argument types should be validated.
             It should return the expected result type for the current op.
@@ -122,15 +150,21 @@ def _register_op_factory(
             for arg in args:
                 if next_param_idx >= len(expected_args):
                     raise ValueError(
-                        f"Too many arguments passed in: {len(args)} > {len(expected_args)}")
+                        f"Too many arguments passed in: {len(args)} > {len(expected_args)}"
+                    )
                 arg_name, arg_param = expected_args[next_param_idx]
                 if arg_param.kind in (
-                    inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.VAR_KEYWORD):
+                    inspect.Parameter.KEYWORD_ONLY,
+                    inspect.Parameter.VAR_KEYWORD,
+                ):
                     raise ValueError(
-                        f"Too many positional arguments passed in: {len(args)} > {next_param_idx}")
+                        f"Too many positional arguments passed in: {len(args)} > {next_param_idx}"
+                    )
                 self._args_decoders.append(
                     make_engine_value_decoder(
-                        [arg_name], arg.value_type['type'], arg_param.annotation))
+                        [arg_name], arg.value_type["type"], arg_param.annotation
+                    )
+                )
                 if arg_param.kind != inspect.Parameter.VAR_POSITIONAL:
                     next_param_idx += 1
 
@@ -138,45 +172,72 @@ def _register_op_factory(
 
             for kwarg_name, kwarg in kwargs.items():
                 expected_arg = next(
-                    (arg for arg in expected_kwargs
-                      if (arg[0] == kwarg_name and arg[1].kind in (
-                          inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD))
-                        or arg[1].kind == inspect.Parameter.VAR_KEYWORD),
-                    None)
+                    (
+                        arg
+                        for arg in expected_kwargs
+                        if (
+                            arg[0] == kwarg_name
+                            and arg[1].kind
+                            in (
+                                inspect.Parameter.KEYWORD_ONLY,
+                                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                            )
+                        )
+                        or arg[1].kind == inspect.Parameter.VAR_KEYWORD
+                    ),
+                    None,
+                )
                 if expected_arg is None:
-                    raise ValueError(f"Unexpected keyword argument passed in: {kwarg_name}")
+                    raise ValueError(
+                        f"Unexpected keyword argument passed in: {kwarg_name}"
+                    )
                 arg_param = expected_arg[1]
                 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
-                            if arg.default is inspect.Parameter.empty
-                               and (arg.kind == inspect.Parameter.POSITIONAL_ONLY or
-                                    (arg.kind in (inspect.Parameter.KEYWORD_ONLY,
-                                                  inspect.Parameter.POSITIONAL_OR_KEYWORD)
-                                    and name not in kwargs))]
+                    [kwarg_name], kwarg.value_type["type"], arg_param.annotation
+                )
+
+            missing_args = [
+                name
+                for (name, arg) in expected_kwargs
+                if arg.default is inspect.Parameter.empty
+                and (
+                    arg.kind == inspect.Parameter.POSITIONAL_ONLY
+                    or (
+                        arg.kind
+                        in (
+                            inspect.Parameter.KEYWORD_ONLY,
+                            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                        )
+                        and name not in kwargs
+                    )
+                )
+            ]
             if len(missing_args) > 0:
                 raise ValueError(f"Missing arguments: {', '.join(missing_args)}")
 
-            prepare_method = getattr(executor_cls, 'analyze', None)
+            prepare_method = getattr(executor_cls, "analyze", None)
             if prepare_method is not None:
                 return prepare_method(self, *args, **kwargs)
             else:
                 return expected_return
 
-        async def prepare(self):
+        async def prepare(self) -> None:
             """
             Prepare for execution.
             It's executed after `analyze` and before any `__call__` execution.
             """
-            setup_method = getattr(super(), 'prepare', None)
+            setup_method = getattr(super(), "prepare", None)
             if setup_method is not None:
                 await _to_async_call(setup_method)()
 
-        async def __call__(self, *args, **kwargs):
-            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()}
+        async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            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:
                 # For GPU executions, data-level parallelism is applied, so we don't want to
@@ -198,13 +259,15 @@ def _register_op_factory(
 
     if category == OpCategory.FUNCTION:
         _engine.register_function_factory(
-            spec_cls.__name__, _FunctionExecutorFactory(spec_cls, _WrappedClass))
+            spec_cls.__name__, _FunctionExecutorFactory(spec_cls, _WrappedClass)
+        )
     else:
         raise ValueError(f"Unsupported executor type {category}")
 
     return _WrappedClass
 
-def executor_class(**args) -> Callable[[type], type]:
+
+def executor_class(**args: Any) -> Callable[[type], type]:
     """
     Decorate a class to provide an executor for an op.
     """
@@ -216,9 +279,9 @@ def executor_class(**args) -> Callable[[type], type]:
         """
         # Use `__annotations__` instead of `get_type_hints`, to avoid resolving forward references.
         type_hints = cls.__annotations__
-        if 'spec' not in type_hints:
+        if "spec" not in type_hints:
             raise TypeError("Expect a `spec` field with type hint")
-        spec_cls = resolve_forward_ref(type_hints['spec'])
+        spec_cls = resolve_forward_ref(type_hints["spec"])
         sig = inspect.signature(cls.__call__)
         return _register_op_factory(
             category=spec_cls._op_category,
@@ -226,34 +289,35 @@ def executor_class(**args) -> Callable[[type], type]:
             expected_return=sig.return_annotation,
             executor_cls=cls,
             spec_cls=spec_cls,
-            op_args=op_args)
+            op_args=op_args,
+        )
 
     return _inner
 
-def function(**args) -> Callable[[Callable], FunctionSpec]:
+
+def function(**args: Any) -> Callable[[Callable[..., Any]], FunctionSpec]:
     """
     Decorate a function to provide a function for an op.
     """
     op_args = OpArgs(**args)
 
-    def _inner(fn: Callable) -> FunctionSpec:
-
+    def _inner(fn: Callable[..., Any]) -> FunctionSpec:
         # Convert snake case to camel case.
-        op_name = ''.join(word.capitalize() for word in fn.__name__.split('_'))
+        op_name = "".join(word.capitalize() for word in fn.__name__.split("_"))
         sig = inspect.signature(fn)
 
         class _Executor:
-            def __call__(self, *args, **kwargs):
+            def __call__(self, *args: Any, **kwargs: Any) -> Any:
                 return fn(*args, **kwargs)
 
         class _Spec(FunctionSpec):
-            pass
+            def __call__(self, *args: Any, **kwargs: Any) -> Any:
+                return fn(*args, **kwargs)
 
         _Spec.__name__ = op_name
         _Spec.__doc__ = fn.__doc__
         _Spec.__module__ = fn.__module__
         _Spec.__qualname__ = fn.__qualname__
-        _Spec.__wrapped__ = fn
 
         _register_op_factory(
             category=OpCategory.FUNCTION,
@@ -261,7 +325,8 @@ def function(**args) -> Callable[[Callable], FunctionSpec]:
             expected_return=sig.return_annotation,
             executor_cls=_Executor,
             spec_cls=_Spec,
-            op_args=op_args)
+            op_args=op_args,
+        )
 
         return _Spec()
 

cocoindex/query.py

@@ -4,25 +4,29 @@ from threading import Lock
 
 from . import flow as fl
 from . import index
-from . import _engine
+from . import _engine  # type: ignore
 
 _handlers_lock = Lock()
 _handlers: dict[str, _engine.SimpleSemanticsQueryHandler] = {}
 
+
 @dataclass
 class SimpleSemanticsQueryInfo:
     """
     Additional information about the query.
     """
+
     similarity_metric: index.VectorSimilarityMetric
     query_vector: list[float]
     vector_field_name: str
 
+
 @dataclass
 class QueryResult:
     """
     A single result from the query.
     """
+
     data: dict[str, Any]
     score: float
 
@@ -31,6 +35,7 @@ class SimpleSemanticsQueryHandler:
     """
     A query handler that uses simple semantics to query the index.
     """
+
     _lazy_query_handler: Callable[[], _engine.SimpleSemanticsQueryHandler]
 
     def __init__(
@@ -38,22 +43,28 @@ class SimpleSemanticsQueryHandler:
         name: str,
         flow: fl.Flow,
         target_name: str,
-        query_transform_flow: Callable[..., fl.DataSlice],
-        default_similarity_metric: index.VectorSimilarityMetric = index.VectorSimilarityMetric.COSINE_SIMILARITY) -> None:
-
+        query_transform_flow: Callable[..., fl.DataSlice[Any]],
+        default_similarity_metric: index.VectorSimilarityMetric = index.VectorSimilarityMetric.COSINE_SIMILARITY,
+    ) -> None:
         engine_handler = None
         lock = Lock()
+
         def _lazy_handler() -> _engine.SimpleSemanticsQueryHandler:
             nonlocal engine_handler, lock
             if engine_handler is None:
                 with lock:
                     if engine_handler is None:
                         engine_handler = _engine.SimpleSemanticsQueryHandler(
-                            flow.internal_flow(), target_name,
-                            fl.TransformFlow(query_transform_flow, [str]).internal_flow(),
-                            default_similarity_metric.value)
+                            flow.internal_flow(),
+                            target_name,
+                            fl.TransformFlow(
+                                query_transform_flow, [str]
+                            ).internal_flow(),
+                            default_similarity_metric.value,
+                        )
                         engine_handler.register_query_handler(name)
             return engine_handler
+
         self._lazy_query_handler = _lazy_handler
 
         with _handlers_lock:
@@ -65,24 +76,36 @@ class SimpleSemanticsQueryHandler:
         """
         return self._lazy_query_handler()
 
-    def search(self, query: str, limit: int, vector_field_name: str | None = None,
-               similarity_metric: index.VectorSimilarityMetric | None = None
-               ) -> tuple[list[QueryResult], SimpleSemanticsQueryInfo]:
+    def search(
+        self,
+        query: str,
+        limit: int,
+        vector_field_name: str | None = None,
+        similarity_metric: index.VectorSimilarityMetric | None = None,
+    ) -> tuple[list[QueryResult], SimpleSemanticsQueryInfo]:
         """
         Search the index with the given query, limit, vector field name, and similarity metric.
         """
         internal_results, internal_info = self.internal_handler().search(
-            query, limit, vector_field_name,
-            similarity_metric.value if similarity_metric is not None else None)
-        results = [QueryResult(data=result['data'],  score=result['score'])
-                   for result in internal_results]
+            query,
+            limit,
+            vector_field_name,
+            similarity_metric.value if similarity_metric is not None else None,
+        )
+        results = [
+            QueryResult(data=result["data"], score=result["score"])
+            for result in internal_results
+        ]
         info = SimpleSemanticsQueryInfo(
-            similarity_metric=index.VectorSimilarityMetric(internal_info['similarity_metric']),
-            query_vector=internal_info['query_vector'],
-            vector_field_name=internal_info['vector_field_name']
+            similarity_metric=index.VectorSimilarityMetric(
+                internal_info["similarity_metric"]
+            ),
+            query_vector=internal_info["query_vector"],
+            vector_field_name=internal_info["vector_field_name"],
         )
         return results, info
 
+
 def ensure_all_handlers_built() -> None:
     """
     Ensure all handlers are built.