ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/prefect.py

@@ -1,63 +0,0 @@
-"""Prefect core features for pipeline orchestration.
-
-This module provides clean re-exports of essential Prefect functionality.
-
-IMPORTANT: You should NEVER use the `task` and `flow` decorators directly
-unless it is 100% impossible to use `pipeline_task` and `pipeline_flow`.
-The standard Prefect decorators are exported here only for extremely
-limited edge cases where the pipeline decorators cannot be used.
-
-Always prefer:
-    >>> from ai_pipeline_core import pipeline_task, pipeline_flow
-    >>>
-    >>> @pipeline_task
-    >>> async def my_task(...): ...
-    >>>
-    >>> @pipeline_flow
-    >>> async def my_flow(...): ...
-
-The `task` and `flow` decorators should only be used when:
-- You absolutely cannot convert to async (pipeline decorators require async)
-- You have a very specific Prefect integration that conflicts with tracing
-- You are writing test utilities or infrastructure code
-
-Exported components:
-    task: Prefect task decorator (AVOID - use pipeline_task instead).
-    flow: Prefect flow decorator (AVOID - use pipeline_flow instead).
-    disable_run_logger: Context manager to suppress Prefect logging.
-    prefect_test_harness: Test harness for unit testing flows/tasks.
-
-Testing utilities (use as fixtures):
-    The disable_run_logger and prefect_test_harness should be used as
-    pytest fixtures as shown in tests/conftest.py:
-
-    >>> @pytest.fixture(autouse=True, scope="session")
-    >>> def prefect_test_fixture():
-    ...     with prefect_test_harness():
-    ...         yield
-    >>>
-    >>> @pytest.fixture(autouse=True)
-    >>> def disable_prefect_logging():
-    ...     with disable_run_logger():
-    ...         yield
-
-Note:
-    The pipeline_task and pipeline_flow decorators from
-    ai_pipeline_core.pipeline provide async-only execution with
-    integrated LMNR tracing and are the standard for this library.
-"""
-
-from prefect import deploy, flow, serve, task
-from prefect.logging import disable_run_logger
-from prefect.testing.utilities import prefect_test_harness
-from prefect.types.entrypoint import EntrypointType
-
-__all__ = [
-    "task",
-    "flow",
-    "disable_run_logger",
-    "prefect_test_harness",
-    "serve",
-    "deploy",
-    "EntrypointType",
-]

ai_pipeline_core/simple_runner/__init__.py

@@ -1,14 +0,0 @@
-"""Simple pipeline execution for local development.
-
-Utilities for running AI pipelines locally without full Prefect orchestration.
-"""
-
-from .cli import run_cli
-from .simple_runner import FlowSequence, run_pipeline, run_pipelines
-
-__all__ = [
-    "run_cli",
-    "run_pipeline",
-    "run_pipelines",
-    "FlowSequence",
-]

ai_pipeline_core/simple_runner/cli.py

@@ -1,254 +0,0 @@
-"""Command-line interface for simple pipeline execution."""
-
-import asyncio
-import os
-import sys
-from contextlib import ExitStack
-from pathlib import Path
-from typing import Callable, Type, TypeVar, cast
-
-from lmnr import Laminar
-from pydantic import ValidationError
-from pydantic_settings import CliPositionalArg, SettingsConfigDict
-
-from ai_pipeline_core.documents import DocumentList
-from ai_pipeline_core.flow.options import FlowOptions
-from ai_pipeline_core.logging import get_pipeline_logger, setup_logging
-from ai_pipeline_core.prefect import disable_run_logger, prefect_test_harness
-from ai_pipeline_core.settings import settings
-
-from .simple_runner import FlowSequence, run_pipelines
-
-logger = get_pipeline_logger(__name__)
-
-TOptions = TypeVar("TOptions", bound=FlowOptions)
-"""Type variable for FlowOptions subclasses used in CLI."""
-
-InitializerFunc = Callable[[FlowOptions], tuple[str, DocumentList]] | None
-"""Function type for custom pipeline initialization.
-
-Initializers can create initial documents or setup project state
-before flow execution begins.
-
-Args:
-    FlowOptions: Parsed CLI options
-
-Returns:
-    Tuple of (project_name, initial_documents) or None
-"""
-
-
-def _initialize_environment() -> None:
-    """Initialize logging and observability systems.
-
-    Sets up the pipeline logging configuration and attempts to
-    initialize LMNR (Laminar) for distributed tracing. Failures
-    in LMNR initialization are logged but don't stop execution.
-
-    Side effects:
-        - Configures Python logging system
-        - Initializes Laminar SDK if API key is available
-        - Logs initialization status
-
-    Note:
-        Called automatically by run_cli before parsing arguments.
-    """
-    setup_logging()
-    try:
-        Laminar.initialize()
-        logger.info("LMNR tracing initialized.")
-    except Exception as e:
-        logger.warning(f"Failed to initialize LMNR tracing: {e}")
-
-
-def _running_under_pytest() -> bool:
-    """Check if code is running under pytest.
-
-    Detects pytest execution context to determine whether test
-    fixtures will provide necessary contexts (like Prefect test
-    harness). This prevents duplicate context setup.
-
-    Returns:
-        True if running under pytest, False otherwise.
-
-    Detection methods:
-        - PYTEST_CURRENT_TEST environment variable (set by pytest)
-        - 'pytest' module in sys.modules (imported by test runner)
-
-    Note:
-        Used to avoid setting up test harness when pytest fixtures
-        already provide it.
-    """
-    return "PYTEST_CURRENT_TEST" in os.environ or "pytest" in sys.modules
-
-
-def run_cli(
-    *,
-    flows: FlowSequence,
-    options_cls: Type[TOptions],
-    initializer: InitializerFunc = None,
-    trace_name: str | None = None,
-) -> None:
-    """Execute pipeline flows from command-line arguments.
-
-    Environment setup:
-        - Initializes logging system
-        - Sets up LMNR tracing (if API key configured)
-        - Creates Prefect test harness (if no API key and not in pytest)
-        - Manages context stack for proper cleanup
-
-    Raises:
-        ValueError: If project name is empty after initialization.
-
-    Example:
-        >>> # In __main__.py
-        >>> from ai_pipeline_core import simple_runner
-        >>> from .flows import AnalysisFlow, SummaryFlow
-        >>> from .config import AnalysisOptions
-        >>>
-        >>> if __name__ == "__main__":
-        ...     simple_runner.run_cli(
-        ...         flows=[AnalysisFlow, SummaryFlow],
-        ...         options_cls=AnalysisOptions,
-        ...         trace_name="document-analysis"
-        ...     )
-
-        Command line:
-        $ python -m my_module ./output --temperature 0.5 --model gpt-5
-        $ python -m my_module ./output --start 2  # Skip first flow
-
-    Note:
-        - Field names are converted to kebab-case for CLI (max_tokens → --max-tokens)
-        - Boolean fields become flags (--verbose/--no-verbose)
-        - Field descriptions from Pydantic become help text
-        - Type hints are enforced during parsing
-        - Validation errors show helpful messages with field names
-        - Includes hints for common error types (numbers, ranges)
-        - Exits with status 1 on error
-        - Shows --help when no arguments provided
-    """
-    # Check if no arguments provided before initialization
-    if len(sys.argv) == 1:
-        # Add --help to show usage when run without arguments
-        sys.argv.append("--help")
-
-    _initialize_environment()
-
-    class _RunnerOptions(  # type: ignore[reportRedeclaration]
-        options_cls,
-        cli_parse_args=True,
-        cli_kebab_case=True,
-        cli_exit_on_error=True,  # Let it exit normally on error
-        cli_prog_name="ai-pipeline",
-        cli_use_class_docs_for_groups=True,
-    ):
-        """Internal options class combining user options with CLI arguments.
-
-        Dynamically created class that inherits from user's options_cls
-        and adds standard CLI arguments for pipeline execution.
-        """
-
-        working_directory: CliPositionalArg[Path]
-        project_name: str | None = None
-        start: int = 1
-        end: int | None = None
-
-        model_config = SettingsConfigDict(frozen=True, extra="ignore")
-
-    try:
-        opts = cast(FlowOptions, _RunnerOptions())  # type: ignore[reportCallIssue]
-    except ValidationError as e:
-        print("\nError: Invalid command line arguments\n", file=sys.stderr)
-        for error in e.errors():
-            field = " -> ".join(str(loc) for loc in error["loc"])
-            msg = error["msg"]
-            value = error.get("input", "")
-
-            # Format the field name nicely (convert from snake_case to kebab-case for CLI)
-            cli_field = field.replace("_", "-")
-
-            print(f"  --{cli_field}: {msg}", file=sys.stderr)
-            if value:
-                print(f"    Provided value: '{value}'", file=sys.stderr)
-
-            # Add helpful hints for common errors
-            if error["type"] == "float_parsing":
-                print("    Hint: Please provide a valid number (e.g., 0.7)", file=sys.stderr)
-            elif error["type"] == "int_parsing":
-                print("    Hint: Please provide a valid integer (e.g., 10)", file=sys.stderr)
-            elif error["type"] == "literal_error":
-                ctx = error.get("ctx", {})
-                expected = ctx.get("expected", "valid options")
-                print(f"    Hint: Valid options are: {expected}", file=sys.stderr)
-            elif error["type"] in [
-                "less_than_equal",
-                "greater_than_equal",
-                "less_than",
-                "greater_than",
-            ]:
-                ctx = error.get("ctx", {})
-                if "le" in ctx:
-                    print(f"    Hint: Value must be ≤ {ctx['le']}", file=sys.stderr)
-                elif "ge" in ctx:
-                    print(f"    Hint: Value must be ≥ {ctx['ge']}", file=sys.stderr)
-                elif "lt" in ctx:
-                    print(f"    Hint: Value must be < {ctx['lt']}", file=sys.stderr)
-                elif "gt" in ctx:
-                    print(f"    Hint: Value must be > {ctx['gt']}", file=sys.stderr)
-
-        print("\nRun with --help to see all available options\n", file=sys.stderr)
-        sys.exit(1)
-
-    wd: Path = cast(Path, getattr(opts, "working_directory"))
-    wd.mkdir(parents=True, exist_ok=True)
-
-    # Get project name from options or use directory basename
-    project_name = getattr(opts, "project_name", None)
-    if not project_name:  # None or empty string
-        project_name = wd.name
-
-    # Ensure project_name is not empty
-    if not project_name:
-        raise ValueError("Project name cannot be empty")
-
-    # Use initializer if provided, otherwise use defaults
-    initial_documents = DocumentList([])
-    if initializer:
-        init_result = initializer(opts)
-        # Always expect tuple format from initializer
-        _, initial_documents = init_result  # Ignore project name from initializer
-
-        # Save initial documents if starting from first step
-        if getattr(opts, "start", 1) == 1 and initial_documents and flows:
-            # Get config from the first flow
-            first_flow_config = getattr(flows[0], "config", None)
-            if first_flow_config:
-                asyncio.run(
-                    first_flow_config.save_documents(
-                        str(wd), initial_documents, validate_output_type=False
-                    )
-                )
-
-    # Setup context stack with optional test harness and tracing
-    with ExitStack() as stack:
-        if trace_name:
-            stack.enter_context(
-                Laminar.start_as_current_span(
-                    name=f"{trace_name}-{project_name}", input=[opts.model_dump_json()]
-                )
-            )
-
-        if not settings.prefect_api_key and not _running_under_pytest():
-            stack.enter_context(prefect_test_harness())
-            stack.enter_context(disable_run_logger())
-
-        asyncio.run(
-            run_pipelines(
-                project_name=project_name,
-                output_dir=wd,
-                flows=flows,
-                flow_options=opts,
-                start_step=getattr(opts, "start", 1),
-                end_step=getattr(opts, "end", None),
-            )
-        )

ai_pipeline_core/simple_runner/simple_runner.py

@@ -1,247 +0,0 @@
-"""Simple pipeline runner for local flow execution.
-
-This module provides the core functionality for running AI pipeline flows
-locally without full Prefect orchestration. It handles document I/O,
-flow sequencing, and error management.
-
-Key components:
-    - Document I/O from/to filesystem directories via FlowConfig
-    - Single and multi-flow execution
-    - Automatic document validation and passing between flows
-    - Step-based execution control (start/end steps)
-
-Directory structure:
-    working_dir/
-    ├── inputdocument/       # Documents of type InputDocument (lowercase)
-    │   ├── file1.txt
-    │   └── file1.txt.description.md   # Optional description
-    └── outputdocument/      # Documents of type OutputDocument (lowercase)
-        └── result.json
-
-Example:
-    >>> from ai_pipeline_core import simple_runner
-    >>>
-    >>> # Run single flow
-    >>> results = await simple_runner.run_pipeline(
-    ...     flow_func=MyFlow,
-    ...     config=MyConfig,
-    ...     project_name="test",
-    ...     output_dir=Path("./output"),
-    ...     flow_options=options
-    ... )
-
-Note:
-    Document directories are organized by document type names (lowercase)
-    for consistent structure and easy access.
-"""
-
-from pathlib import Path
-from typing import Any, Callable, Sequence
-
-from ai_pipeline_core.documents import DocumentList
-from ai_pipeline_core.flow.options import FlowOptions
-from ai_pipeline_core.logging import get_pipeline_logger
-
-logger = get_pipeline_logger(__name__)
-
-FlowSequence = Sequence[Callable[..., Any]]
-"""Type alias for a sequence of flow functions."""
-
-
-async def run_pipeline(
-    flow_func: Callable[..., Any],
-    project_name: str,
-    output_dir: Path,
-    flow_options: FlowOptions,
-    flow_name: str | None = None,
-) -> DocumentList:
-    """Execute a single pipeline flow with document I/O.
-
-    Runs a flow function with automatic document loading, validation,
-    and saving. The flow receives input documents from the filesystem
-    and saves its output for subsequent flows.
-
-    The execution proceeds through these steps:
-    1. Load input documents from output_dir subdirectories
-    2. Validate input documents against flow's config requirements
-    3. Execute flow function with documents and options
-    4. Validate output documents match config.OUTPUT_DOCUMENT_TYPE
-    5. Save output documents to output_dir subdirectories
-
-    Args:
-        flow_func: Async flow function decorated with @pipeline_flow.
-                  Must accept (project_name, documents, flow_options).
-                  The flow must have a config attribute set by @pipeline_flow.
-
-        project_name: Name of the project/pipeline for logging and tracking.
-
-        output_dir: Directory for loading input and saving output documents.
-                   Document subdirectories are created as needed.
-
-        flow_options: Configuration options passed to the flow function.
-                     Can be FlowOptions or any subclass.
-
-        flow_name: Optional display name for logging. If None, uses
-                  flow_func.name or flow_func.__name__.
-
-    Returns:
-        DocumentList containing the flow's output documents.
-
-    Raises:
-        RuntimeError: If required input documents are missing or if
-                     flow doesn't have a config attribute.
-
-    Example:
-        >>> from my_flows import AnalysisFlow
-        >>>
-        >>> results = await run_pipeline(
-        ...     flow_func=AnalysisFlow,
-        ...     project_name="analysis_001",
-        ...     output_dir=Path("./results"),
-        ...     flow_options=FlowOptions(temperature=0.7)
-        ... )
-        >>> print(f"Generated {len(results)} documents")
-
-    Note:
-        - Flow must be async (decorated with @pipeline_flow with config)
-        - Input documents are loaded based on flow's config.INPUT_DOCUMENT_TYPES
-        - Output is validated against config.OUTPUT_DOCUMENT_TYPE
-        - All I/O is logged for debugging
-    """
-    if flow_name is None:
-        # For Prefect Flow objects, use their name attribute
-        # For regular functions, fall back to __name__
-        flow_name = getattr(flow_func, "name", None) or getattr(flow_func, "__name__", "flow")
-
-    logger.info(f"Running Flow: {flow_name}")
-
-    # Get config from the flow function (attached by @pipeline_flow decorator)
-    config = getattr(flow_func, "config", None)
-    if config is None:
-        raise RuntimeError(
-            f"Flow {flow_name} does not have a config attribute. "
-            "Ensure it's decorated with @pipeline_flow(config=YourConfig)"
-        )
-
-    # Load input documents using FlowConfig's new async method
-    input_documents = await config.load_documents(str(output_dir))
-
-    if not config.has_input_documents(input_documents):
-        raise RuntimeError(f"Missing input documents for flow {flow_name}")
-
-    result_documents = await flow_func(project_name, input_documents, flow_options)
-
-    config.validate_output_documents(result_documents)
-
-    # Save output documents using FlowConfig's new async method
-    await config.save_documents(str(output_dir), result_documents)
-
-    logger.info(f"Completed Flow: {flow_name}")
-
-    return result_documents
-
-
-async def run_pipelines(
-    project_name: str,
-    output_dir: Path,
-    flows: FlowSequence,
-    flow_options: FlowOptions,
-    start_step: int = 1,
-    end_step: int | None = None,
-) -> None:
-    """Execute multiple pipeline flows in sequence.
-
-    Runs a series of flows where each flow's output becomes the input
-    for the next flow. Supports partial execution with start/end steps
-    for debugging and resuming failed pipelines.
-
-    Execution proceeds by:
-    1. Validating step indices
-    2. For each flow in range [start_step, end_step]:
-       a. Loading input documents from output_dir
-       b. Executing flow with documents
-       c. Saving output documents to output_dir
-       d. Output becomes input for next flow
-    3. Logging progress and any failures
-
-    Steps are 1-based for user convenience. Step 1 is the first flow,
-    Step N is the Nth flow. Use start_step > 1 to skip initial flows
-    and end_step < N to stop early.
-
-    Args:
-        project_name: Name of the overall pipeline/project.
-        output_dir: Directory for document I/O between flows.
-                   Shared by all flows in the sequence.
-        flows: Sequence of flow functions to execute in order.
-              Must all be async functions decorated with @pipeline_flow
-              with a config parameter.
-        flow_options: Options passed to all flows in the sequence.
-                     Individual flows can use different fields.
-        start_step: First flow to execute (1-based index).
-                   Default 1 starts from the beginning.
-        end_step: Last flow to execute (1-based index).
-                 None runs through the last flow.
-
-    Raises:
-        ValueError: If start_step or end_step are out of range.
-        RuntimeError: If any flow doesn't have a config attribute.
-
-    Example:
-        >>> # Run full pipeline
-        >>> await run_pipelines(
-        ...     project_name="analysis",
-        ...     output_dir=Path("./work"),
-        ...     flows=[ExtractFlow, AnalyzeFlow, SummarizeFlow],
-        ...     flow_options=options
-        ... )
-        >>>
-        >>> # Run only steps 2-3 (skip extraction)
-        >>> await run_pipelines(
-        ...     ...,
-        ...     start_step=2,
-        ...     end_step=3
-        ... )
-
-    Note:
-        - Each flow must be decorated with @pipeline_flow(config=...)
-        - Each flow's output must match the next flow's input types
-        - Failed flows stop the entire pipeline
-        - Progress is logged with step numbers for debugging
-        - Documents persist in output_dir between runs
-    """
-    num_steps = len(flows)
-    start_index = start_step - 1
-    end_index = (end_step if end_step is not None else num_steps) - 1
-
-    if (
-        not (0 <= start_index < num_steps)
-        or not (0 <= end_index < num_steps)
-        or start_index > end_index
-    ):
-        raise ValueError("Invalid start/end steps.")
-
-    logger.info(f"Starting pipeline '{project_name}' (Steps {start_step} to {end_index + 1})")
-
-    for i in range(start_index, end_index + 1):
-        flow_func = flows[i]
-        # For Prefect Flow objects, use their name attribute; for functions, use __name__
-        flow_name = getattr(flow_func, "name", None) or getattr(
-            flow_func, "__name__", f"flow_{i + 1}"
-        )
-
-        logger.info(f"--- [Step {i + 1}/{num_steps}] Running Flow: {flow_name} ---")
-
-        try:
-            await run_pipeline(
-                flow_func=flow_func,
-                project_name=project_name,
-                output_dir=output_dir,
-                flow_options=flow_options,
-                flow_name=f"[Step {i + 1}/{num_steps}] {flow_name}",
-            )
-
-        except Exception as e:
-            logger.error(
-                f"--- [Step {i + 1}/{num_steps}] Flow {flow_name} Failed: {e} ---", exc_info=True
-            )
-            raise

ai_pipeline_core/storage/__init__.py

@@ -1,8 +0,0 @@
-"""Storage module for ai_pipeline_core.
-
-@public
-"""
-
-from ai_pipeline_core.storage.storage import ObjectInfo, RetryPolicy, Storage
-
-__all__ = ["Storage", "ObjectInfo", "RetryPolicy"]