ai-pipeline-core 0.1.8__py3-none-any.whl → 0.1.11__py3-none-any.whl

ai_pipeline_core/prompt_manager.py

@@ -1,3 +1,49 @@
+"""Jinja2-based prompt template management system.
+
+@public
+
+This module provides the PromptManager class for loading and rendering
+Jinja2 templates used as prompts for language models. It implements a
+smart search strategy that looks for templates in both local and shared
+directories.
+
+Search strategy:
+    1. Local directory (same as calling module)
+    2. Local 'prompts' subdirectory
+    3. Parent 'prompts' directories (up to package boundary)
+
+Key features:
+    - Automatic template discovery
+    - Jinja2 template rendering with context
+    - Smart path resolution (.jinja2/.jinja extension handling)
+    - Clear error messages for missing templates
+
+Example:
+    >>> from ai_pipeline_core import PromptManager
+    >>>
+    >>> # Initialize at module level (not inside functions)
+    >>> pm = PromptManager(__file__)
+    >>>
+    >>> # Render a template
+    >>> prompt = pm.get(
+    ...     "analyze.jinja2",
+    ...     document=doc,
+    ...     instructions="Extract key points"
+    ... )
+
+Template organization:
+    project/
+    ├── my_module.py        # Can use local templates
+    ├── analyze.jinja2      # Local template (same directory)
+    └── prompts/           # Shared prompts directory
+        ├── summarize.jinja2
+        └── extract.jinja2
+
+Note:
+    Templates should use .jinja2 or .jinja extension.
+    The extension can be omitted when calling get().
+"""
+
 from pathlib import Path
 from typing import Any
 
@@ -5,29 +51,125 @@ import jinja2
 
 from ai_pipeline_core.logging import get_pipeline_logger
 
-from .exceptions import PromptNotFoundError, PromptRenderError
+from .exceptions import PromptError, PromptNotFoundError, PromptRenderError
 
 logger = get_pipeline_logger(__name__)
 
 
 class PromptManager:
-    """A utility to load and render prompts from a structured directory.
+    """Manages Jinja2 prompt templates with smart path resolution.
+
+    @public
+
+    PromptManager provides a convenient interface for loading and rendering
+    Jinja2 templates used as prompts for LLMs. It automatically searches for
+    templates in multiple locations, supporting both local (module-specific)
+    and shared (project-wide) templates.
 
-    Searches for 'prompts' directory in the current directory and parent directories
-    (as long as __init__.py exists in parent directories).
+    Search hierarchy:
+        1. Same directory as the calling module (for local templates)
+        2. 'prompts' subdirectory in the calling module's directory
+        3. 'prompts' directories in parent packages (up to package boundary)
+
+    Attributes:
+        search_paths: List of directories where templates are searched.
+        env: Jinja2 Environment configured for prompt rendering.
+
+    Example:
+        >>> # BEST PRACTICE: Instantiate at module scope (top level), not inside functions
+        >>> # In flow/my_flow.py
+        >>> from ai_pipeline_core import PromptManager
+        >>> pm = PromptManager(__file__)  # Module-level initialization
+        >>>
+        >>> # WRONG - Don't instantiate inside handlers or hot paths:
+        >>> # async def process():
+        >>> #     pm = PromptManager(__file__)  # NO! Creates new instance each call
+        >>>
+        >>> # Uses flow/prompts/analyze.jinja2 if it exists,
+        >>> # otherwise searches parent directories
+        >>> prompt = pm.get("analyze", context=data)
+        >>>
+        >>> # Can also use templates in same directory as module
+        >>> prompt = pm.get("local_template.jinja2")
+
+    Template format:
+        Templates use standard Jinja2 syntax:
+        ```jinja2
+        Analyze the following document:
+        {{ document.name }}
+
+        {% if instructions %}
+        Instructions: {{ instructions }}
+        {% endif %}
+        ```
+
+    Note:
+        - Autoescape is disabled for prompts (raw text output)
+        - Whitespace control is enabled (trim_blocks, lstrip_blocks)
+
+    Template Inheritance:
+        Templates support standard Jinja2 inheritance. Templates are searched
+        in order of search_paths, so templates in earlier paths override later ones.
+        Precedence (first match wins):
+        1. Same directory as module
+        2. Module's prompts/ subdirectory
+        3. Parent prompts/ directories (nearest to farthest)
+        - Templates are cached by Jinja2 for performance
     """
 
-    def __init__(self, current_dir: str, prompts_dir: str = "prompts"):
-        """Initialize PromptManager with the current file path.
+    def __init__(self, current_file: str, prompts_dir: str = "prompts"):
+        """Initialize PromptManager with smart template discovery.
+
+        @public
+
+        Sets up the Jinja2 environment with a FileSystemLoader that searches
+        multiple directories for templates. The search starts from the calling
+        module's location and extends to parent package directories.
 
         Args:
-            current_dir: The __file__ path of the calling module (required)
-            prompts_dir: Name of the prompts directory to search for (default: "prompts")
+            current_file: The __file__ path of the calling module. Must be
+                         a valid file path (not __name__). Used as the
+                         starting point for template discovery.
+            prompts_dir: Name of the prompts subdirectory to search for
+                        in each package level. Defaults to "prompts".
+                        Do not pass prompts_dir='prompts' because it is already the default.
+
+        Raises:
+            PromptError: If current_file is not a valid file path (e.g.,
+                        if __name__ was passed instead of __file__).
+
+        Note:
+            Search behavior - Given a module at /project/flows/my_flow.py:
+            1. /project/flows/ (local templates)
+            2. /project/flows/prompts/ (if exists)
+            3. /project/prompts/ (if /project has __init__.py)
+
+            Search stops when no __init__.py is found (package boundary).
+
+        Example:
+            >>> # Correct usage
+            >>> pm = PromptManager(__file__)
+            >>>
+            >>> # Custom prompts directory name
+            >>> pm = PromptManager(__file__, prompts_dir="templates")
+            >>>
+            >>> # Common mistake (will raise PromptError)
+            >>> pm = PromptManager(__name__)  # Wrong!
+
+        Note:
+            The search is limited to 4 parent levels to prevent
+            excessive filesystem traversal.
         """
         search_paths: list[Path] = []
 
         # Start from the directory containing the calling file
-        current_path = Path(current_dir).resolve()
+        current_path = Path(current_file).resolve()
+        if not current_path.exists():
+            raise PromptError(
+                f"PromptManager expected __file__ (a valid file path), "
+                f"but got {current_file!r}. Did you pass __name__ instead?"
+            )
+
         if current_path.is_file():
             current_path = current_path.parent
 
@@ -74,32 +216,81 @@ class PromptManager:
         )
 
     def get(self, prompt_path: str, **kwargs: Any) -> str:
-        """
-        Renders a specific prompt with the given context.
+        """Load and render a Jinja2 template with the given context.
+
+        @public
+
+        Searches for the template in all configured search paths and renders
+        it with the provided context variables. Automatically tries adding
+        .jinja2 or .jinja extensions if the file is not found.
 
         Args:
-            prompt_path: The path to the prompt file relative to the `prompts`
-                         directory (e.g., 'step_01_process_inputs/summarize_document.jinja2').
-                         The .jinja2 extension will be added automatically if missing.
-            **kwargs: Variables to be injected into the template.
+            prompt_path: Path to the template file, relative to any search
+                        directory. Can be a simple filename ("analyze")
+                        or include subdirectories ("tasks/summarize").
+                        Extensions (.jinja2, .jinja) are optional.
+            **kwargs: Context variables passed to the template. These become
+                     available as variables within the Jinja2 template.
 
         Returns:
-            The rendered prompt string.
+            The rendered template as a string, ready to be sent to an LLM.
+
+        Raises:
+            PromptNotFoundError: If the template file cannot be found in
+                               any search path.
+            PromptRenderError: If the template contains errors or if
+                              rendering fails (e.g., missing variables,
+                              syntax errors).
+
+        Note:
+            Template resolution - Given prompt_path="analyze":
+            1. Try "analyze" as-is
+            2. Try "analyze.jinja2"
+            3. Try "analyze.jinja"
+
+            The first matching file is used.
+
+        Example:
+            >>> pm = PromptManager(__file__)
+            >>>
+            >>> # Simple rendering
+            >>> prompt = pm.get("summarize", text="Long document...")
+            >>>
+            >>> # With complex context
+            >>> prompt = pm.get(
+            ...     "analyze",
+            ...     document=doc,
+            ...     max_length=500,
+            ...     style="technical",
+            ...     options={"include_metadata": True}
+            ... )
+            >>>
+            >>> # Nested template path
+            >>> prompt = pm.get("flows/extraction/extract_entities")
+
+        Template example:
+            ```jinja2
+            Summarize the following text in {{ max_length }} words:
+
+            {{ text }}
+
+            {% if style %}
+            Style: {{ style }}
+            {% endif %}
+            ```
+
+        Note:
+            All Jinja2 features are available: loops, conditionals,
+            filters, macros, inheritance, etc.
         """
         try:
             template = self.env.get_template(prompt_path)
             return template.render(**kwargs)
         except jinja2.TemplateNotFound:
             # If the template wasn't found and doesn't end with .jinja2, try adding the extension
-            if not prompt_path.endswith(".jinja2"):
-                try:
-                    template = self.env.get_template(prompt_path + ".jinja2")
-                    return template.render(**kwargs)
-                except jinja2.TemplateNotFound:
-                    pass  # Fall through to the original error
-            if not prompt_path.endswith(".jinja"):
+            for extension in [".jinja2", ".jinja", ".j2"]:
                 try:
-                    template = self.env.get_template(prompt_path + ".jinja")
+                    template = self.env.get_template(prompt_path + extension)
                     return template.render(**kwargs)
                 except jinja2.TemplateNotFound:
                     pass  # Fall through to the original error

ai_pipeline_core/settings.py

@@ -1,12 +1,115 @@
-"""Core configuration settings for pipeline operations."""
+"""Core configuration settings for pipeline operations.
+
+@public
+
+This module provides the Settings base class for configuration management.
+Applications should inherit from Settings to create their own ProjectSettings
+class with additional configuration fields.
+
+Environment variables:
+    OPENAI_BASE_URL: LiteLLM proxy endpoint (e.g., http://localhost:4000)
+    OPENAI_API_KEY: API key for LiteLLM proxy authentication
+    PREFECT_API_URL: Prefect server endpoint for flow orchestration
+    PREFECT_API_KEY: Prefect API authentication key
+    LMNR_PROJECT_API_KEY: Laminar project key for observability
+
+Configuration precedence:
+    1. Environment variables (highest priority)
+    2. .env file in current directory
+    3. Default values (empty strings)
+
+Example:
+    >>> from ai_pipeline_core import Settings
+    >>>
+    >>> # Create your project's settings class
+    >>> class ProjectSettings(Settings):
+    ...     app_name: str = "my-app"
+    ...     debug_mode: bool = False
+    >>>
+    >>> # Create singleton instance
+    >>> settings = ProjectSettings()
+    >>>
+    >>> # Access configuration
+    >>> print(settings.openai_base_url)
+    >>> print(settings.app_name)
+
+.env file format:
+    OPENAI_BASE_URL=http://localhost:4000
+    OPENAI_API_KEY=sk-1234567890
+    PREFECT_API_URL=http://localhost:4200/api
+    PREFECT_API_KEY=pnu_abc123
+    LMNR_PROJECT_API_KEY=lmnr_proj_xyz
+    APP_NAME=production-app
+    DEBUG_MODE=false
+
+Note:
+    Settings are loaded once at initialization and frozen. There is no
+    built-in reload mechanism - the process must be restarted to pick up
+    changes to environment variables or .env file. This is by design to
+    ensure consistency during execution.
+"""
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class Settings(BaseSettings):
-    """Core settings for pipeline operations."""
+    """Base configuration class for AI Pipeline applications.
+
+    @public
+
+    Settings is designed to be inherited by your application's configuration
+    class. It provides core AI Pipeline settings and type-safe configuration
+    management with automatic loading from environment variables and .env files.
+    All settings are immutable after initialization.
+
+    Inherit from Settings to add your application-specific configuration:
+
+        >>> from ai_pipeline_core import Settings
+        >>>
+        >>> class ProjectSettings(Settings):
+        ...     # Your custom settings
+        ...     app_name: str = "my-app"
+        ...     max_retries: int = 3
+        ...     enable_cache: bool = True
+        >>>
+        >>> # Create singleton instance for your app
+        >>> settings = ProjectSettings()
+
+    Core Attributes:
+        openai_base_url: LiteLLM proxy URL for OpenAI-compatible API.
+                        Required for all LLM operations. Usually
+                        http://localhost:4000 for local development.
+
+        openai_api_key: Authentication key for LiteLLM proxy. Required
+                       for LLM operations. Format depends on proxy config.
+
+        prefect_api_url: Prefect server API endpoint. Required for flow
+                        deployment and remote execution. Leave empty for
+                        local-only execution.
+
+        prefect_api_key: Prefect API authentication key. Required only
+                        when connecting to Prefect Cloud or secured server.
+
+        lmnr_project_api_key: Laminar (LMNR) project API key for tracing
+                              and observability. Optional but recommended
+                              for production monitoring.
+
+    Configuration sources:
+        - Environment variables (highest priority)
+        - .env file in current directory
+        - Default values in class definition
+
+    Note:
+        Empty strings are used as defaults to allow optional services.
+        Check for empty values before using service-specific settings.
+    """
 
-    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore")
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+        frozen=True,  # Settings are immutable after initialization
+    )
 
     # LLM API Configuration
     openai_base_url: str = ""
@@ -20,5 +123,6 @@ class Settings(BaseSettings):
     lmnr_project_api_key: str = ""
 
 
-# Create a single, importable instance of the settings
+# Legacy: Module-level instance for backwards compatibility
+# Applications should create their own settings instance
 settings = Settings()

ai_pipeline_core/simple_runner/__init__.py

@@ -1,3 +1,8 @@
+"""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 (
     ConfigSequence,

ai_pipeline_core/simple_runner/cli.py

@@ -1,3 +1,5 @@
+"""Command-line interface for simple pipeline execution."""
+
 from __future__ import annotations
 
 import asyncio
@@ -8,6 +10,7 @@ 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
@@ -21,10 +24,37 @@ from .simple_runner import ConfigSequence, FlowSequence, run_pipelines, save_doc
 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()
@@ -33,8 +63,24 @@ def _initialize_environment() -> None:
         logger.warning(f"Failed to initialize LMNR tracing: {e}")
 
 
-def _running_under_pytest() -> bool:  # NEW
-    """Return True when invoked by pytest (so fixtures will supply test contexts)."""
+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
 
 
@@ -46,17 +92,51 @@ def run_cli(
     initializer: InitializerFunc = None,
     trace_name: str | None = None,
 ) -> None:
-    """
-    Parse CLI+env into options, then run the pipeline.
-
-    - working_directory: required positional arg
-    - --project-name: optional, defaults to directory name
-    - --start/--end: optional, 1-based step bounds
-    - all other flags come from options_cls (fields & Field descriptions)
+    """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.simple_runner import run_cli
+        >>> from .flows import AnalysisFlow, SummaryFlow
+        >>> from .config import AnalysisConfig, AnalysisOptions
+        >>>
+        >>> if __name__ == "__main__":
+        ...     run_cli(
+        ...         flows=[AnalysisFlow, SummaryFlow],
+        ...         flow_configs=[
+        ...             (AnalysisConfig, AnalysisOptions),
+        ...             (AnalysisConfig, AnalysisOptions)
+        ...         ],
+        ...         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
+        # Add --help to show usage when run without arguments
         sys.argv.append("--help")
 
     _initialize_environment()
@@ -69,6 +149,12 @@ def run_cli(
         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
@@ -76,7 +162,49 @@ def run_cli(
 
         model_config = SettingsConfigDict(frozen=True, extra="ignore")
 
-    opts = cast(FlowOptions, _RunnerOptions())  # type: ignore[reportCallIssue]
+    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)
@@ -97,23 +225,23 @@ def run_cli(
         # 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:
             save_documents_to_directory(wd, initial_documents)
 
     # Setup context stack with optional test harness and tracing
-
     with ExitStack() as stack:
-        if not settings.prefect_api_key and not _running_under_pytest():
-            stack.enter_context(prefect_test_harness())
-            stack.enter_context(disable_run_logger())
-
         if trace_name:
             stack.enter_context(
-                Laminar.start_span(
+                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,