diff --git a/python/cocoindex/cli.py b/python/cocoindex/cli.py index 4b2a8540..f757cfbc 100644 --- a/python/cocoindex/cli.py +++ b/python/cocoindex/cli.py @@ -1,4 +1,3 @@ -import asyncio import click import datetime from rich.console import Console @@ -21,7 +20,7 @@ def ls(show_all: bool): """ List all flows. """ - current_flow_names = [fl.name for fl in flow.flows()] + current_flow_names = flow.flow_names() persisted_flow_names = flow_names_with_setup() remaining_persisted_flow_names = set(persisted_flow_names) diff --git a/python/cocoindex/functions.py b/python/cocoindex/functions.py index 7f12ca12..8d65c9a7 100644 --- a/python/cocoindex/functions.py +++ b/python/cocoindex/functions.py @@ -1,10 +1,13 @@ """All builtin functions.""" -from typing import Annotated, Any +from typing import Annotated, Any, TYPE_CHECKING -import sentence_transformers from .typing import Float32, Vector, TypeAttr from . import op, llm +# Libraries that are heavy to import. Lazily import them later. +if TYPE_CHECKING: + import sentence_transformers + class ParseJson(op.FunctionSpec): """Parse a text into a JSON object.""" @@ -35,9 +38,10 @@ class SentenceTransformerEmbedExecutor: """Executor for SentenceTransformerEmbed.""" spec: SentenceTransformerEmbed - _model: sentence_transformers.SentenceTransformer + _model: "sentence_transformers.SentenceTransformer" def analyze(self, text): + 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() diff --git a/python/cocoindex/op.py b/python/cocoindex/op.py index bf211b01..911b6d30 100644 --- a/python/cocoindex/op.py +++ b/python/cocoindex/op.py @@ -5,10 +5,10 @@ import dataclasses import inspect -from typing import get_type_hints, Protocol, Any, Callable, Awaitable, dataclass_transform +from typing import Protocol, Any, Callable, Awaitable, dataclass_transform from enum import Enum -from .typing import encode_enriched_type +from .typing import encode_enriched_type, resolve_forward_ref from .convert import encode_engine_value, make_engine_value_decoder from . import _engine @@ -214,10 +214,11 @@ def _inner(cls: type[Executor]) -> type: """ Decorate a class to provide an executor for an op. """ - type_hints = get_type_hints(cls) + # Use `__annotations__` instead of `get_type_hints`, to avoid resolving forward references. + type_hints = cls.__annotations__ if 'spec' not in type_hints: raise TypeError("Expect a `spec` field with type hint") - spec_cls = 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, diff --git a/python/cocoindex/typing.py b/python/cocoindex/typing.py index accb1c36..a6ec351b 100644 --- a/python/cocoindex/typing.py +++ b/python/cocoindex/typing.py @@ -245,3 +245,8 @@ def encode_enriched_type(t) -> dict[str, Any] | None: return None return encode_enriched_type_info(analyze_type_info(t)) + +def resolve_forward_ref(t): + if t is str: + return eval(t) # pylint: disable=eval-used + return t