feat(setup-api): update and officially expose setup API

Author: georgeh0

python/cocoindex/__init__.py

@@ -10,7 +10,8 @@
 from .flow import FlowBuilder, DataScope, DataSlice, Flow, transform_flow
 from .flow import flow_def
 from .flow import EvaluateAndDumpOptions, GeneratedField
-from .flow import update_all_flows_async, FlowLiveUpdater, FlowLiveUpdaterOptions
+from .flow import FlowLiveUpdater, FlowLiveUpdaterOptions
+from .flow import update_all_flows_async, setup_all_flows, drop_all_flows
 from .lib import init, start_server, stop, main_fn
 from .llm import LlmSpec, LlmApiType
 from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
@@ -40,9 +41,11 @@
     "flow_def",
     "EvaluateAndDumpOptions",
     "GeneratedField",
-    "update_all_flows_async",
     "FlowLiveUpdater",
     "FlowLiveUpdaterOptions",
+    "update_all_flows_async",
+    "setup_all_flows",
+    "drop_all_flows",
     # Lib
     "init",
     "start_server",

python/cocoindex/cli.py

@@ -7,7 +7,7 @@
 import threading
 import types
 from types import FrameType
-from typing import Any
+from typing import Any, Iterable
 
 import click
 import watchfiles
@@ -17,7 +17,7 @@
 from rich.table import Table
 
 from . import flow, lib, setting
-from .setup import make_setup_bundle, make_drop_bundle, flow_names_with_setup
+from .setup import flow_names_with_setup
 
 # Create ServerSettings lazily upon first call, as environment variables may be loaded from files, etc.
 COCOINDEX_HOST = "https://cocoindex.io"
@@ -244,15 +244,15 @@ def show(app_flow_specifier: str, color: bool, verbose: bool) -> None:
 
 
 def _setup_flows(
+    flow_iter: Iterable[flow.Flow],
     *,
-    flow_full_names: list[str],
     force: bool,
     quiet: bool = False,
     always_show_setup: bool = False,
 ) -> None:
-    setup_bundle = make_setup_bundle(flow_full_names=flow_full_names)
+    setup_bundle = flow.make_setup_bundle(flow_iter)
     description, is_up_to_date = setup_bundle.describe()
-    if always_show_setup or not (quiet and (force or is_up_to_date)):
+    if always_show_setup or not is_up_to_date:
         click.echo(description)
     if is_up_to_date:
         if not quiet:
@@ -285,9 +285,7 @@ def setup(app_target: str, force: bool) -> None:
     """
     app_ref = _get_app_ref_from_specifier(app_target)
     _load_user_app(app_ref)
-    _setup_flows(
-        flow_full_names=flow.flow_full_names(), force=force, always_show_setup=True
-    )
+    _setup_flows(flow.flows().values(), force=force, always_show_setup=True)
 
 
 @cli.command("drop")
@@ -311,7 +309,6 @@ def drop(app_target: str | None, flow_name: tuple[str, ...], force: bool) -> Non
     2. Drop specific named flows: `cocoindex drop <APP_TARGET> [FLOW_NAME...]`
     """
     app_ref = None
-    flow_names: list[str] = []
 
     if not app_target:
         raise click.UsageError(
@@ -321,44 +318,39 @@ def drop(app_target: str | None, flow_name: tuple[str, ...], force: bool) -> Non
 
     app_ref = _get_app_ref_from_specifier(app_target)
     _load_user_app(app_ref)
+
+    flows: Iterable[flow.Flow]
     if flow_name:
-        flow_names = list(flow_name)
-        click.echo(
-            f"Preparing to drop specified flows: {', '.join(flow_names)} (in '{app_ref}').",
-            err=True,
-        )
+        flows = []
+        for name in flow_name:
+            try:
+                flows.append(flow.flow_by_name(name))
+            except KeyError:
+                click.echo(
+                    f"Warning: Failed to get flow `{name}`. Ignored.",
+                    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,
-        )
+        flows = flow.flows().values()
 
-    flow_full_names = []
-    for name in flow_names:
-        try:
-            flow_full_names.append(flow.flow_by_name(name).full_name)
-        except KeyError:
-            click.echo(
-                f"Warning: Failed to get flow `{name}`. Ignored.",
-                err=True,
-            )
+    flow_full_names = ", ".join(fl.full_name for fl in flows)
+    click.echo(
+        f"Preparing to drop specified flows: {flow_full_names} (in '{app_ref}').",
+        err=True,
+    )
 
-    if not flow_full_names:
+    if not flows:
         click.echo("No flows identified for the drop operation.")
         return
 
-    setup_bundle = make_drop_bundle(flow_full_names=flow_full_names)
+    setup_bundle = flow.make_drop_bundle(flows)
     description, is_up_to_date = setup_bundle.describe()
     click.echo(description)
     if is_up_to_date:
         click.echo("No flows need to be dropped.")
         return
     if not force and not click.confirm(
-        f"\nThis will apply changes to drop setup for: {', '.join(flow_full_names)}. Continue? [yes/N]",
+        f"\nThis will apply changes to drop setup for: {flow_full_names}. Continue? [yes/N]",
         default=False,
         show_default=False,
     ):
@@ -420,19 +412,15 @@ def update(
     if flow_name is None:
         if setup:
             _setup_flows(
-                flow_full_names=flow.flow_full_names(),
+                flow.flows().values(),
                 force=force,
                 quiet=quiet,
             )
         return flow.update_all_flows(options)
     else:
         fl = flow.flow_by_name(flow_name)
         if setup:
-            _setup_flows(
-                flow_full_names=[fl.full_name],
-                force=force,
-                quiet=quiet,
-            )
+            _setup_flows((fl,), force=force, quiet=quiet)
         with flow.FlowLiveUpdater(fl, options) as updater:
             updater.wait()
             return updater.update_stats()

python/cocoindex/flow.py

@@ -9,7 +9,6 @@
 import inspect
 import datetime
 import functools
-
 from typing import (
     Any,
     Callable,
@@ -20,6 +19,7 @@
     get_origin,
     NamedTuple,
     cast,
+    Iterable,
 )
 from threading import Lock
 from enum import Enum
@@ -34,7 +34,7 @@
 from .convert import dump_engine_object, encode_engine_value, make_engine_value_decoder
 from .typing import encode_enriched_type
 from .runtime import execution_context
-from .setup import make_setup_bundle, make_drop_bundle
+from .setup import SetupChangeBundle
 
 
 class _NameBuilder:
@@ -668,7 +668,7 @@ async def setup_async(self, report_to_stdout: bool = False) -> None:
         """
         Setup the flow. The async version.
         """
-        await make_setup_bundle(flow_full_names=[self.full_name]).apply_async(
+        await make_setup_bundle([self]).describe_and_apply_async(
             report_to_stdout=report_to_stdout
         )
 
@@ -682,7 +682,7 @@ async def drop_async(self, report_to_stdout: bool = False) -> None:
         """
         Drop the flow. The async version.
         """
-        await make_drop_bundle(flow_full_names=[self.full_name]).apply_async(
+        await make_drop_bundle([self]).describe_and_apply_async(
             report_to_stdout=report_to_stdout
         )
 
@@ -751,14 +751,6 @@ def flow_names() -> list[str]:
         return list(_flows.keys())
 
 
-def flow_full_names() -> list[str]:
-    """
-    Get the full names of all flows.
-    """
-    with _flows_lock:
-        return [fl.full_name for fl in _flows.values()]
-
-
 def flows() -> dict[str, Flow]:
     """
     Get all flows.
@@ -1004,3 +996,43 @@ def _transform_flow_wrapper(fn: Callable[..., DataSlice[T]]) -> TransformFlow[T]
         return _transform_flow
 
     return _transform_flow_wrapper
+
+
+def make_setup_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+    """
+    Make a bundle to setup flows with the given names.
+    """
+    full_names = []
+    for fl in flow_iter:
+        fl.internal_flow()
+        full_names.append(fl.full_name)
+    return SetupChangeBundle(_engine.make_setup_bundle(full_names))
+
+
+def make_drop_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+    """
+    Make a bundle to drop flows with the given names.
+    """
+    full_names = []
+    for fl in flow_iter:
+        fl.internal_flow()
+        full_names.append(fl.full_name)
+    return SetupChangeBundle(_engine.make_drop_bundle(full_names))
+
+
+def setup_all_flows(report_to_stdout: bool = False) -> None:
+    """
+    Setup all flows registered in the current process.
+    """
+    with _flows_lock:
+        flow_list = list(_flows.values())
+    make_setup_bundle(flow_list).describe_and_apply(report_to_stdout=report_to_stdout)
+
+
+def drop_all_flows(report_to_stdout: bool = False) -> None:
+    """
+    Drop all flows registered in the current process.
+    """
+    with _flows_lock:
+        flow_list = list(_flows.values())
+    make_drop_bundle(flow_list).describe_and_apply(report_to_stdout=report_to_stdout)

python/cocoindex/setup.py

@@ -2,7 +2,6 @@
 This module provides APIs to manage the setup of flows.
 """
 
-from . import flow
 from . import setting
 from . import _engine  # type: ignore
 from .runtime import execution_context
@@ -49,39 +48,28 @@ async def describe_async(self) -> tuple[str, bool]:
         """
         return await self._engine_bundle.describe_async()  # type: ignore
 
+    def describe_and_apply(self, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise.
+        """
+        execution_context.run(
+            self.describe_and_apply_async(report_to_stdout=report_to_stdout)
+        )
 
-def make_setup_bundle(*, flow_full_names: list[str]) -> SetupChangeBundle:
-    """
-    Make a bundle to setup flows with the given names.
-    """
-    flow.ensure_all_flows_built()
-    return SetupChangeBundle(_engine.make_setup_bundle(flow_full_names))
-
-
-def make_drop_bundle(*, flow_full_names: list[str]) -> SetupChangeBundle:
-    """
-    Make a bundle to drop flows with the given names.
-    """
-    flow.ensure_all_flows_built()
-    return SetupChangeBundle(_engine.make_drop_bundle(flow_full_names))
-
-
-def setup_all_flows(report_to_stdout: bool = False) -> None:
-    """
-    Setup all flows registered in the current process.
-    """
-    make_setup_bundle(flow_full_names=flow.flow_full_names()).apply(
-        report_to_stdout=report_to_stdout
-    )
-
-
-def drop_all_flows(report_to_stdout: bool = False) -> None:
-    """
-    Drop all flows registered in the current process.
-    """
-    make_drop_bundle(flow_full_names=flow.flow_full_names()).apply(
-        report_to_stdout=report_to_stdout
-    )
+    async def describe_and_apply_async(self, *, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise. Async version of `describe_and_apply`.
+        """
+        if report_to_stdout:
+            desc, is_up_to_date = await self.describe_async()
+            print("Setup status:\n")
+            print(desc)
+            if is_up_to_date:
+                print("No setup changes to apply.")
+                return
+        await self.apply_async(report_to_stdout=report_to_stdout)
 
 
 def flow_names_with_setup() -> list[str]: