cocoindex 0.1.44__cp312-cp312-macosx_11_0_arm64.whl → 0.1.45__cp312-cp312-macosx_11_0_arm64.whl

cocoindex/__init__.py

@@ -1,6 +1,7 @@
 """
 Cocoindex is a framework for building and running indexing pipelines.
 """
+
 from . import functions, query, sources, storages, cli, utils
 
 from .auth_registry import AuthEntryReference, add_auth_entry, ref_auth_entry
@@ -13,4 +14,4 @@ from .llm import LlmSpec, LlmApiType
 from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
 from .setting import DatabaseConnectionSpec, Settings, ServerSettings
 from .setting import get_app_namespace
-from .typing import Float32, Float64, LocalDateTime, OffsetDateTime, Range, Vector, Json
+from .typing import Float32, Float64, LocalDateTime, OffsetDateTime, Range, Vector, Json

cocoindex/auth_registry.py

@@ -5,21 +5,25 @@ Auth registry is used to register and reference auth entries.
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
-from . import _engine
+from . import _engine  # type: ignore
 from .convert import dump_engine_object
 
 T = TypeVar("T")
 
+
 @dataclass
 class AuthEntryReference(Generic[T]):
     """Reference an auth entry by its key."""
+
     key: str
 
+
 def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]:
     """Add an auth entry to the registry. Returns its reference."""
     _engine.add_auth_entry(key, dump_engine_object(value))
     return AuthEntryReference(key)
 
-def ref_auth_entry(key: str) -> AuthEntryReference:
+
+def ref_auth_entry(key: str) -> AuthEntryReference[T]:
     """Reference an auth entry by its key."""
-    return AuthEntryReference(key)
+    return AuthEntryReference(key)

cocoindex/cli.py

@@ -8,13 +8,16 @@ import types
 
 from dotenv import load_dotenv, find_dotenv
 from rich.console import Console
+from rich.panel import Panel
 from rich.table import Table
+from typing import Any
 
-from . import flow, lib, setting, query
+from . import flow, lib, setting
 from .setup import sync_setup, drop_setup, flow_names_with_setup, apply_setup_changes
 
 # Create ServerSettings lazily upon first call, as environment variables may be loaded from files, etc.
-COCOINDEX_HOST = 'https://cocoindex.io'
+COCOINDEX_HOST = "https://cocoindex.io"
+
 
 def _parse_app_flow_specifier(specifier: str) -> tuple[str, str | None]:
     """Parses 'module_or_path[:flow_name]' into (module_or_path, flow_name | None)."""
@@ -25,7 +28,7 @@ def _parse_app_flow_specifier(specifier: str) -> tuple[str, str | None]:
         raise click.BadParameter(
             f"Application module/path part is missing or invalid in specifier: '{specifier}'. "
             "Expected format like 'myapp.py' or 'myapp:MyFlow'.",
-            param_hint="APP_SPECIFIER"
+            param_hint="APP_SPECIFIER",
         )
 
     if len(parts) == 1:
@@ -33,7 +36,7 @@ def _parse_app_flow_specifier(specifier: str) -> tuple[str, str | None]:
 
     flow_ref_part = parts[1]
 
-    if not flow_ref_part: # Handles empty string after colon
+    if not flow_ref_part:  # Handles empty string after colon
         return app_ref, None
 
     if not flow_ref_part.isidentifier():
@@ -41,10 +44,11 @@ def _parse_app_flow_specifier(specifier: str) -> tuple[str, str | None]:
             f"Invalid format for flow name part ('{flow_ref_part}') in specifier '{specifier}'. "
             "If a colon separates the application from the flow name, the flow name should typically be "
             "a valid identifier (e.g., alphanumeric with underscores, not starting with a number).",
-            param_hint="APP_SPECIFIER"
+            param_hint="APP_SPECIFIER",
         )
     return app_ref, flow_ref_part
 
+
 def _get_app_ref_from_specifier(
     specifier: str,
 ) -> str:
@@ -59,12 +63,13 @@ def _get_app_ref_from_specifier(
             click.style(
                 f"Ignoring flow name '{flow_ref}' in '{specifier}': "
                 f"this command operates on the entire app/module '{app_ref}'.",
-                fg='yellow'
+                fg="yellow",
             ),
-            err=True
+            err=True,
         )
     return app_ref
 
+
 def _load_user_app(app_target: str) -> types.ModuleType:
     """
     Loads the user's application, which can be a file path or an installed module name.
@@ -81,7 +86,7 @@ def _load_user_app(app_target: str) -> types.ModuleType:
         app_path = os.path.abspath(app_target)
         app_dir = os.path.dirname(app_path)
         module_name = os.path.splitext(os.path.basename(app_path))[0]
-        
+
         if app_dir not in sys.path:
             sys.path.insert(0, app_dir)
         try:
@@ -89,7 +94,9 @@ def _load_user_app(app_target: str) -> types.ModuleType:
             if spec is None:
                 raise ImportError(f"Could not create spec for file: {app_path}")
             module = importlib.util.module_from_spec(spec)
-            sys.modules[spec.name] = module 
+            sys.modules[spec.name] = module
+            if spec.loader is None:
+                raise ImportError(f"Could not create loader for file: {app_path}")
             spec.loader.exec_module(module)
             return module
         except (ImportError, FileNotFoundError, PermissionError) as e:
@@ -104,19 +111,24 @@ def _load_user_app(app_target: str) -> types.ModuleType:
     except ImportError as e:
         raise click.ClickException(f"Failed to load module '{app_target}': {e}")
     except Exception as e:
-        raise click.ClickException(f"Unexpected error importing module '{app_target}': {e}")
+        raise click.ClickException(
+            f"Unexpected error importing module '{app_target}': {e}"
+        )
+
 
 @click.group()
 @click.version_option(package_name="cocoindex", message="%(prog)s version %(version)s")
 @click.option(
     "--env-file",
-    type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True),
+    type=click.Path(
+        exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True
+    ),
     help="Path to a .env file to load environment variables from. "
-         "If not provided, attempts to load '.env' from the current directory.",
+    "If not provided, attempts to load '.env' from the current directory.",
     default=None,
-    show_default=False
+    show_default=False,
 )
-def cli(env_file: str | None):
+def cli(env_file: str | None) -> None:
     """
     CLI for Cocoindex.
     """
@@ -133,9 +145,10 @@ def cli(env_file: str | None):
     except Exception as e:
         raise click.ClickException(f"Failed to initialize CocoIndex library: {e}")
 
+
 @cli.command()
 @click.argument("app_target", type=str, required=False)
-def ls(app_target: str | None):
+def ls(app_target: str | None) -> None:
     """
     List all flows.
 
@@ -145,45 +158,49 @@ def ls(app_target: str | None):
     If APP_TARGET is omitted, lists all flows that have a persisted
     setup in the backend.
     """
+    persisted_flow_names = flow_names_with_setup()
     if app_target:
         app_ref = _get_app_ref_from_specifier(app_target)
         _load_user_app(app_ref)
 
         current_flow_names = set(flow.flow_names())
-        persisted_flow_names = set(flow_names_with_setup())
 
         if not current_flow_names:
             click.echo(f"No flows are defined in '{app_ref}'.")
             return
 
         has_missing = False
+        persisted_flow_names_set = set(persisted_flow_names)
         for name in sorted(current_flow_names):
-            if name in persisted_flow_names:
+            if name in persisted_flow_names_set:
                 click.echo(name)
             else:
                 click.echo(f"{name} [+]")
                 has_missing = True
 
         if has_missing:
-            click.echo('')
-            click.echo('Notes:')
-            click.echo('  [+]: Flows present in the current process, but missing setup.')
+            click.echo("")
+            click.echo("Notes:")
+            click.echo(
+                "  [+]: Flows present in the current process, but missing setup."
+            )
 
     else:
-        persisted_flow_names = sorted(flow_names_with_setup())
-
         if not persisted_flow_names:
             click.echo("No persisted flow setups found in the backend.")
             return
 
-        for name in persisted_flow_names:
+        for name in sorted(persisted_flow_names):
             click.echo(name)
 
+
 @cli.command()
 @click.argument("app_flow_specifier", type=str)
-@click.option("--color/--no-color", default=True, help="Enable or disable colored output.")
+@click.option(
+    "--color/--no-color", default=True, help="Enable or disable colored output."
+)
 @click.option("--verbose", is_flag=True, help="Show verbose output with full details.")
-def show(app_flow_specifier: str, color: bool, verbose: bool):
+def show(app_flow_specifier: str, color: bool, verbose: bool) -> None:
     """
     Show the flow spec and schema.
 
@@ -208,7 +225,7 @@ def show(app_flow_specifier: str, color: bool, verbose: bool):
     table = Table(
         title=f"Schema for Flow: {fl.name}",
         title_style="cyan",
-        header_style="bold magenta"
+        header_style="bold magenta",
     )
     table.add_column("Field", style="cyan")
     table.add_column("Type", style="green")
@@ -217,9 +234,10 @@ def show(app_flow_specifier: str, color: bool, verbose: bool):
         table.add_row(field_name, field_type, attr_str)
     console.print(table)
 
+
 @cli.command()
 @click.argument("app_target", type=str)
-def setup(app_target: str):
+def setup(app_target: str) -> None:
     """
     Check and apply backend setup changes for flows, including the internal and target storage
     (to export).
@@ -235,19 +253,29 @@ def setup(app_target: str):
         click.echo("No changes need to be pushed.")
         return
     if not click.confirm(
-        "Changes need to be pushed. Continue? [yes/N]", default=False, show_default=False):
+        "Changes need to be pushed. Continue? [yes/N]",
+        default=False,
+        show_default=False,
+    ):
         return
     apply_setup_changes(setup_status)
 
+
 @cli.command("drop")
 @click.argument("app_target", type=str, required=False)
 @click.argument("flow_name", type=str, nargs=-1)
 @click.option(
-    "-a", "--all", "drop_all", is_flag=True, show_default=True, default=False,
+    "-a",
+    "--all",
+    "drop_all",
+    is_flag=True,
+    show_default=True,
+    default=False,
     help="Drop the backend setup for all flows with persisted setup, "
-         "even if not defined in the current process."
-         "If used, APP_TARGET and any listed flow names are ignored.")
-def drop(app_target: str | None, flow_name: tuple[str, ...], drop_all: bool):
+    "even if not defined in the current process."
+    "If used, APP_TARGET and any listed flow names are ignored.",
+)
+def drop(app_target: str | None, flow_name: tuple[str, ...], drop_all: bool) -> None:
     """
     Drop the backend setup for flows.
 
@@ -262,20 +290,29 @@ def drop(app_target: str | None, flow_name: tuple[str, ...], drop_all: bool):
 
     if drop_all:
         if app_target or flow_name:
-            click.echo("Warning: When --all is used, APP_TARGET and any individual flow names are ignored.", err=True)
+            click.echo(
+                "Warning: When --all is used, APP_TARGET and any individual flow names are ignored.",
+                err=True,
+            )
         flow_names = flow_names_with_setup()
     elif app_target:
         app_ref = _get_app_ref_from_specifier(app_target)
         _load_user_app(app_ref)
         if flow_name:
             flow_names = list(flow_name)
-            click.echo(f"Preparing to drop specified flows: {', '.join(flow_names)} (in '{app_ref}').", err=True)
+            click.echo(
+                f"Preparing to drop specified flows: {', '.join(flow_names)} (in '{app_ref}').",
+                err=True,
+            )
         else:
             flow_names = flow.flow_names()
             if not flow_names:
                 click.echo(f"No flows found defined in '{app_ref}' to drop.")
                 return
-            click.echo(f"Preparing to drop all flows defined in '{app_ref}': {', '.join(flow_names)}.", err=True)
+            click.echo(
+                f"Preparing to drop all flows defined in '{app_ref}': {', '.join(flow_names)}.",
+                err=True,
+            )
     else:
         raise click.UsageError(
             "Missing arguments. You must either provide an APP_TARGET (to target app-specific flows) "
@@ -293,20 +330,33 @@ def drop(app_target: str | None, flow_name: tuple[str, ...], drop_all: bool):
         return
     if not click.confirm(
         f"\nThis will apply changes to drop setup for: {', '.join(flow_names)}. Continue? [yes/N]",
-        default=False, show_default=False):
+        default=False,
+        show_default=False,
+    ):
         click.echo("Drop operation aborted by user.")
         return
     apply_setup_changes(setup_status)
 
+
 @cli.command()
 @click.argument("app_flow_specifier", type=str)
 @click.option(
-    "-L", "--live", is_flag=True, show_default=True, default=False,
-    help="Continuously watch changes from data sources and apply to the target index.")
+    "-L",
+    "--live",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Continuously watch changes from data sources and apply to the target index.",
+)
 @click.option(
-    "-q", "--quiet", is_flag=True, show_default=True, default=False,
-    help="Avoid printing anything to the standard output, e.g. statistics.")
-def update(app_flow_specifier: str, live: bool, quiet: bool):
+    "-q",
+    "--quiet",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Avoid printing anything to the standard output, e.g. statistics.",
+)
+def update(app_flow_specifier: str, live: bool, quiet: bool) -> Any:
     """
     Update the index to reflect the latest data from data sources.
 
@@ -324,15 +374,26 @@ def update(app_flow_specifier: str, live: bool, quiet: bool):
             updater.wait()
             return updater.update_stats()
 
+
 @cli.command()
 @click.argument("app_flow_specifier", type=str)
 @click.option(
-    "-o", "--output-dir", type=str, required=False,
-    help="The directory to dump the output to.")
+    "-o",
+    "--output-dir",
+    type=str,
+    required=False,
+    help="The directory to dump the output to.",
+)
 @click.option(
-    "--cache/--no-cache", is_flag=True, show_default=True, default=True,
-    help="Use already-cached intermediate data if available.")
-def evaluate(app_flow_specifier: str, output_dir: str | None, cache: bool = True):
+    "--cache/--no-cache",
+    is_flag=True,
+    show_default=True,
+    default=True,
+    help="Use already-cached intermediate data if available.",
+)
+def evaluate(
+    app_flow_specifier: str, output_dir: str | None, cache: bool = True
+) -> None:
     """
     Evaluate the flow and dump flow outputs to files.
 
@@ -358,32 +419,64 @@ def evaluate(app_flow_specifier: str, output_dir: str | None, cache: bool = True
     options = flow.EvaluateAndDumpOptions(output_dir=output_dir, use_cache=cache)
     fl.evaluate_and_dump(options)
 
+
 @cli.command()
 @click.argument("app_target", type=str)
 @click.option(
-    "-a", "--address", type=str,
+    "-a",
+    "--address",
+    type=str,
     help="The address to bind the server to, in the format of IP:PORT. "
-         "If unspecified, the address specified in COCOINDEX_SERVER_ADDRESS will be used.")
+    "If unspecified, the address specified in COCOINDEX_SERVER_ADDRESS will be used.",
+)
 @click.option(
-    "-c", "--cors-origin", type=str,
+    "-c",
+    "--cors-origin",
+    type=str,
     help="The origins of the clients (e.g. CocoInsight UI) to allow CORS from. "
-         "Multiple origins can be specified as a comma-separated list. "
-         "e.g. `https://cocoindex.io,http://localhost:3000`. "
-         "Origins specified in COCOINDEX_SERVER_CORS_ORIGINS will also be included.")
+    "Multiple origins can be specified as a comma-separated list. "
+    "e.g. `https://cocoindex.io,http://localhost:3000`. "
+    "Origins specified in COCOINDEX_SERVER_CORS_ORIGINS will also be included.",
+)
 @click.option(
-    "-ci", "--cors-cocoindex", is_flag=True, show_default=True, default=False,
-    help=f"Allow {COCOINDEX_HOST} to access the server.")
+    "-ci",
+    "--cors-cocoindex",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help=f"Allow {COCOINDEX_HOST} to access the server.",
+)
 @click.option(
-    "-cl", "--cors-local", type=int,
-    help="Allow http://localhost:<port> to access the server.")
+    "-cl",
+    "--cors-local",
+    type=int,
+    help="Allow http://localhost:<port> to access the server.",
+)
 @click.option(
-    "-L", "--live-update", is_flag=True, show_default=True, default=False,
-    help="Continuously watch changes from data sources and apply to the target index.")
+    "-L",
+    "--live-update",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Continuously watch changes from data sources and apply to the target index.",
+)
 @click.option(
-    "-q", "--quiet", is_flag=True, show_default=True, default=False,
-    help="Avoid printing anything to the standard output, e.g. statistics.")
-def server(app_target: str, address: str | None, live_update: bool, quiet: bool,
-           cors_origin: str | None, cors_cocoindex: bool, cors_local: int | None):
+    "-q",
+    "--quiet",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Avoid printing anything to the standard output, e.g. statistics.",
+)
+def server(
+    app_target: str,
+    address: str | None,
+    live_update: bool,
+    quiet: bool,
+    cors_origin: str | None,
+    cors_cocoindex: bool,
+    cors_local: int | None,
+) -> None:
     """
     Start a HTTP server providing REST APIs.
 
@@ -420,20 +513,46 @@ def server(app_target: str, address: str | None, live_update: bool, quiet: bool,
 
 def _flow_name(name: str | None) -> str:
     names = flow.flow_names()
-    available = ', '.join(sorted(names))
+    available = ", ".join(sorted(names))
     if name is not None:
         if name not in names:
-            raise click.BadParameter(f"Flow '{name}' not found.\nAvailable: {available if names else 'None'}")
+            raise click.BadParameter(
+                f"Flow '{name}' not found.\nAvailable: {available if names else 'None'}"
+            )
         return name
     if len(names) == 0:
         raise click.UsageError("No flows available in the loaded application.")
     elif len(names) == 1:
         return names[0]
     else:
-        raise click.UsageError(f"Multiple flows available, please specify which flow to target by appending :FlowName to the APP_TARGET.\nAvailable: {available}")
+        console = Console()
+        index = 0
+
+        while True:
+            console.clear()
+            console.print(
+                Panel.fit("Select a Flow", title_align="left", border_style="cyan")
+            )
+            for i, fname in enumerate(names):
+                console.print(
+                    f"> [bold green]{fname}[/bold green]"
+                    if i == index
+                    else f"  {fname}"
+                )
+
+            key = click.getchar()
+            if key == "\x1b[A":  # Up arrow
+                index = (index - 1) % len(names)
+            elif key == "\x1b[B":  # Down arrow
+                index = (index + 1) % len(names)
+            elif key in ("\r", "\n"):  # Enter
+                console.clear()
+                return names[index]
+
 
 def _flow_by_name(name: str | None) -> flow.Flow:
     return flow.flow_by_name(_flow_name(name))
 
+
 if __name__ == "__main__":
-    cli()
+    cli()