ai-pipeline-core 0.1.1__py3-none-any.whl

ai_pipeline_core/__init__.py

@@ -0,0 +1,36 @@
+"""Pipeline Core - Shared infrastructure for AI pipelines."""
+
+from .documents import Document, DocumentList, FlowDocument, TaskDocument
+from .flow import FlowConfig
+from .logging import (
+    LoggerMixin,
+    LoggingConfig,
+    StructuredLoggerMixin,
+    get_pipeline_logger,
+    setup_logging,
+)
+from .logging import (
+    get_pipeline_logger as get_logger,
+)
+from .prompt_manager import PromptManager
+from .settings import settings
+from .tracing import trace
+
+__version__ = "0.1.1"
+
+__all__ = [
+    "Document",
+    "DocumentList",
+    "FlowConfig",
+    "FlowDocument",
+    "get_logger",
+    "get_pipeline_logger",
+    "LoggerMixin",
+    "LoggingConfig",
+    "PromptManager",
+    "settings",
+    "setup_logging",
+    "StructuredLoggerMixin",
+    "TaskDocument",
+    "trace",
+]

ai_pipeline_core/documents/__init__.py

@@ -0,0 +1,11 @@
+from .document import Document
+from .document_list import DocumentList
+from .flow_document import FlowDocument
+from .task_document import TaskDocument
+
+__all__ = [
+    "Document",
+    "DocumentList",
+    "FlowDocument",
+    "TaskDocument",
+]

ai_pipeline_core/documents/document.py

@@ -0,0 +1,252 @@
+import base64
+import hashlib
+import json
+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 pydantic import BaseModel, ConfigDict, field_serializer, field_validator
+from ruamel.yaml import YAML
+
+from ai_pipeline_core.exceptions import DocumentNameError, DocumentSizeError
+
+from .mime_type import (
+    detect_mime_type,
+    is_image_mime_type,
+    is_pdf_mime_type,
+    is_text_mime_type,
+)
+
+
+class Document(BaseModel, ABC):
+    """Abstract base class for all documents"""
+
+    MAX_CONTENT_SIZE: ClassVar[int] = 10 * 1024 * 1024  # 10MB default
+    DESCRIPTION_EXTENSION: ClassVar[str] = ".description.md"
+
+    # Optional enum of allowed file names. Subclasses may set this.
+    # This is used to validate the document name.
+    FILES: ClassVar[type[StrEnum] | None] = None
+
+    name: str
+    description: str | None = None
+    content: bytes
+
+    # Pydantic configuration
+    model_config = ConfigDict(
+        frozen=True,  # Make documents immutable
+        arbitrary_types_allowed=True,
+    )
+
+    @abstractmethod
+    def get_base_type(self) -> Literal["flow", "task"]:
+        """Get the type of the document - must be implemented by subclasses"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def base_type(self) -> Literal["flow", "task"]:
+        """Alias for document_type for backward compatibility"""
+        return self.get_base_type()
+
+    @property
+    def is_flow(self) -> bool:
+        """Check if document is a flow document"""
+        return self.get_base_type() == "flow"
+
+    @property
+    def is_task(self) -> bool:
+        """Check if document is a task document"""
+        return self.get_base_type() == "task"
+
+    @classmethod
+    def get_expected_files(cls) -> list[str] | None:
+        """
+        Return the list of allowed file names for this document class, or None if unrestricted.
+        """
+        if cls.FILES is None:
+            return None
+        try:
+            values = [member.value for member in cls.FILES]
+        except TypeError:
+            raise DocumentNameError(f"{cls.__name__}.FILES must be an Enum of string values")
+        if len(values) == 0:
+            return None
+        return values
+
+    @classmethod
+    def validate_file_name(cls, name: str) -> None:
+        """
+        Optional file-name validation hook.
+
+        Default behavior:
+        - If `FILES` enum is defined on the subclass, ensure the **basename** of `name`
+          equals one of the enum values (exact string match).
+        - If `FILES` is None, do nothing.
+
+        Override this method in subclasses for custom conventions (regex, prefixes, etc.).
+        Raise DocumentNameError when invalid.
+        """
+        if cls.FILES is None:
+            return
+
+        try:
+            allowed = {str(member.value) for member in cls.FILES}  # type: ignore[arg-type]
+        except TypeError:
+            raise DocumentNameError(f"{cls.__name__}.FILES must be an Enum of string values")
+
+        if name not in allowed:
+            allowed_str = ", ".join(sorted(allowed))
+            raise DocumentNameError(f"Invalid filename '{name}'. Allowed names: {allowed_str}")
+
+    @field_validator("name")
+    def validate_name(cls, v: str) -> str:
+        """Validate document name matches expected patterns and is secure"""
+        if v.endswith(cls.DESCRIPTION_EXTENSION):
+            raise DocumentNameError(
+                f"Document names cannot end with {cls.DESCRIPTION_EXTENSION}: {v}"
+            )
+
+        if ".." in v or "\\" in v or "/" in v:
+            raise DocumentNameError(f"Invalid filename - contains path traversal characters: {v}")
+
+        if not v or v.startswith(" ") or v.endswith(" "):
+            raise DocumentNameError(f"Invalid filename format: {v}")
+
+        cls.validate_file_name(v)
+
+        return v
+
+    @field_validator("content")
+    def validate_content(cls, v: bytes) -> bytes:
+        """Validate content size"""
+        # Check content size limit
+        max_size = getattr(cls, "MAX_CONTENT_SIZE", 100 * 1024 * 1024)
+        if len(v) > max_size:
+            raise DocumentSizeError(
+                f"Document size ({len(v)} bytes) exceeds maximum allowed size ({max_size} bytes)"
+            )
+
+        return v
+
+    @field_serializer("content")
+    def serialize_content(self, v: bytes) -> str:
+        """Serialize bytes content to string for JSON serialization"""
+        try:
+            return v.decode("utf-8")
+        except UnicodeDecodeError:
+            # Fall back to base64 for binary content
+            return base64.b64encode(v).decode("ascii")
+
+    @property
+    def id(self) -> str:
+        """Return the first 6 characters of the SHA256 hash of the content, encoded in base32"""
+        return self.sha256[:6]
+
+    @cached_property
+    def sha256(self) -> str:
+        """Full SHA256 hash of content, encoded in base32"""
+        return b32encode(hashlib.sha256(self.content).digest()).decode("ascii").upper()
+
+    @property
+    def size(self) -> int:
+        """Size of content in bytes"""
+        return len(self.content)
+
+    @cached_property
+    def detected_mime_type(self) -> str:
+        """Detect MIME type from content using python-magic"""
+        return detect_mime_type(self.content, self.name)
+
+    @property
+    def mime_type(self) -> str:
+        """Get MIME type - uses content detection with fallback to extension"""
+        return self.detected_mime_type
+
+    @property
+    def is_text(self) -> bool:
+        """Check if document is text based on MIME type"""
+        return is_text_mime_type(self.mime_type)
+
+    @property
+    def is_pdf(self) -> bool:
+        """Check if document is PDF"""
+        return is_pdf_mime_type(self.mime_type)
+
+    @property
+    def is_image(self) -> bool:
+        """Check if document is an image"""
+        return is_image_mime_type(self.mime_type)
+
+    @property
+    def should_be_cached(self) -> bool:
+        """Check if document should be cached"""
+        return False
+
+    def as_text(self) -> str:
+        """Parse document as text"""
+        if not self.is_text:
+            raise ValueError(f"Document is not text: {self.name}")
+        return self.content.decode("utf-8")
+
+    def as_yaml(self) -> Any:
+        """Parse document as YAML"""
+        if not self.is_text:
+            raise ValueError(f"Document is not text: {self.name}")
+        return YAML().load(self.content.decode("utf-8"))  # type: ignore
+
+    def as_json(self) -> Any:
+        """Parse document as JSON"""
+        if not self.is_text:
+            raise ValueError(f"Document is not text: {self.name}")
+        return json.loads(self.content.decode("utf-8"))
+
+    def serialize_model(self) -> dict[str, Any]:
+        """Serialize document to a dictionary with proper encoding."""
+        result = {
+            "name": self.name,
+            "description": self.description,
+            "base_type": self.get_base_type(),
+            "size": self.size,
+            "id": self.id,
+            "sha256": self.sha256,
+            "mime_type": self.mime_type,
+        }
+
+        # Try to encode content as UTF-8, fall back to base64
+        if self.is_text or self.mime_type.startswith("text/"):
+            try:
+                result["content"] = self.content.decode("utf-8")
+                result["content_encoding"] = "utf-8"
+            except UnicodeDecodeError:
+                # For text files with encoding issues, use UTF-8 with replacement
+                result["content"] = self.content.decode("utf-8", errors="replace")
+                result["content_encoding"] = "utf-8"
+        else:
+            # Binary content - use base64
+            result["content"] = base64.b64encode(self.content).decode("ascii")
+            result["content_encoding"] = "base64"
+
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Self:
+        """Deserialize document from dictionary."""
+        # Extract content and encoding
+        content_str = data.get("content", "")
+        content_encoding = data.get("content_encoding", "utf-8")
+
+        # Decode content based on encoding
+        if content_encoding == "base64":
+            content = base64.b64decode(content_str)
+        else:
+            # Default to UTF-8
+            content = content_str.encode("utf-8")
+
+        # Create document with the required fields
+        return cls(
+            name=data["name"],
+            content=content,
+            description=data.get("description"),
+        )

ai_pipeline_core/documents/document_list.py

@@ -0,0 +1,131 @@
+from typing import Any, Iterable, SupportsIndex, Union, overload
+
+from typing_extensions import Self
+
+from .document import Document
+
+
+class DocumentList(list[Document]):
+    """
+    A specialized list for Document objects with built-in validation.
+
+    Features:
+    - Optionally ensures no duplicate filenames within the list
+    - Optionally validates that all documents have the same type (for flow outputs)
+    - Provides convenience methods for document operations
+    - Works with both FlowDocument and TaskDocument classes
+    """
+
+    def __init__(
+        self,
+        documents: list[Document] | None = None,
+        validate_same_type: bool = False,
+        validate_duplicates: bool = False,
+    ) -> None:
+        """
+        Initialize DocumentList with optional initial documents.
+
+        Args:
+                documents: Initial list of documents
+                validate_same_type: If True, validates that all documents have the same type.
+                                                 Should be True for flow outputs, False for inputs.
+        """
+        super().__init__()
+        self._validate_same_type = validate_same_type
+        self._validate_duplicates = validate_duplicates
+        if documents:
+            self.extend(documents)
+
+    def _validate_no_duplicates(self) -> None:
+        """Validate that there are no duplicate filenames."""
+        if not self._validate_duplicates:
+            return
+
+        filenames = [doc.name for doc in self]
+        seen: set[str] = set()
+        duplicates: list[str] = []
+        for name in filenames:
+            if name in seen:
+                duplicates.append(name)
+            seen.add(name)
+        if duplicates:
+            unique_duplicates = list(set(duplicates))
+            raise ValueError(f"Duplicate document names found: {unique_duplicates}")
+
+    def _validate_no_description_files(self) -> None:
+        """Validate that no documents have DESCRIPTION_EXTENSION suffix."""
+        description_files = [
+            doc.name for doc in self if doc.name.endswith(Document.DESCRIPTION_EXTENSION)
+        ]
+        if description_files:
+            raise ValueError(
+                f"Documents with {Document.DESCRIPTION_EXTENSION} suffix are not allowed: "
+                f"{description_files}"
+            )
+
+    def _validate_types(self) -> None:
+        """Validate that all documents have the same class type if required."""
+        if not self._validate_same_type or not self:
+            return
+
+        first_class = type(self[0])
+        different_types = [doc for doc in self if type(doc) is not first_class]
+        if different_types:
+            types = list({type(doc).__name__ for doc in self})
+            raise ValueError(f"All documents must have the same type. Found types: {types}")
+
+    def _validate(self) -> None:
+        """Run all validations."""
+        self._validate_no_duplicates()
+        self._validate_no_description_files()
+        self._validate_types()
+
+    def append(self, document: Document) -> None:
+        """Add a document to the list with validation."""
+        super().append(document)
+        self._validate()
+
+    def extend(self, documents: Iterable[Document]) -> None:
+        """Extend the list with multiple documents with validation."""
+        super().extend(documents)
+        self._validate()
+
+    def insert(self, index: SupportsIndex, document: Document) -> None:
+        """Insert a document at the specified index with validation."""
+        super().insert(index, document)
+        self._validate()
+
+    @overload
+    def __setitem__(self, index: SupportsIndex, value: Document) -> None: ...
+
+    @overload
+    def __setitem__(self, index: slice, value: Iterable[Document]) -> None: ...
+
+    def __setitem__(self, index: Union[SupportsIndex, slice], value: Any) -> None:
+        """Set item with validation."""
+        super().__setitem__(index, value)
+        self._validate()
+
+    def __iadd__(self, other: Any) -> "Self":
+        """In-place addition with validation."""
+        result = super().__iadd__(other)
+        self._validate()
+        return result
+
+    def filter_by_type(self, document_type: type[Document]) -> "DocumentList":
+        """Return a new DocumentList containing only instances of the specified document class."""
+        return DocumentList([doc for doc in self if type(doc) is document_type])
+
+    def filter_by_types(self, document_types: list[type[Document]]) -> "DocumentList":
+        """Return a new DocumentList containing only instances of the specified document classes."""
+        documents = DocumentList()
+        for document_type in document_types:
+            documents.extend(self.filter_by_type(document_type))
+        return documents
+
+    def get_by_name(self, name: str) -> Document | None:
+        """Get a document by its name."""
+        for doc in self:
+            if doc.name == name:
+                return doc
+        return None

ai_pipeline_core/documents/flow_document.py

@@ -0,0 +1,21 @@
+"""Flow-specific document base class."""
+
+from typing import Literal, final
+
+from .document import Document
+
+
+class FlowDocument(Document):
+    """
+    Abstract base class for flow-specific documents.
+
+    Flow documents represent inputs, outputs, and intermediate results
+    within a Prefect flow execution context.
+
+    Compared to TaskDocument, FlowDocument are persistent across Prefect flow runs.
+    """
+
+    @final
+    def get_base_type(self) -> Literal["flow"]:
+        """Get the document type."""
+        return "flow"

ai_pipeline_core/documents/mime_type.py

@@ -0,0 +1,78 @@
+"""MIME type detection utilities for documents"""
+
+import magic
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+logger = get_pipeline_logger(__name__)
+
+
+def detect_mime_type(content: bytes, name: str) -> str:
+    """Detect MIME type from content using python-magic"""
+
+    try:
+        if name.endswith(".md") and content.decode("utf-8"):
+            return "text/markdown"
+    except UnicodeDecodeError:
+        pass
+
+    if len(content) <= 4:
+        return "application/x-empty"
+
+    try:
+        mime = magic.from_buffer(content[:1024], mime=True)
+        return mime
+    except (AttributeError, OSError, magic.MagicException) as e:
+        logger.warning(f"MIME detection failed for {name}: {e}, falling back to extension")
+        return mime_type_from_extension(name)
+    except Exception as e:
+        logger.error(f"Unexpected error in MIME detection for {name}: {e}")
+        return mime_type_from_extension(name)
+
+
+def mime_type_from_extension(name: str) -> str:
+    """Get MIME type based on file extension"""
+    ext = name.lower().split(".")[-1] if "." in name else ""
+
+    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",
+    }
+
+    return 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_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/documents/task_document.py

@@ -0,0 +1,22 @@
+"""Task-specific document base class."""
+
+from typing import Literal, final
+
+from .document import Document
+
+
+class TaskDocument(Document):
+    """
+    Abstract base class for task-specific documents.
+
+    Task documents represent inputs, outputs, and intermediate results
+    within a Prefect task execution context.
+
+    Compared to FlowDocument, TaskDocument are not persisted across Prefect task runs.
+    They are used for intermediate results that are not needed after the task completes.
+    """
+
+    @final
+    def get_base_type(self) -> Literal["task"]:
+        """Get the document type."""
+        return "task"

ai_pipeline_core/documents/utils.py

@@ -0,0 +1,33 @@
+import re
+from urllib.parse import urlparse
+
+
+def sanitize_url(url: str) -> str:
+    """
+    Sanitize URL or query string for use in filenames.
+    Removes or replaces characters that are invalid in filenames.
+    """
+    # Remove protocol if it's a URL
+    if url.startswith(("http://", "https://")):
+        parsed = urlparse(url)
+        # Use domain + path
+        url = parsed.netloc + parsed.path
+
+    # Replace invalid filename characters
+    sanitized = re.sub(r'[<>:"/\\|?*]', "_", url)
+
+    # Replace multiple underscores with single one
+    sanitized = re.sub(r"_+", "_", sanitized)
+
+    # Remove leading/trailing underscores and dots
+    sanitized = sanitized.strip("_.")
+
+    # Limit length to prevent too long filenames
+    if len(sanitized) > 100:
+        sanitized = sanitized[:100]
+
+    # Ensure we have something
+    if not sanitized:
+        sanitized = "unnamed"
+
+    return sanitized

ai_pipeline_core/exceptions.py

@@ -0,0 +1,61 @@
+"""Exception hierarchy for AI Pipeline Core."""
+
+
+class PipelineCoreError(Exception):
+    """Base exception for all pipeline errors."""
+
+    pass
+
+
+class DocumentError(PipelineCoreError):
+    """Document-related errors."""
+
+    pass
+
+
+class DocumentValidationError(DocumentError):
+    """Document validation failed."""
+
+    pass
+
+
+class DocumentSizeError(DocumentValidationError):
+    """Document size exceeds limits."""
+
+    pass
+
+
+class DocumentNameError(DocumentValidationError):
+    """Invalid document name."""
+
+    pass
+
+
+class LLMError(PipelineCoreError):
+    """LLM-related errors."""
+
+    pass
+
+
+class PromptError(PipelineCoreError):
+    """Prompt-related errors."""
+
+    pass
+
+
+class PromptRenderError(PromptError):
+    """Failed to render prompt template."""
+
+    pass
+
+
+class PromptNotFoundError(PromptError):
+    """Prompt template not found."""
+
+    pass
+
+
+class MimeTypeError(DocumentError):
+    """MIME type detection or validation error."""
+
+    pass

ai_pipeline_core/flow/__init__.py

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