Split up doc collection and data extraction

compass/_cli/collect.py

@@ -0,0 +1,62 @@
+"""COMPASS CLI collect subcommand"""
+
+import click
+
+from compass._cli.common import run_async_command, OUT_DIR_POLICY_CHOICES
+from compass.plugin import create_schema_based_one_shot_extraction_plugin
+from compass.pipeline import CollectionRequest
+from compass.utilities.io import load_config
+
+
+@click.command
+@click.option(
+    "--config",
+    "-c",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to a collection configuration JSON or JSON5 file. This file "
+    "should contain any/all the arguments to pass to "
+    ":class:`~compass.pipeline.data_classes.CollectionRequest`.",
+)
+@click.option(
+    "-v",
+    "--verbose",
+    count=True,
+    help="Show logs on the terminal.",
+)
+@click.option(
+    "-np",
+    "--no-progress",
+    is_flag=True,
+    help="Flag to hide progress bars during collection.",
+)
+@click.option(
+    "--plugin",
+    "-p",
+    required=False,
+    default=None,
+    help="One-shot plugin configuration to add to COMPASS before collection",
+)
+@click.option(
+    "--out-dir-exists",
+    "-o",
+    required=False,
+    default=None,
+    type=click.Choice(OUT_DIR_POLICY_CHOICES, case_sensitive=False),
+    help="How to handle an existing output directory."
+    " Choices: fail, increment, overwrite, prompt."
+    " If omitted, prompts interactively when running in a terminal,"
+    " or fails when running non-interactively (e.g. CI).",
+)
+def collect(config, verbose, no_progress, plugin, out_dir_exists):
+    """Collect ordinance documents for a list of jurisdictions"""
+    config = load_config(config)
+
+    if plugin is not None:
+        create_schema_based_one_shot_extraction_plugin(
+            config=plugin, tech=config["tech"]
+        )
+
+    run_async_command(
+        config, CollectionRequest, verbose, no_progress, out_dir_exists
+    )

compass/_cli/common.py

@@ -1,27 +1,99 @@
 """Shared helpers for COMPASS CLI subcommands"""
 
+import sys
+import shutil
+import asyncio
 import logging
+import warnings
+import contextlib
+import multiprocessing
+from pathlib import Path
 
+import click
+from rich.console import Console
+from rich.live import Live
 from rich.logging import RichHandler
+from rich.theme import Theme
 
+from compass.pb import COMPASS_PB
 from compass.utilities.logs import AddLocationFilter
+from compass.pipeline.coordinator import run_compass
 
 
-def setup_cli_logging(console, verbosity_level, log_level="INFO"):
-    """Attach a Rich log handler to selected libraries
+OUT_DIR_POLICY_CHOICES = ["fail", "increment", "overwrite", "prompt"]
+
+
+def run_async_command(
+    config, request_class, verbose, no_progress, out_dir_exists=None
+):
+    """Run a COMPASS async command with shared CLI behavior
 
     Parameters
     ----------
-    console : rich.console.Console
-        Console instance used by the Rich log handler.
-    verbosity_level : int
-        Number of ``-v`` flags supplied on the command line. Each
-        increment opts an additional set of libraries into terminal
-        logging.
-    log_level : str, optional
-        Log level applied to each attached library logger and handler.
-        By default, ``"INFO"``.
+    config : dict
+        Configuration dictionary passed as keyword arguments to
+        `command`. This mapping must include an ``"out_dir"`` entry,
+        which is resolved according to `out_dir_exists` before command
+        execution.
+    request_class : callable
+        The COMPASS request class to instantiate and pass to the command
+        function, e.g.
+        :class:`~compass.pipeline.data_classes.CollectionRequest`.
+    verbose : int
+        CLI verbosity level controlling which library loggers are shown
+        in the console. Higher values enable logs from more underlying
+        libraries.
+    no_progress : bool
+        Option to disable the Rich live progress display. If ``True``,
+        the command is executed directly without attaching COMPASS
+        progress bars.
+    out_dir_exists : str, optional
+        Policy controlling how an existing output directory should be
+        handled. Supported values are ``"fail"``, ``"increment"``,
+        ``"overwrite"``, and ``"prompt"``. If ``None``, the policy is
+        chosen automatically based on whether the session is
+        interactive. By default, ``None``.
     """
+    custom_theme = Theme({"logging.level.trace": "rgb(94,79,162)"})
+    console = Console(theme=custom_theme)
+
+    setup_cli_logging(
+        console, verbose, log_level=config.get("log_level", "INFO")
+    )
+
+    config["out_dir"] = _resolve_out_dir_conflict(
+        config["out_dir"], out_dir_exists
+    )
+
+    with contextlib.suppress(RuntimeError):
+        multiprocessing.set_start_method("spawn")
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    request = request_class(**config)
+    if no_progress:
+        loop.run_until_complete(run_compass(request))
+        return
+
+    warnings.filterwarnings("ignore")
+
+    COMPASS_PB.console = console
+    live_display = Live(
+        COMPASS_PB.group,
+        console=console,
+        refresh_per_second=20,
+        transient=True,
+    )
+    with live_display:
+        run_msg = loop.run_until_complete(run_compass(request))
+
+    console.print(run_msg)
+    COMPASS_PB.console = None
+
+
+def setup_cli_logging(console, verbosity_level, log_level="INFO"):
+    """[NOT PUBLIC API] Setup logging for CLI"""
     libs = []
     if verbosity_level >= 1:
         libs.append("compass")
@@ -49,3 +121,98 @@ def setup_cli_logging(console, verbosity_level, log_level="INFO"):
         handler.addFilter(AddLocationFilter())
         logger.addHandler(handler)
         logger.setLevel(log_level)
+
+
+def _resolve_out_dir_conflict(out_dir, policy):
+    """Handle existing output directory using the selected policy"""
+    out_dir = Path(out_dir)
+    policy = _resolve_out_dir_policy(policy)
+
+    if not out_dir.exists() or policy == "fail":
+        return out_dir
+
+    if policy == "increment":
+        new_out_dir = _next_versioned_directory(out_dir)
+        click.echo(
+            "Output directory exists. "
+            f"Using incremented directory: {new_out_dir!s}"
+        )
+        return new_out_dir
+
+    if policy == "overwrite":
+        click.echo(f"Overwriting existing output directory: {out_dir!s}")
+        shutil.rmtree(out_dir)
+        return out_dir
+
+    if policy == "prompt":
+        return _resolve_prompt_out_dir_conflict(out_dir)
+
+    msg = (
+        f"Unknown out_dir_exists policy '{policy}'. "
+        f"Supported values: {OUT_DIR_POLICY_CHOICES}."
+    )
+    raise click.ClickException(msg)
+
+
+def _next_versioned_directory(out_dir):
+    """Create the next available output directory with versioning"""
+    idx = 2
+    max_idx = 1_000_000
+    while idx <= max_idx:
+        candidate = out_dir.parent / f"{out_dir.name}_v{idx}"
+        if not candidate.exists():
+            return candidate
+        idx += 1
+
+    msg = (
+        f"Unable to find an available versioned directory for '{out_dir!s}' "
+        f"up to suffix _v{max_idx}."
+    )
+    raise click.ClickException(msg)
+
+
+def _resolve_out_dir_policy(policy):
+    """Resolve output directory policy from explicit input
+
+    Falls back to terminal mode defaults when no policy is set.
+    """
+    if policy is not None:
+        return policy.lower()
+    if sys.stdin.isatty():
+        return "prompt"
+    return "fail"
+
+
+def _resolve_prompt_out_dir_conflict(out_dir):
+    """Handle interactive prompt flow for existing output directory"""
+    if not sys.stdin.isatty():
+        msg = (
+            "Cannot use out_dir_exists='prompt' in non-interactive mode. "
+            "Use one of: fail, increment, overwrite."
+        )
+        raise click.ClickException(msg)
+
+    create_incremented = click.confirm(
+        f"Output directory '{out_dir!s}' already exists. "
+        "Create a new incremented directory automatically?",
+        default=True,
+    )
+    if create_incremented:
+        new_out_dir = _next_versioned_directory(out_dir)
+        click.echo(f"Using incremented directory: {new_out_dir!s}")
+        return new_out_dir
+
+    overwrite = click.confirm(
+        f"Overwrite '{out_dir!s}' by deleting it and continuing?",
+        default=False,
+    )
+    if overwrite:
+        click.echo(f"Overwriting existing output directory: {out_dir!s}")
+        shutil.rmtree(out_dir)
+        return out_dir
+
+    msg = (
+        "Run cancelled. Please update out_dir in config, or rerun with "
+        "--out_dir_exists increment/overwrite."
+    )
+    raise click.ClickException(msg)

compass/_cli/extract.py

@@ -0,0 +1,62 @@
+"""COMPASS CLI extract subcommand"""
+
+import click
+
+from compass._cli.common import run_async_command, OUT_DIR_POLICY_CHOICES
+from compass.plugin import create_schema_based_one_shot_extraction_plugin
+from compass.pipeline import ExtractionRequest
+from compass.utilities.io import load_config
+
+
+@click.command
+@click.option(
+    "--config",
+    "-c",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to an extraction configuration JSON or JSON5 file. This file "
+    "should contain any/all the arguments to pass to "
+    ":class:`~compass.pipeline.data_classes.ExtractionRequest`.",
+)
+@click.option(
+    "-v",
+    "--verbose",
+    count=True,
+    help="Show logs on the terminal.",
+)
+@click.option(
+    "-np",
+    "--no-progress",
+    is_flag=True,
+    help="Flag to hide progress bars during extraction.",
+)
+@click.option(
+    "--plugin",
+    "-p",
+    required=False,
+    default=None,
+    help="One-shot plugin configuration to add to COMPASS before extraction",
+)
+@click.option(
+    "--out-dir-exists",
+    "-o",
+    required=False,
+    default=None,
+    type=click.Choice(OUT_DIR_POLICY_CHOICES, case_sensitive=False),
+    help="How to handle an existing output directory."
+    " Choices: fail, increment, overwrite, prompt."
+    " If omitted, prompts interactively when running in a terminal,"
+    " or fails when running non-interactively (e.g. CI).",
+)
+def extract(config, verbose, no_progress, plugin, out_dir_exists):
+    """Extract structured data from a saved collection manifest"""
+    config = load_config(config)
+
+    if plugin is not None:
+        create_schema_based_one_shot_extraction_plugin(
+            config=plugin, tech=config["tech"]
+        )
+
+    run_async_command(
+        config, ExtractionRequest, verbose, no_progress, out_dir_exists
+    )

compass/_cli/finalize.py

@@ -10,7 +10,7 @@
 from compass.utilities.io import load_config
 from compass.utilities.jurisdictions import Jurisdiction
 from compass.utilities.finalize import save_run_meta, doc_infos_to_db, save_db
-from compass.scripts.process import _initialize_model_params
+from compass.pipeline.coordinator import _build_models
 
 
 @click.command
@@ -21,7 +21,7 @@
     type=click.Path(exists=True),
     help="Path to COMPASS run configuration JSON or JSON5 file. This file "
     "should contain any/all the arguments to pass to "
-    ":func:`compass.scripts.process.process_jurisdictions_with_openai`. "
+    ":class:`~compass.pipeline.data_classes.ProcessRequest`. "
     "The output directory that this config points to will be finalized.",
 )
 def finalize(config):
@@ -57,7 +57,7 @@ def finalize(config):
     console = Console(theme=custom_theme)
     console.print(f"Finalizing COMPASS run in {dirs.out!s}...")
 
-    models = _initialize_model_params(config.get("model", "gpt-4o-mini"))
+    models = _build_models(config.get("model", "gpt-4o-mini"))
     start_datetime = datetime.fromtimestamp(dirs.out.stat().st_ctime)
     end_datetime = datetime.fromtimestamp(jurisdictions_fp.stat().st_mtime)
 

compass/_cli/main.py

@@ -3,6 +3,8 @@
 import click
 
 from compass import __version__
+from compass._cli.collect import collect
+from compass._cli.extract import extract
 from compass._cli.process import process
 from compass._cli.finalize import finalize
 from compass._cli.search import search
@@ -16,6 +18,8 @@ def main(ctx):
     ctx.ensure_object(dict)
 
 
+main.add_command(collect)
+main.add_command(extract)
 main.add_command(process)
 main.add_command(finalize)
 main.add_command(search)