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

ai_pipeline_core/__init__.py

@@ -1,7 +1,23 @@
 """Pipeline Core - Shared infrastructure for AI pipelines."""
 
-from .documents import Document, DocumentList, FlowDocument, TaskDocument
-from .flow import FlowConfig
+from . import llm
+from .documents import (
+    Document,
+    DocumentList,
+    FlowDocument,
+    TaskDocument,
+    canonical_name_key,
+    sanitize_url,
+)
+from .flow import FlowConfig, FlowOptions
+from .llm import (
+    AIMessages,
+    AIMessageType,
+    ModelName,
+    ModelOptions,
+    ModelResponse,
+    StructuredModelResponse,
+)
 from .logging import (
     LoggerMixin,
     LoggingConfig,
@@ -9,28 +25,53 @@ from .logging import (
     get_pipeline_logger,
     setup_logging,
 )
-from .logging import (
-    get_pipeline_logger as get_logger,
-)
+from .logging import get_pipeline_logger as get_logger
+from .pipeline import pipeline_flow, pipeline_task
+from .prefect import disable_run_logger, prefect_test_harness
 from .prompt_manager import PromptManager
 from .settings import settings
-from .tracing import trace
+from .tracing import TraceInfo, TraceLevel, trace
 
-__version__ = "0.1.6"
+__version__ = "0.1.8"
 
 __all__ = [
-    "Document",
-    "DocumentList",
-    "FlowConfig",
-    "FlowDocument",
+    # Config/Settings
+    "settings",
+    # Logging
     "get_logger",
     "get_pipeline_logger",
     "LoggerMixin",
     "LoggingConfig",
-    "PromptManager",
-    "settings",
     "setup_logging",
     "StructuredLoggerMixin",
+    # Documents
+    "Document",
+    "DocumentList",
+    "FlowDocument",
     "TaskDocument",
+    "canonical_name_key",
+    "sanitize_url",
+    # Flow/Task
+    "FlowConfig",
+    "FlowOptions",
+    # Pipeline decorators (with tracing)
+    "pipeline_task",
+    "pipeline_flow",
+    # Prefect decorators (clean, no tracing)
+    "prefect_test_harness",
+    "disable_run_logger",
+    # LLM
+    "llm",
+    "ModelName",
+    "ModelOptions",
+    "ModelResponse",
+    "StructuredModelResponse",
+    "AIMessages",
+    "AIMessageType",
+    # Tracing
     "trace",
+    "TraceLevel",
+    "TraceInfo",
+    # Utils
+    "PromptManager",
 ]

ai_pipeline_core/documents/__init__.py

@@ -2,10 +2,13 @@ from .document import Document
 from .document_list import DocumentList
 from .flow_document import FlowDocument
 from .task_document import TaskDocument
+from .utils import canonical_name_key, sanitize_url
 
 __all__ = [
     "Document",
     "DocumentList",
     "FlowDocument",
     "TaskDocument",
+    "canonical_name_key",
+    "sanitize_url",
 ]

ai_pipeline_core/documents/document.py

@@ -26,12 +26,27 @@ TModel = TypeVar("TModel", bound=BaseModel)
 
 
 class Document(BaseModel, ABC):
-    """Abstract base class for all documents"""
+    """Abstract base class for all documents.
+
+    Warning: Document subclasses should NOT start with 'Test' prefix as this
+    causes conflicts with pytest test discovery. Classes with 'Test' prefix
+    will be rejected at definition time.
+    """
 
     MAX_CONTENT_SIZE: ClassVar[int] = 25 * 1024 * 1024  # 25MB default
     DESCRIPTION_EXTENSION: ClassVar[str] = ".description.md"
     MARKDOWN_LIST_SEPARATOR: ClassVar[str] = "\n\n---\n\n"
 
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Validate subclass names to prevent pytest conflicts."""
+        super().__init_subclass__(**kwargs)
+        if cls.__name__.startswith("Test"):
+            raise TypeError(
+                f"Document subclass '{cls.__name__}' cannot start with 'Test' prefix. "
+                "This causes conflicts with pytest test discovery. "
+                "Please use a different name (e.g., 'SampleDocument', 'ExampleDocument')."
+            )
+
     def __init__(self, **data: Any) -> None:
         """Prevent direct instantiation of abstract Document class."""
         if type(self) is Document:

ai_pipeline_core/flow/__init__.py

@@ -1,3 +1,7 @@
 from .config import FlowConfig
+from .options import FlowOptions
 
-__all__ = ["FlowConfig"]
+__all__ = [
+    "FlowConfig",
+    "FlowOptions",
+]

ai_pipeline_core/flow/config.py

@@ -14,6 +14,27 @@ class FlowConfig(ABC):
     INPUT_DOCUMENT_TYPES: ClassVar[list[type[FlowDocument]]]
     OUTPUT_DOCUMENT_TYPE: ClassVar[type[FlowDocument]]
 
+    def __init_subclass__(cls, **kwargs):
+        """Validate that OUTPUT_DOCUMENT_TYPE is not in INPUT_DOCUMENT_TYPES."""
+        super().__init_subclass__(**kwargs)
+
+        # Skip validation for the abstract base class itself
+        if cls.__name__ == "FlowConfig":
+            return
+
+        # Ensure required attributes are defined
+        if not hasattr(cls, "INPUT_DOCUMENT_TYPES"):
+            raise TypeError(f"FlowConfig {cls.__name__} must define INPUT_DOCUMENT_TYPES")
+        if not hasattr(cls, "OUTPUT_DOCUMENT_TYPE"):
+            raise TypeError(f"FlowConfig {cls.__name__} must define OUTPUT_DOCUMENT_TYPE")
+
+        # Validate that output type is not in input types
+        if cls.OUTPUT_DOCUMENT_TYPE in cls.INPUT_DOCUMENT_TYPES:
+            raise TypeError(
+                f"FlowConfig {cls.__name__}: OUTPUT_DOCUMENT_TYPE "
+                f"({cls.OUTPUT_DOCUMENT_TYPE.__name__}) cannot be in INPUT_DOCUMENT_TYPES"
+            )
+
     @classmethod
     def get_input_document_types(cls) -> list[type[FlowDocument]]:
         """

ai_pipeline_core/flow/options.py

@@ -0,0 +1,26 @@
+from typing import TypeVar
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from ai_pipeline_core.llm import ModelName
+
+T = TypeVar("T", bound="FlowOptions")
+
+
+class FlowOptions(BaseSettings):
+    """Base configuration for AI Pipeline flows."""
+
+    core_model: ModelName | str = Field(
+        default="gpt-5",
+        description="Primary model for complex analysis and generation tasks.",
+    )
+    small_model: ModelName | str = Field(
+        default="gpt-5-mini",
+        description="Fast, cost-effective model for simple tasks and orchestration.",
+    )
+
+    model_config = SettingsConfigDict(frozen=True, extra="ignore")
+
+
+__all__ = ["FlowOptions"]

ai_pipeline_core/llm/client.py

@@ -118,11 +118,13 @@ async def _generate_with_retry(
                 span.set_attributes(response.get_laminar_metadata())
                 Laminar.set_span_output(response.content)
                 if not response.content:
-                    # disable cache in case of empty response
-                    completion_kwargs["extra_body"]["cache"] = {"no-cache": True}
                     raise ValueError(f"Model {model} returned an empty response.")
                 return response
         except (asyncio.TimeoutError, ValueError, Exception) as e:
+            if not isinstance(e, asyncio.TimeoutError):
+                # disable cache if it's not a timeout because it may cause an error
+                completion_kwargs["extra_body"]["cache"] = {"no-cache": True}
+
             logger.warning(
                 "LLM generation failed (attempt %d/%d): %s",
                 attempt + 1,
@@ -167,7 +169,7 @@ T = TypeVar("T", bound=BaseModel)
 
 @trace(ignore_inputs=["context"])
 async def generate_structured(
-    model: ModelName,
+    model: ModelName | str,
     response_format: type[T],
     *,
     context: AIMessages = AIMessages(),

ai_pipeline_core/llm/model_options.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 
 
 class ModelOptions(BaseModel):
+    temperature: float | None = None
     system_prompt: str | None = None
     search_context_size: Literal["low", "medium", "high"] | None = None
     reasoning_effort: Literal["low", "medium", "high"] | None = None
@@ -21,6 +22,9 @@ class ModelOptions(BaseModel):
             "extra_body": {},
         }
 
+        if self.temperature:
+            kwargs["temperature"] = self.temperature
+
         if self.max_completion_tokens:
             kwargs["max_completion_tokens"] = self.max_completion_tokens
 

ai_pipeline_core/pipeline.py

@@ -0,0 +1,414 @@
+"""
+ai_pipeline_core.pipeline
+=========================
+
+Tiny wrappers around Prefect's public ``@task`` and ``@flow`` that add our
+``trace`` decorator and **require async functions**.
+
+Why this exists
+---------------
+Prefect tasks/flows are awaitable at runtime, but their public type stubs
+don’t declare that clearly. We therefore:
+
+1) Return the **real Prefect objects** (so you keep every Prefect method).
+2) Type them as small Protocols that say “this is awaitable and has common
+   helpers like `.submit`/`.map`”.
+
+This keeps Pyright happy without altering runtime behavior and avoids
+leaking advanced typing constructs (like ``ParamSpec``) that confuse tools
+that introspect callables (e.g., Pydantic).
+
+Quick start
+-----------
+from ai_pipeline_core.pipeline import pipeline_task, pipeline_flow
+from ai_pipeline_core.documents import DocumentList
+from ai_pipeline_core.flow.options import FlowOptions
+
+@pipeline_task
+async def add(x: int, y: int) -> int:
+    return x + y
+
+@pipeline_flow
+async def my_flow(project_name: str, docs: DocumentList, opts: FlowOptions) -> DocumentList:
+    await add(1, 2)  # awaitable and typed
+    return docs
+
+Rules
+-----
+• Your decorated function **must** be ``async def``.
+• ``@pipeline_flow`` functions must accept at least:
+  (project_name: str, documents: DocumentList, flow_options: FlowOptions | subclass).
+• Both wrappers return the same Prefect objects you’d get from Prefect directly.
+"""
+
+from __future__ import annotations
+
+import datetime
+import inspect
+from typing import Any, Callable, Coroutine, Iterable, Protocol, TypeVar, Union, cast, overload
+
+from prefect.assets import Asset
+from prefect.cache_policies import CachePolicy
+from prefect.context import TaskRunContext
+from prefect.flows import FlowStateHook
+from prefect.flows import flow as _prefect_flow  # public import
+from prefect.futures import PrefectFuture
+from prefect.results import ResultSerializer, ResultStorage
+from prefect.task_runners import TaskRunner
+from prefect.tasks import task as _prefect_task  # public import
+from prefect.utilities.annotations import NotSet
+from typing_extensions import TypeAlias
+
+from ai_pipeline_core.documents import DocumentList
+from ai_pipeline_core.flow.options import FlowOptions
+from ai_pipeline_core.tracing import TraceLevel, trace
+
+# --------------------------------------------------------------------------- #
+# Public callback aliases (Prefect stubs omit these exact types)
+# --------------------------------------------------------------------------- #
+RetryConditionCallable: TypeAlias = Callable[[Any, Any, Any], bool]
+StateHookCallable: TypeAlias = Callable[[Any, Any, Any], None]
+TaskRunNameValueOrCallable: TypeAlias = Union[str, Callable[[], str]]
+
+# --------------------------------------------------------------------------- #
+# Typing helpers
+# --------------------------------------------------------------------------- #
+R_co = TypeVar("R_co", covariant=True)
+FO_contra = TypeVar("FO_contra", bound=FlowOptions, contravariant=True)
+"""Flow options are an *input* type, so contravariant fits the callable model."""
+
+
+class _TaskLike(Protocol[R_co]):
+    """Minimal 'task-like' view: awaitable call + common helpers."""
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Coroutine[Any, Any, R_co]: ...
+
+    submit: Callable[..., Any]
+    map: Callable[..., Any]
+    name: str | None
+
+    def __getattr__(self, name: str) -> Any: ...  # allow unknown helpers without type errors
+
+
+class _DocumentsFlowCallable(Protocol[FO_contra]):
+    """User async flow signature (first three params fixed)."""
+
+    def __call__(
+        self,
+        project_name: str,
+        documents: DocumentList,
+        flow_options: FO_contra,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Coroutine[Any, Any, DocumentList]: ...
+
+
+class _FlowLike(Protocol[FO_contra]):
+    """Callable returned by Prefect ``@flow`` wrapper that we expose to users."""
+
+    def __call__(
+        self,
+        project_name: str,
+        documents: DocumentList,
+        flow_options: FO_contra,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Coroutine[Any, Any, DocumentList]: ...
+
+    name: str | None
+
+    def __getattr__(self, name: str) -> Any: ...  # allow unknown helpers without type errors
+
+
+# --------------------------------------------------------------------------- #
+# Small helper: safely get a callable's name without upsetting the type checker
+# --------------------------------------------------------------------------- #
+def _callable_name(obj: Any, fallback: str) -> str:
+    try:
+        n = getattr(obj, "__name__", None)
+        return n if isinstance(n, str) else fallback
+    except Exception:
+        return fallback
+
+
+# --------------------------------------------------------------------------- #
+# @pipeline_task — async-only, traced, returns Prefect's Task object
+# --------------------------------------------------------------------------- #
+@overload
+def pipeline_task(__fn: Callable[..., Coroutine[Any, Any, R_co]], /) -> _TaskLike[R_co]: ...
+@overload
+def pipeline_task(
+    *,
+    # tracing
+    trace_level: TraceLevel = "always",
+    trace_ignore_input: bool = False,
+    trace_ignore_output: bool = False,
+    trace_ignore_inputs: list[str] | None = None,
+    trace_input_formatter: Callable[..., str] | None = None,
+    trace_output_formatter: Callable[..., str] | None = None,
+    # prefect passthrough
+    name: str | None = None,
+    description: str | None = None,
+    tags: Iterable[str] | None = None,
+    version: str | None = None,
+    cache_policy: CachePolicy | type[NotSet] = NotSet,
+    cache_key_fn: Callable[[TaskRunContext, dict[str, Any]], str | None] | None = None,
+    cache_expiration: datetime.timedelta | None = None,
+    task_run_name: TaskRunNameValueOrCallable | None = None,
+    retries: int | None = None,
+    retry_delay_seconds: int | float | list[float] | Callable[[int], list[float]] | None = None,
+    retry_jitter_factor: float | None = None,
+    persist_result: bool | None = None,
+    result_storage: ResultStorage | str | None = None,
+    result_serializer: ResultSerializer | str | None = None,
+    result_storage_key: str | None = None,
+    cache_result_in_memory: bool = True,
+    timeout_seconds: int | float | None = None,
+    log_prints: bool | None = False,
+    refresh_cache: bool | None = None,
+    on_completion: list[StateHookCallable] | None = None,
+    on_failure: list[StateHookCallable] | None = None,
+    retry_condition_fn: RetryConditionCallable | None = None,
+    viz_return_value: bool | None = None,
+    asset_deps: list[str | Asset] | None = None,
+) -> Callable[[Callable[..., Coroutine[Any, Any, R_co]]], _TaskLike[R_co]]: ...
+
+
+def pipeline_task(
+    __fn: Callable[..., Coroutine[Any, Any, R_co]] | None = None,
+    /,
+    *,
+    # tracing
+    trace_level: TraceLevel = "always",
+    trace_ignore_input: bool = False,
+    trace_ignore_output: bool = False,
+    trace_ignore_inputs: list[str] | None = None,
+    trace_input_formatter: Callable[..., str] | None = None,
+    trace_output_formatter: Callable[..., str] | None = None,
+    # prefect passthrough
+    name: str | None = None,
+    description: str | None = None,
+    tags: Iterable[str] | None = None,
+    version: str | None = None,
+    cache_policy: CachePolicy | type[NotSet] = NotSet,
+    cache_key_fn: Callable[[TaskRunContext, dict[str, Any]], str | None] | None = None,
+    cache_expiration: datetime.timedelta | None = None,
+    task_run_name: TaskRunNameValueOrCallable | None = None,
+    retries: int | None = None,
+    retry_delay_seconds: int | float | list[float] | Callable[[int], list[float]] | None = None,
+    retry_jitter_factor: float | None = None,
+    persist_result: bool | None = None,
+    result_storage: ResultStorage | str | None = None,
+    result_serializer: ResultSerializer | str | None = None,
+    result_storage_key: str | None = None,
+    cache_result_in_memory: bool = True,
+    timeout_seconds: int | float | None = None,
+    log_prints: bool | None = False,
+    refresh_cache: bool | None = None,
+    on_completion: list[StateHookCallable] | None = None,
+    on_failure: list[StateHookCallable] | None = None,
+    retry_condition_fn: RetryConditionCallable | None = None,
+    viz_return_value: bool | None = None,
+    asset_deps: list[str | Asset] | None = None,
+) -> _TaskLike[R_co] | Callable[[Callable[..., Coroutine[Any, Any, R_co]]], _TaskLike[R_co]]:
+    """Decorate an **async** function as a traced Prefect task."""
+    task_decorator: Callable[..., Any] = _prefect_task  # helps the type checker
+
+    def _apply(fn: Callable[..., Coroutine[Any, Any, R_co]]) -> _TaskLike[R_co]:
+        if not inspect.iscoroutinefunction(fn):
+            raise TypeError(
+                f"@pipeline_task target '{_callable_name(fn, 'task')}' must be 'async def'"
+            )
+
+        traced_fn = trace(
+            level=trace_level,
+            name=name or _callable_name(fn, "task"),
+            ignore_input=trace_ignore_input,
+            ignore_output=trace_ignore_output,
+            ignore_inputs=trace_ignore_inputs,
+            input_formatter=trace_input_formatter,
+            output_formatter=trace_output_formatter,
+        )(fn)
+
+        return cast(
+            _TaskLike[R_co],
+            task_decorator(
+                name=name,
+                description=description,
+                tags=tags,
+                version=version,
+                cache_policy=cache_policy,
+                cache_key_fn=cache_key_fn,
+                cache_expiration=cache_expiration,
+                task_run_name=task_run_name,
+                retries=0 if retries is None else retries,
+                retry_delay_seconds=retry_delay_seconds,
+                retry_jitter_factor=retry_jitter_factor,
+                persist_result=persist_result,
+                result_storage=result_storage,
+                result_serializer=result_serializer,
+                result_storage_key=result_storage_key,
+                cache_result_in_memory=cache_result_in_memory,
+                timeout_seconds=timeout_seconds,
+                log_prints=log_prints,
+                refresh_cache=refresh_cache,
+                on_completion=on_completion,
+                on_failure=on_failure,
+                retry_condition_fn=retry_condition_fn,
+                viz_return_value=viz_return_value,
+                asset_deps=asset_deps,
+            )(traced_fn),
+        )
+
+    return _apply(__fn) if __fn else _apply
+
+
+# --------------------------------------------------------------------------- #
+# @pipeline_flow — async-only, traced, returns Prefect’s flow wrapper
+# --------------------------------------------------------------------------- #
+@overload
+def pipeline_flow(__fn: _DocumentsFlowCallable[FO_contra], /) -> _FlowLike[FO_contra]: ...
+@overload
+def pipeline_flow(
+    *,
+    # tracing
+    trace_level: TraceLevel = "always",
+    trace_ignore_input: bool = False,
+    trace_ignore_output: bool = False,
+    trace_ignore_inputs: list[str] | None = None,
+    trace_input_formatter: Callable[..., str] | None = None,
+    trace_output_formatter: Callable[..., str] | None = None,
+    # prefect passthrough
+    name: str | None = None,
+    version: str | None = None,
+    flow_run_name: Union[Callable[[], str], str] | None = None,
+    retries: int | None = None,
+    retry_delay_seconds: int | float | None = None,
+    task_runner: TaskRunner[PrefectFuture[Any]] | None = None,
+    description: str | None = None,
+    timeout_seconds: int | float | None = None,
+    validate_parameters: bool = True,
+    persist_result: bool | None = None,
+    result_storage: ResultStorage | str | None = None,
+    result_serializer: ResultSerializer | str | None = None,
+    cache_result_in_memory: bool = True,
+    log_prints: bool | None = None,
+    on_completion: list[FlowStateHook[Any, Any]] | None = None,
+    on_failure: list[FlowStateHook[Any, Any]] | None = None,
+    on_cancellation: list[FlowStateHook[Any, Any]] | None = None,
+    on_crashed: list[FlowStateHook[Any, Any]] | None = None,
+    on_running: list[FlowStateHook[Any, Any]] | None = None,
+) -> Callable[[_DocumentsFlowCallable[FO_contra]], _FlowLike[FO_contra]]: ...
+
+
+def pipeline_flow(
+    __fn: _DocumentsFlowCallable[FO_contra] | None = None,
+    /,
+    *,
+    # tracing
+    trace_level: TraceLevel = "always",
+    trace_ignore_input: bool = False,
+    trace_ignore_output: bool = False,
+    trace_ignore_inputs: list[str] | None = None,
+    trace_input_formatter: Callable[..., str] | None = None,
+    trace_output_formatter: Callable[..., str] | None = None,
+    # prefect passthrough
+    name: str | None = None,
+    version: str | None = None,
+    flow_run_name: Union[Callable[[], str], str] | None = None,
+    retries: int | None = None,
+    retry_delay_seconds: int | float | None = None,
+    task_runner: TaskRunner[PrefectFuture[Any]] | None = None,
+    description: str | None = None,
+    timeout_seconds: int | float | None = None,
+    validate_parameters: bool = True,
+    persist_result: bool | None = None,
+    result_storage: ResultStorage | str | None = None,
+    result_serializer: ResultSerializer | str | None = None,
+    cache_result_in_memory: bool = True,
+    log_prints: bool | None = None,
+    on_completion: list[FlowStateHook[Any, Any]] | None = None,
+    on_failure: list[FlowStateHook[Any, Any]] | None = None,
+    on_cancellation: list[FlowStateHook[Any, Any]] | None = None,
+    on_crashed: list[FlowStateHook[Any, Any]] | None = None,
+    on_running: list[FlowStateHook[Any, Any]] | None = None,
+) -> _FlowLike[FO_contra] | Callable[[_DocumentsFlowCallable[FO_contra]], _FlowLike[FO_contra]]:
+    """Decorate an **async** flow.
+
+    Required signature:
+        async def flow_fn(
+            project_name: str,
+            documents: DocumentList,
+            flow_options: FlowOptions,  # or any subclass
+            *args,
+            **kwargs
+        ) -> DocumentList
+
+    Returns the same callable object Prefect’s ``@flow`` would return.
+    """
+    flow_decorator: Callable[..., Any] = _prefect_flow
+
+    def _apply(fn: _DocumentsFlowCallable[FO_contra]) -> _FlowLike[FO_contra]:
+        fname = _callable_name(fn, "flow")
+
+        if not inspect.iscoroutinefunction(fn):
+            raise TypeError(f"@pipeline_flow '{fname}' must be declared with 'async def'")
+        if len(inspect.signature(fn).parameters) < 3:
+            raise TypeError(
+                f"@pipeline_flow '{fname}' must accept "
+                "'project_name, documents, flow_options' as its first three parameters"
+            )
+
+        async def _wrapper(
+            project_name: str,
+            documents: DocumentList,
+            flow_options: FO_contra,
+            *args: Any,
+            **kwargs: Any,
+        ) -> DocumentList:
+            result = await fn(project_name, documents, flow_options, *args, **kwargs)
+            if not isinstance(result, DocumentList):  # pyright: ignore[reportUnnecessaryIsInstance]
+                raise TypeError(
+                    f"Flow '{fname}' must return DocumentList, got {type(result).__name__}"
+                )
+            return result
+
+        traced = trace(
+            level=trace_level,
+            name=name or fname,
+            ignore_input=trace_ignore_input,
+            ignore_output=trace_ignore_output,
+            ignore_inputs=trace_ignore_inputs,
+            input_formatter=trace_input_formatter,
+            output_formatter=trace_output_formatter,
+        )(_wrapper)
+
+        return cast(
+            _FlowLike[FO_contra],
+            flow_decorator(
+                name=name,
+                version=version,
+                flow_run_name=flow_run_name,
+                retries=0 if retries is None else retries,
+                retry_delay_seconds=retry_delay_seconds,
+                task_runner=task_runner,
+                description=description,
+                timeout_seconds=timeout_seconds,
+                validate_parameters=validate_parameters,
+                persist_result=persist_result,
+                result_storage=result_storage,
+                result_serializer=result_serializer,
+                cache_result_in_memory=cache_result_in_memory,
+                log_prints=log_prints,
+                on_completion=on_completion,
+                on_failure=on_failure,
+                on_cancellation=on_cancellation,
+                on_crashed=on_crashed,
+                on_running=on_running,
+            )(traced),
+        )
+
+    return _apply(__fn) if __fn else _apply
+
+
+__all__ = ["pipeline_task", "pipeline_flow"]

ai_pipeline_core/prefect.py

@@ -0,0 +1,7 @@
+"""Prefect core features."""
+
+from prefect import flow, task
+from prefect.logging import disable_run_logger
+from prefect.testing.utilities import prefect_test_harness
+
+__all__ = ["task", "flow", "disable_run_logger", "prefect_test_harness"]

ai_pipeline_core/simple_runner/__init__.py

@@ -0,0 +1,19 @@
+from .cli import run_cli
+from .simple_runner import (
+    ConfigSequence,
+    FlowSequence,
+    load_documents_from_directory,
+    run_pipeline,
+    run_pipelines,
+    save_documents_to_directory,
+)
+
+__all__ = [
+    "run_cli",
+    "run_pipeline",
+    "run_pipelines",
+    "load_documents_from_directory",
+    "save_documents_to_directory",
+    "FlowSequence",
+    "ConfigSequence",
+]