ai-pipeline-core 0.1.5__tar.gz → 0.1.7__tar.gz

{ai_pipeline_core-0.1.5 → ai_pipeline_core-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.1.5
+Version: 0.1.7
 Summary: Core utilities for AI-powered processing pipelines using prefect
 Project-URL: Homepage, https://github.com/bbarwik/ai-pipeline-core
 Project-URL: Repository, https://github.com/bbarwik/ai-pipeline-core
@@ -20,7 +20,7 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.12
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: jinja2>=3.1.6
-Requires-Dist: lmnr>=0.7.5
+Requires-Dist: lmnr>=0.7.6
 Requires-Dist: openai>=1.99.9
 Requires-Dist: prefect>=3.4.13
 Requires-Dist: pydantic-settings>=2.10.1
@@ -151,40 +151,76 @@ async def process_document(doc: Document):
     return response.parsed
 ```
 
-### Prefect Flow Integration
+### Enhanced Pipeline Decorators (New in v0.1.7)
 ```python
-from prefect import flow, task
-from ai_pipeline_core.documents import Document, DocumentList, FlowDocument
-from ai_pipeline_core.flow import FlowConfig
-from ai_pipeline_core.tracing import trace
+from ai_pipeline_core import pipeline_flow, pipeline_task
+from ai_pipeline_core.flow import FlowOptions
+from ai_pipeline_core.documents import DocumentList, FlowDocument
 
-class OutputDocument(FlowDocument):
-    """Custom output document type"""
-    def get_type(self) -> str:
-        return "output"
+class CustomFlowOptions(FlowOptions):
+    """Extend base options with your custom fields"""
+    batch_size: int = 100
+    temperature: float = 0.7
 
-class MyFlowConfig(FlowConfig):
-    INPUT_DOCUMENT_TYPES = [InputDocument]
-    OUTPUT_DOCUMENT_TYPE = OutputDocument
-
-@task
-@trace
+@pipeline_task(trace_level="always", retries=3)
 async def process_task(doc: Document) -> Document:
-    # Task-level processing with automatic tracing
+    # Task with automatic tracing and retries
     result = await process_document(doc)
-    # Convert result to JSON string for document content
-    import json
-    return OutputDocument(name="result", content=json.dumps(result.model_dump()).encode())
+    return OutputDocument(name="result", content=result.encode())
+
+@pipeline_flow(trace_level="always")
+async def my_pipeline(
+    project_name: str,
+    documents: DocumentList,
+    flow_options: CustomFlowOptions  # Type-safe custom options
+) -> DocumentList:
+    # Pipeline flow with enforced signature and tracing
+    results = []
+    for doc in documents:
+        result = await process_task(doc)
+        results.append(result)
+    return DocumentList(results)
+```
 
-@flow
-async def my_pipeline(documents: DocumentList):
-    config = MyFlowConfig()
-    input_docs = config.get_input_documents(documents)
+### Simple Runner Utility (New in v0.1.7)
+```python
+from ai_pipeline_core.simple_runner import run_cli, run_pipeline
+from ai_pipeline_core.flow import FlowOptions
+
+# CLI-based pipeline execution
+if __name__ == "__main__":
+    run_cli(
+        flows=[my_pipeline],
+        flow_configs=[MyFlowConfig],
+        options_cls=CustomFlowOptions
+    )
 
-    results = await process_task.map(input_docs)
+# Or programmatic execution
+async def main():
+    result = await run_pipeline(
+        project_name="my-project",
+        output_dir=Path("./output"),
+        flow=my_pipeline,
+        flow_config=MyFlowConfig,
+        flow_options=CustomFlowOptions(batch_size=50)
+    )
+```
 
-    config.validate_output_documents(results)
-    return results
+### Clean Prefect Decorators (New in v0.1.7)
+```python
+# Import clean Prefect decorators without tracing
+from ai_pipeline_core.prefect import flow, task
+
+# Or use pipeline decorators with tracing
+from ai_pipeline_core import pipeline_flow, pipeline_task
+
+@task  # Clean Prefect task
+def compute(x: int) -> int:
+    return x * 2
+
+@pipeline_task(trace_level="always")  # With tracing
+def compute_traced(x: int) -> int:
+    return x * 2
 ```
 
 ## Core Modules
@@ -291,8 +327,14 @@ ai_pipeline_core/
 │   ├── client.py      # Async client implementation
 │   └── model_options.py # Configuration models
 ├── flow/              # Prefect flow utilities
-│   └── config.py      # Type-safe flow configuration
+│   ├── config.py      # Type-safe flow configuration
+│   └── options.py     # FlowOptions base class (v0.1.7)
+├── simple_runner/     # Pipeline execution utilities (v0.1.7)
+│   ├── cli.py         # CLI interface
+│   └── simple_runner.py # Core runner logic
 ├── logging/           # Structured logging
+├── pipeline.py        # Enhanced decorators (v0.1.7)
+├── prefect.py         # Clean Prefect exports (v0.1.7)
 ├── tracing.py         # Observability decorators
 └── settings.py        # Centralized configuration
 ```
@@ -469,9 +511,29 @@ Built with:
 - [LiteLLM](https://litellm.ai/) - LLM proxy
 - [Pydantic](https://pydantic-docs.helpmanual.io/) - Data validation
 
+## What's New in v0.1.7
+
+### Major Additions
+- **Enhanced Pipeline Decorators**: New `pipeline_flow` and `pipeline_task` decorators combining Prefect functionality with automatic LMNR tracing
+- **FlowOptions Base Class**: Extensible configuration system for flows with type-safe inheritance
+- **Simple Runner Module**: CLI and programmatic utilities for easy pipeline execution
+- **Clean Prefect Exports**: Separate imports for Prefect decorators with and without tracing
+- **Expanded Exports**: All major components now accessible from top-level package import
+
+### API Improvements
+- Better type inference for document flows with custom options
+- Support for custom FlowOptions inheritance in pipeline flows
+- Improved error messages for invalid flow signatures
+- Enhanced document utility functions (`canonical_name_key`, `sanitize_url`)
+
+### Developer Experience
+- Simplified imports - most components available from `ai_pipeline_core` directly
+- Better separation of concerns between clean Prefect and traced pipeline decorators
+- More intuitive flow configuration with `FlowOptions` inheritance
+
 ## Stability Notice
 
-**Current Version**: 0.1.5
+**Current Version**: 0.1.7
 **Status**: Internal Preview
 **API Stability**: Unstable - Breaking changes expected
 **Recommended Use**: Learning and reference only

{ai_pipeline_core-0.1.5 → ai_pipeline_core-0.1.7}/README.md

@@ -109,40 +109,76 @@ async def process_document(doc: Document):
     return response.parsed
 ```
 
-### Prefect Flow Integration
+### Enhanced Pipeline Decorators (New in v0.1.7)
 ```python
-from prefect import flow, task
-from ai_pipeline_core.documents import Document, DocumentList, FlowDocument
-from ai_pipeline_core.flow import FlowConfig
-from ai_pipeline_core.tracing import trace
+from ai_pipeline_core import pipeline_flow, pipeline_task
+from ai_pipeline_core.flow import FlowOptions
+from ai_pipeline_core.documents import DocumentList, FlowDocument
 
-class OutputDocument(FlowDocument):
-    """Custom output document type"""
-    def get_type(self) -> str:
-        return "output"
+class CustomFlowOptions(FlowOptions):
+    """Extend base options with your custom fields"""
+    batch_size: int = 100
+    temperature: float = 0.7
 
-class MyFlowConfig(FlowConfig):
-    INPUT_DOCUMENT_TYPES = [InputDocument]
-    OUTPUT_DOCUMENT_TYPE = OutputDocument
-
-@task
-@trace
+@pipeline_task(trace_level="always", retries=3)
 async def process_task(doc: Document) -> Document:
-    # Task-level processing with automatic tracing
+    # Task with automatic tracing and retries
     result = await process_document(doc)
-    # Convert result to JSON string for document content
-    import json
-    return OutputDocument(name="result", content=json.dumps(result.model_dump()).encode())
+    return OutputDocument(name="result", content=result.encode())
+
+@pipeline_flow(trace_level="always")
+async def my_pipeline(
+    project_name: str,
+    documents: DocumentList,
+    flow_options: CustomFlowOptions  # Type-safe custom options
+) -> DocumentList:
+    # Pipeline flow with enforced signature and tracing
+    results = []
+    for doc in documents:
+        result = await process_task(doc)
+        results.append(result)
+    return DocumentList(results)
+```
 
-@flow
-async def my_pipeline(documents: DocumentList):
-    config = MyFlowConfig()
-    input_docs = config.get_input_documents(documents)
+### Simple Runner Utility (New in v0.1.7)
+```python
+from ai_pipeline_core.simple_runner import run_cli, run_pipeline
+from ai_pipeline_core.flow import FlowOptions
+
+# CLI-based pipeline execution
+if __name__ == "__main__":
+    run_cli(
+        flows=[my_pipeline],
+        flow_configs=[MyFlowConfig],
+        options_cls=CustomFlowOptions
+    )
 
-    results = await process_task.map(input_docs)
+# Or programmatic execution
+async def main():
+    result = await run_pipeline(
+        project_name="my-project",
+        output_dir=Path("./output"),
+        flow=my_pipeline,
+        flow_config=MyFlowConfig,
+        flow_options=CustomFlowOptions(batch_size=50)
+    )
+```
 
-    config.validate_output_documents(results)
-    return results
+### Clean Prefect Decorators (New in v0.1.7)
+```python
+# Import clean Prefect decorators without tracing
+from ai_pipeline_core.prefect import flow, task
+
+# Or use pipeline decorators with tracing
+from ai_pipeline_core import pipeline_flow, pipeline_task
+
+@task  # Clean Prefect task
+def compute(x: int) -> int:
+    return x * 2
+
+@pipeline_task(trace_level="always")  # With tracing
+def compute_traced(x: int) -> int:
+    return x * 2
 ```
 
 ## Core Modules
@@ -249,8 +285,14 @@ ai_pipeline_core/
 │   ├── client.py      # Async client implementation
 │   └── model_options.py # Configuration models
 ├── flow/              # Prefect flow utilities
-│   └── config.py      # Type-safe flow configuration
+│   ├── config.py      # Type-safe flow configuration
+│   └── options.py     # FlowOptions base class (v0.1.7)
+├── simple_runner/     # Pipeline execution utilities (v0.1.7)
+│   ├── cli.py         # CLI interface
+│   └── simple_runner.py # Core runner logic
 ├── logging/           # Structured logging
+├── pipeline.py        # Enhanced decorators (v0.1.7)
+├── prefect.py         # Clean Prefect exports (v0.1.7)
 ├── tracing.py         # Observability decorators
 └── settings.py        # Centralized configuration
 ```
@@ -427,9 +469,29 @@ Built with:
 - [LiteLLM](https://litellm.ai/) - LLM proxy
 - [Pydantic](https://pydantic-docs.helpmanual.io/) - Data validation
 
+## What's New in v0.1.7
+
+### Major Additions
+- **Enhanced Pipeline Decorators**: New `pipeline_flow` and `pipeline_task` decorators combining Prefect functionality with automatic LMNR tracing
+- **FlowOptions Base Class**: Extensible configuration system for flows with type-safe inheritance
+- **Simple Runner Module**: CLI and programmatic utilities for easy pipeline execution
+- **Clean Prefect Exports**: Separate imports for Prefect decorators with and without tracing
+- **Expanded Exports**: All major components now accessible from top-level package import
+
+### API Improvements
+- Better type inference for document flows with custom options
+- Support for custom FlowOptions inheritance in pipeline flows
+- Improved error messages for invalid flow signatures
+- Enhanced document utility functions (`canonical_name_key`, `sanitize_url`)
+
+### Developer Experience
+- Simplified imports - most components available from `ai_pipeline_core` directly
+- Better separation of concerns between clean Prefect and traced pipeline decorators
+- More intuitive flow configuration with `FlowOptions` inheritance
+
 ## Stability Notice
 
-**Current Version**: 0.1.5
+**Current Version**: 0.1.7
 **Status**: Internal Preview
 **API Stability**: Unstable - Breaking changes expected
 **Recommended Use**: Learning and reference only

ai_pipeline_core-0.1.7/ai_pipeline_core/__init__.py

@@ -0,0 +1,77 @@
+"""Pipeline Core - Shared infrastructure for AI pipelines."""
+
+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,
+    StructuredLoggerMixin,
+    get_pipeline_logger,
+    setup_logging,
+)
+from .logging import get_pipeline_logger as get_logger
+from .pipeline import pipeline_flow, pipeline_task
+from .prefect import flow, task
+from .prompt_manager import PromptManager
+from .settings import settings
+from .tracing import TraceInfo, TraceLevel, trace
+
+__version__ = "0.1.7"
+
+__all__ = [
+    # Config/Settings
+    "settings",
+    # Logging
+    "get_logger",
+    "get_pipeline_logger",
+    "LoggerMixin",
+    "LoggingConfig",
+    "setup_logging",
+    "StructuredLoggerMixin",
+    # Documents
+    "Document",
+    "DocumentList",
+    "FlowDocument",
+    "TaskDocument",
+    "canonical_name_key",
+    "sanitize_url",
+    # Flow/Task
+    "FlowConfig",
+    "FlowOptions",
+    # Prefect decorators (clean, no tracing)
+    "task",
+    "flow",
+    # Pipeline decorators (with tracing)
+    "pipeline_task",
+    "pipeline_flow",
+    # LLM
+    "llm",
+    "ModelName",
+    "ModelOptions",
+    "ModelResponse",
+    "StructuredModelResponse",
+    "AIMessages",
+    "AIMessageType",
+    # Tracing
+    "trace",
+    "TraceLevel",
+    "TraceInfo",
+    # Utils
+    "PromptManager",
+]

{ai_pipeline_core-0.1.5 → ai_pipeline_core-0.1.7}/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-0.1.5 → ai_pipeline_core-0.1.7}/ai_pipeline_core/documents/document.py

@@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
 from base64 import b32encode
 from enum import StrEnum
 from functools import cached_property
-from typing import Any, ClassVar, Literal, Self
+from typing import Any, ClassVar, Literal, Self, TypeVar
 
 from pydantic import BaseModel, ConfigDict, field_serializer, field_validator
 from ruamel.yaml import YAML
@@ -19,8 +19,11 @@ from .mime_type import (
     is_image_mime_type,
     is_pdf_mime_type,
     is_text_mime_type,
+    is_yaml_mime_type,
 )
 
+TModel = TypeVar("TModel", bound=BaseModel)
+
 
 class Document(BaseModel, ABC):
     """Abstract base class for all documents"""
@@ -207,15 +210,40 @@ class Document(BaseModel, ABC):
         """Parse document as JSON"""
         return json.loads(self.as_text())
 
+    def as_pydantic_model(self, model_type: type[TModel]) -> TModel:
+        """Parse document as a pydantic model and return the validated instance"""
+        data = self.as_yaml() if is_yaml_mime_type(self.mime_type) else self.as_json()
+        return model_type.model_validate(data)
+
     def as_markdown_list(self) -> list[str]:
         """Parse document as a markdown list"""
         return self.as_text().split(self.MARKDOWN_LIST_SEPARATOR)
 
     @classmethod
-    def create(cls, name: str, description: str | None, content: bytes | str) -> Self:
+    def create(
+        cls,
+        name: str,
+        description: str | None,
+        content: bytes | str | BaseModel | list[str] | Any,
+    ) -> Self:
         """Create a document from a name, description, and content"""
-        if isinstance(content, str):
+        is_yaml_extension = name.endswith(".yaml") or name.endswith(".yml")
+        is_json_extension = name.endswith(".json")
+        is_markdown_extension = name.endswith(".md")
+        is_str_list = isinstance(content, list) and all(isinstance(item, str) for item in content)
+        if isinstance(content, bytes):
+            pass
+        elif isinstance(content, str):
             content = content.encode("utf-8")
+        elif is_str_list and is_markdown_extension:
+            return cls.create_as_markdown_list(name, description, content)  # type: ignore[arg-type]
+        elif is_yaml_extension:
+            return cls.create_as_yaml(name, description, content)
+        elif is_json_extension:
+            return cls.create_as_json(name, description, content)
+        else:
+            raise ValueError(f"Unsupported content type: {type(content)} for {name}")
+
         return cls(name=name, description=description, content=content)
 
     @classmethod
@@ -230,6 +258,32 @@ class Document(BaseModel, ABC):
         content = Document.MARKDOWN_LIST_SEPARATOR.join(cleaned_items)
         return cls.create(name, description, content)
 
+    @classmethod
+    def create_as_json(cls, name: str, description: str | None, data: Any) -> Self:
+        """Create a document from a name, description, and JSON data"""
+        assert name.endswith(".json"), f"Document name must end with .json: {name}"
+        if isinstance(data, BaseModel):
+            data = data.model_dump(mode="json")
+        content = json.dumps(data, indent=2).encode("utf-8")
+        return cls.create(name, description, content)
+
+    @classmethod
+    def create_as_yaml(cls, name: str, description: str | None, data: Any) -> Self:
+        """Create a document from a name, description, and YAML data"""
+        assert name.endswith(".yaml") or name.endswith(".yml"), (
+            f"Document name must end with .yaml or .yml: {name}"
+        )
+        if isinstance(data, BaseModel):
+            data = data.model_dump()
+        yaml = YAML()
+        yaml.indent(mapping=2, sequence=4, offset=2)
+        from io import BytesIO
+
+        stream = BytesIO()
+        yaml.dump(data, stream)
+        content = stream.getvalue()
+        return cls.create(name, description, content)
+
     def serialize_model(self) -> dict[str, Any]:
         """Serialize document to a dictionary with proper encoding."""
         result = {

ai_pipeline_core-0.1.7/ai_pipeline_core/documents/mime_type.py

@@ -0,0 +1,110 @@
+"""MIME type detection utilities for documents"""
+
+import magic
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+logger = get_pipeline_logger(__name__)
+
+# Extension to MIME type mapping for common formats
+# These are formats where extension-based detection is more reliable
+EXTENSION_MIME_MAP = {
+    "md": "text/markdown",
+    "txt": "text/plain",
+    "pdf": "application/pdf",
+    "png": "image/png",
+    "jpg": "image/jpeg",
+    "jpeg": "image/jpeg",
+    "gif": "image/gif",
+    "bmp": "image/bmp",
+    "webp": "image/webp",
+    "json": "application/json",
+    "yaml": "application/yaml",
+    "yml": "application/yaml",
+    "xml": "text/xml",
+    "html": "text/html",
+    "htm": "text/html",
+    "py": "text/x-python",
+    "css": "text/css",
+    "js": "application/javascript",
+    "ts": "application/typescript",
+    "tsx": "application/typescript",
+    "jsx": "application/javascript",
+}
+
+
+def detect_mime_type(content: bytes, name: str) -> str:
+    """Detect MIME type from content and filename
+
+    Uses a hybrid approach:
+    1. Check for empty content
+    2. Try extension-based detection for known formats
+    3. Fall back to magic content detection
+    4. Final fallback to application/octet-stream
+    """
+
+    # Check for empty content
+    if len(content) == 0:
+        return "application/x-empty"
+
+    # Try extension-based detection first for known formats
+    # This is more reliable for text formats that magic might misidentify
+    ext = name.lower().split(".")[-1] if "." in name else ""
+    if ext in EXTENSION_MIME_MAP:
+        return EXTENSION_MIME_MAP[ext]
+
+    # Try content-based detection with magic
+    try:
+        mime = magic.from_buffer(content[:1024], mime=True)
+        # If magic returns a valid mime type, use it
+        if mime and mime != "application/octet-stream":
+            return mime
+    except (AttributeError, OSError, magic.MagicException) as e:
+        logger.warning(f"MIME detection failed for {name}: {e}")
+    except Exception as e:
+        logger.error(f"Unexpected error in MIME detection for {name}: {e}")
+
+    # Final fallback based on extension or default
+    return EXTENSION_MIME_MAP.get(ext, "application/octet-stream")
+
+
+def mime_type_from_extension(name: str) -> str:
+    """Get MIME type based on file extension
+
+    Legacy function kept for compatibility
+    """
+    ext = name.lower().split(".")[-1] if "." in name else ""
+    return EXTENSION_MIME_MAP.get(ext, "application/octet-stream")
+
+
+def is_text_mime_type(mime_type: str) -> bool:
+    """Check if MIME type represents text content"""
+    text_types = [
+        "text/",
+        "application/json",
+        "application/xml",
+        "application/javascript",
+        "application/yaml",
+        "application/x-yaml",
+    ]
+    return any(mime_type.startswith(t) for t in text_types)
+
+
+def is_json_mime_type(mime_type: str) -> bool:
+    """Check if MIME type is JSON"""
+    return mime_type == "application/json"
+
+
+def is_yaml_mime_type(mime_type: str) -> bool:
+    """Check if MIME type is YAML"""
+    return mime_type == "application/yaml" or mime_type == "application/x-yaml"
+
+
+def is_pdf_mime_type(mime_type: str) -> bool:
+    """Check if MIME type is PDF"""
+    return mime_type == "application/pdf"
+
+
+def is_image_mime_type(mime_type: str) -> bool:
+    """Check if MIME type is an image"""
+    return mime_type.startswith("image/")

ai_pipeline_core-0.1.7/ai_pipeline_core/flow/__init__.py

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

ai_pipeline_core-0.1.7/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-0.1.5 → ai_pipeline_core-0.1.7}/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(),