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

ai_pipeline_core/documents/document_list.py

@@ -1,3 +1,8 @@
+"""Type-safe list container for Document objects.
+
+@public
+"""
+
 from typing import Any, Iterable, SupportsIndex, Union, overload
 
 from typing_extensions import Self
@@ -6,14 +11,25 @@ from .document import Document
 
 
 class DocumentList(list[Document]):
-    """
-    A specialized list for Document objects with built-in validation.
+    """Type-safe container for Document objects.
+
+    @public
+
+    Specialized list with validation and filtering for documents.
+
+    Best Practice: Use default constructor in 90% of cases. Only enable
+    validate_same_type or validate_duplicates when you explicitly need them.
 
-    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
+    Example:
+        >>> # RECOMMENDED - default constructor for most cases
+        >>> docs = DocumentList([doc1, doc2])
+        >>> # Or empty initialization
+        >>> docs = DocumentList()
+        >>> docs.append(MyDocument(name="file.txt", content=b"data"))
+        >>>
+        >>> # Only use validation flags when specifically needed:
+        >>> docs = DocumentList(validate_same_type=True)  # Rare use case
+        >>> doc = docs.get_by("file.txt")  # Get by name
     """
 
     def __init__(
@@ -22,13 +38,14 @@ class DocumentList(list[Document]):
         validate_same_type: bool = False,
         validate_duplicates: bool = False,
     ) -> None:
-        """
-        Initialize DocumentList with optional initial documents.
+        """Initialize DocumentList.
+
+        @public
 
         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.
+            documents: Initial list of documents.
+            validate_same_type: Enforce same document type.
+            validate_duplicates: Prevent duplicate filenames.
         """
         super().__init__()
         self._validate_same_type = validate_same_type
@@ -37,7 +54,11 @@ class DocumentList(list[Document]):
             self.extend(documents)
 
     def _validate_no_duplicates(self) -> None:
-        """Validate that there are no duplicate filenames."""
+        """Check for duplicate document names.
+
+        Raises:
+            ValueError: If duplicate document names are found.
+        """
         if not self._validate_duplicates:
             return
 
@@ -53,7 +74,11 @@ class DocumentList(list[Document]):
             raise ValueError(f"Duplicate document names found: {unique_duplicates}")
 
     def _validate_no_description_files(self) -> None:
-        """Validate that no documents have DESCRIPTION_EXTENSION suffix."""
+        """Ensure no documents use reserved description file extension.
+
+        Raises:
+            ValueError: If any document uses the reserved description file extension.
+        """
         description_files = [
             doc.name for doc in self if doc.name.endswith(Document.DESCRIPTION_EXTENSION)
         ]
@@ -64,7 +89,11 @@ class DocumentList(list[Document]):
             )
 
     def _validate_types(self) -> None:
-        """Validate that all documents have the same class type if required."""
+        """Ensure all documents are of the same class type.
+
+        Raises:
+            ValueError: If documents have different class types.
+        """
         if not self._validate_same_type or not self:
             return
 
@@ -75,23 +104,23 @@ class DocumentList(list[Document]):
             raise ValueError(f"All documents must have the same type. Found types: {types}")
 
     def _validate(self) -> None:
-        """Run all validations."""
+        """Run all configured validation checks."""
         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."""
+        """Add a document to the end of the list."""
         super().append(document)
         self._validate()
 
     def extend(self, documents: Iterable[Document]) -> None:
-        """Extend the list with multiple documents with validation."""
+        """Add multiple documents to the list."""
         super().extend(documents)
         self._validate()
 
     def insert(self, index: SupportsIndex, document: Document) -> None:
-        """Insert a document at the specified index with validation."""
+        """Insert a document at the specified position."""
         super().insert(index, document)
         self._validate()
 
@@ -102,30 +131,110 @@ class DocumentList(list[Document]):
     def __setitem__(self, index: slice, value: Iterable[Document]) -> None: ...
 
     def __setitem__(self, index: Union[SupportsIndex, slice], value: Any) -> None:
-        """Set item with validation."""
+        """Set item or slice with validation."""
         super().__setitem__(index, value)
         self._validate()
 
     def __iadd__(self, other: Any) -> "Self":
-        """In-place addition with validation."""
+        """In-place addition (+=) with validation.
+
+        Returns:
+            Self: This DocumentList after modification.
+        """
         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
+    @overload
+    def filter_by(self, arg: str) -> "DocumentList": ...
+
+    @overload
+    def filter_by(self, arg: type[Document]) -> "DocumentList": ...
+
+    @overload
+    def filter_by(self, arg: list[type[Document]]) -> "DocumentList": ...
+
+    def filter_by(self, arg: str | type[Document] | list[type[Document]]) -> "DocumentList":
+        """Filter documents by name or type(s).
+
+        @public
+
+        Args:
+            arg: Document name (str), single document type, or list of document types.
+
+        Returns:
+            New DocumentList with filtered documents.
+
+        Raises:
+            TypeError: If arg is not a valid type (str, Document type, or list of Document types).
+
+        Example:
+            >>> docs.filter_by("file.txt")  # Filter by name
+            >>> docs.filter_by(MyDocument)  # Filter by type
+            >>> docs.filter_by([Doc1, Doc2])  # Filter by multiple types
+        """
+        if isinstance(arg, str):
+            # Filter by name
+            return DocumentList([doc for doc in self if doc.name == arg])
+        elif isinstance(arg, type):
+            # Filter by single type (including subclasses)
+            return DocumentList([doc for doc in self if isinstance(doc, arg)])
+        elif isinstance(arg, list):  # type: ignore[reportUnnecessaryIsInstance]
+            # Filter by multiple types
+            documents = DocumentList()
+            for document_type in arg:
+                documents.extend([doc for doc in self if isinstance(doc, document_type)])
+            return documents
+        else:
+            raise TypeError(f"Invalid argument type for filter_by: {type(arg)}")
+
+    @overload
+    def get_by(self, arg: str) -> Document: ...
+
+    @overload
+    def get_by(self, arg: type[Document]) -> Document: ...
+
+    @overload
+    def get_by(self, arg: str, required: bool = True) -> Document | None: ...
+
+    @overload
+    def get_by(self, arg: type[Document], required: bool = True) -> Document | None: ...
+
+    def get_by(self, arg: str | type[Document], required: bool = True) -> Document | None:
+        """Get a single document by name or type.
+
+        @public
+
+        Args:
+            arg: Document name (str) or document type.
+            required: If True, raises ValueError when not found. If False, returns None.
+
+        Returns:
+            The first matching document, or None if not found and required=False.
+
+        Raises:
+            ValueError: If required=True and document not found.
+            TypeError: If arg is not a string or Document type.
+
+        Example:
+            >>> doc = docs.get_by("file.txt")  # Get by name, raises if not found
+            >>> doc = docs.get_by(MyDocument, required=False)  # Returns None if not found
+        """
+        if isinstance(arg, str):
+            # Get by name
+            for doc in self:
+                if doc.name == arg:
+                    return doc
+            if required:
+                raise ValueError(f"Document with name '{arg}' not found")
+            return None
+        elif isinstance(arg, type):  # type: ignore[reportUnnecessaryIsInstance]
+            # Get by type (including subclasses)
+            for doc in self:
+                if isinstance(doc, arg):
+                    return doc
+            if required:
+                raise ValueError(f"Document of type '{arg.__name__}' not found")
+            return None
+        else:
+            raise TypeError(f"Invalid argument type for get_by: {type(arg)}")

ai_pipeline_core/documents/flow_document.py

@@ -1,27 +1,128 @@
-"""Flow-specific document base class."""
+"""Flow-specific document base class for persistent pipeline data.
 
-from typing import Any, Literal, final
+@public
+
+This module provides the FlowDocument abstract base class for documents
+that need to persist across Prefect flow runs and between pipeline steps.
+"""
+
+from typing import Literal, final
 
 from .document import Document
 
 
 class FlowDocument(Document):
-    """
-    Abstract base class for flow-specific documents.
+    """Abstract base class for documents that persist across flow runs.
+
+    @public
+
+    FlowDocument is used for data that needs to be saved between pipeline
+    steps and across multiple flow executions. These documents are typically
+    written to the file system using the simple_runner utilities.
+
+    Key characteristics:
+    - Persisted to file system between pipeline steps
+    - Survives across multiple flow runs
+    - Used for flow inputs and outputs
+    - Saved in directories named after the document's canonical name
+
+    Creating FlowDocuments:
+        **Use the `create` classmethod** for most use cases. It handles automatic
+        conversion of various content types. Only use __init__ when you have bytes.
+
+        >>> from enum import StrEnum
+        >>>
+        >>> # Simple document with pass:
+        >>> class MyDoc(FlowDocument):
+        ...     pass
+        >>>
+        >>> # Document with restricted file names:
+        >>> class ConfigDoc(FlowDocument):
+        ...     class FILES(StrEnum):
+        ...         CONFIG = "config.yaml"
+        ...         SETTINGS = "settings.json"
+        >>>
+        >>> # RECOMMENDED - automatic conversion:
+        >>> doc = MyDoc.create(name="data.json", content={"key": "value"})
+        >>> doc = ConfigDoc.create(name="config.yaml", content={"host": "localhost"})
+
+    Persistence:
+        Documents are saved to: {output_dir}/{canonical_name}/{filename}
+        For example: output/my_doc/data.json
 
-    Flow documents represent inputs, outputs, and intermediate results
-    within a Prefect flow execution context.
+    Note:
+        - Cannot instantiate FlowDocument directly - must subclass
+        - Used with FlowConfig to define flow input/output types
+        - No additional abstract methods to implement
 
-    Compared to TaskDocument, FlowDocument are persistent across Prefect flow runs.
+    See Also:
+        TaskDocument: For temporary documents within task execution
+        TemporaryDocument: For documents that are never persisted
     """
 
-    def __init__(self, **data: Any) -> None:
-        """Prevent direct instantiation of abstract FlowDocument class."""
+    def __init__(
+        self,
+        *,
+        name: str,
+        content: bytes,
+        description: str | None = None,
+    ) -> None:
+        """Initialize a FlowDocument with raw bytes content.
+
+        Important:
+            **Most users should use the `create` classmethod instead of __init__.**
+            The create method provides automatic content conversion for various types
+            (str, dict, list, Pydantic models) while __init__ only accepts bytes.
+
+        Prevents direct instantiation of the abstract FlowDocument class.
+        FlowDocument must be subclassed for specific document types.
+
+        Args:
+            name: Document filename (required, keyword-only)
+            content: Document content as raw bytes (required, keyword-only)
+            description: Optional human-readable description (keyword-only)
+
+        Raises:
+            TypeError: If attempting to instantiate FlowDocument directly
+                      instead of using a concrete subclass.
+
+        Example:
+            >>> from enum import StrEnum
+            >>>
+            >>> # Simple subclass:
+            >>> class MyFlowDoc(FlowDocument):
+            ...     pass
+            >>>
+            >>> # With FILES restriction:
+            >>> class RestrictedDoc(FlowDocument):
+            ...     class FILES(StrEnum):
+            ...         DATA = "data.json"
+            ...         METADATA = "metadata.yaml"
+            >>>
+            >>> # Direct constructor - only for bytes:
+            >>> doc = MyFlowDoc(name="test.bin", content=b"raw data")
+            >>>
+            >>> # RECOMMENDED - use create for automatic conversion:
+            >>> doc = RestrictedDoc.create(name="data.json", content={"key": "value"})
+            >>> # This would raise DocumentNameError:
+            >>> # doc = RestrictedDoc.create(name="other.json", content={})
+        """
         if type(self) is FlowDocument:
             raise TypeError("Cannot instantiate abstract FlowDocument class directly")
-        super().__init__(**data)
+        super().__init__(name=name, content=content, description=description)
 
     @final
     def get_base_type(self) -> Literal["flow"]:
-        """Get the document type."""
+        """Return the base type identifier for flow documents.
+
+        This method is final and cannot be overridden by subclasses.
+        It identifies this document as a flow-persistent document.
+
+        Returns:
+            "flow" - Indicates this document persists across flow runs.
+
+        Note:
+            This determines the document's lifecycle and persistence behavior
+            in the pipeline system.
+        """
         return "flow"

ai_pipeline_core/documents/mime_type.py

@@ -1,4 +1,10 @@
-"""MIME type detection utilities for documents"""
+"""@internal MIME type detection utilities for documents.
+
+This module provides functions for detecting and validating MIME types
+from document content and filenames. It uses a hybrid approach combining
+extension-based detection for known formats and content analysis via
+python-magic for unknown files.
+"""
 
 import magic
 
@@ -34,15 +40,45 @@ EXTENSION_MIME_MAP = {
 
 
 def detect_mime_type(content: bytes, name: str) -> str:
-    """Detect MIME type from content and filename
+    r"""Detect MIME type from document 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
-    """
+    Uses a multi-stage detection strategy for maximum accuracy:
+    1. Returns 'application/x-empty' for empty content
+    2. Uses extension-based detection for known formats (most reliable)
+    3. Falls back to python-magic content analysis
+    4. Final fallback to extension or 'application/octet-stream'
+
+    Args:
+        content: Document content as bytes.
+        name: Filename with extension.
+
+    Returns:
+        MIME type string (e.g., 'text/plain', 'application/json').
+        Never returns None or empty string.
+
+    Fallback behavior:
+        - Empty content: 'application/x-empty'
+        - Unknown extension with binary content: 'application/octet-stream'
+        - Magic library failure: Falls back to extension or 'application/octet-stream'
 
+    Performance:
+        Only the first 1024 bytes are analyzed for content detection.
+        Extension-based detection is O(1) lookup.
+
+    Note:
+        Extension-based detection is preferred for text formats as
+        content analysis can sometimes misidentify structured text.
+
+    Example:
+        >>> detect_mime_type(b'{"key": "value"}', "data.json")
+        'application/json'
+        >>> detect_mime_type(b'Hello World', "text.txt")
+        'text/plain'
+        >>> detect_mime_type(b'', "empty.txt")
+        'application/x-empty'
+        >>> detect_mime_type(b'\\x89PNG', "image.xyz")
+        'image/png'  # Magic detects PNG despite wrong extension
+    """
     # Check for empty content
     if len(content) == 0:
         return "application/x-empty"
@@ -69,16 +105,60 @@ def detect_mime_type(content: bytes, name: str) -> str:
 
 
 def mime_type_from_extension(name: str) -> str:
-    """Get MIME type based on file extension
+    """Get MIME type based solely on file extension.
+
+    Simple extension-based MIME type detection without content analysis.
+    This is a legacy function maintained for backward compatibility.
+
+    Args:
+        name: Filename with extension.
 
-    Legacy function kept for compatibility
+    Returns:
+        MIME type based on extension, or 'application/octet-stream'
+        if extension is unknown.
+
+    Note:
+        Prefer detect_mime_type() for more accurate detection.
+        This function only checks the file extension.
+
+    Example:
+        >>> mime_type_from_extension("document.pdf")
+        'application/pdf'
+        >>> mime_type_from_extension("unknown.xyz")
+        'application/octet-stream'
     """
     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"""
+    """Check if MIME type represents text-based content.
+
+    Determines if content can be safely decoded as text.
+    Includes common text formats and structured text like JSON/YAML.
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if MIME type indicates text content, False otherwise.
+
+    Recognized as text:
+        - Any type starting with 'text/'
+        - application/json
+        - application/xml
+        - application/javascript
+        - application/yaml
+        - application/x-yaml
+
+    Example:
+        >>> is_text_mime_type('text/plain')
+        True
+        >>> is_text_mime_type('application/json')
+        True
+        >>> is_text_mime_type('image/png')
+        False
+    """
     text_types = [
         "text/",
         "application/json",
@@ -91,20 +171,98 @@ def is_text_mime_type(mime_type: str) -> bool:
 
 
 def is_json_mime_type(mime_type: str) -> bool:
-    """Check if MIME type is JSON"""
+    """Check if MIME type is JSON.
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if MIME type is 'application/json', False otherwise.
+
+    Note:
+        Only matches exact 'application/json', not variants like
+        'application/ld+json' or 'application/vnd.api+json'.
+
+    Example:
+        >>> is_json_mime_type('application/json')
+        True
+        >>> is_json_mime_type('text/json')  # Not standard JSON MIME
+        False
+    """
     return mime_type == "application/json"
 
 
 def is_yaml_mime_type(mime_type: str) -> bool:
-    """Check if MIME type is YAML"""
+    """Check if MIME type is YAML.
+
+    Recognizes both standard YAML MIME types.
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if MIME type is YAML, False otherwise.
+
+    Recognized types:
+        - application/yaml (standard)
+        - application/x-yaml (legacy)
+
+    Example:
+        >>> is_yaml_mime_type('application/yaml')
+        True
+        >>> is_yaml_mime_type('application/x-yaml')
+        True
+    """
     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"""
+    """Check if MIME type is PDF.
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if MIME type is 'application/pdf', False otherwise.
+
+    Note:
+        PDF documents require special handling in the LLM module
+        and are supported by certain vision-capable models.
+
+    Example:
+        >>> is_pdf_mime_type('application/pdf')
+        True
+        >>> is_pdf_mime_type('text/plain')
+        False
+    """
     return mime_type == "application/pdf"
 
 
 def is_image_mime_type(mime_type: str) -> bool:
-    """Check if MIME type is an image"""
+    """Check if MIME type represents an image.
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if MIME type starts with 'image/', False otherwise.
+
+    Recognized formats:
+        Any MIME type starting with 'image/' including:
+        - image/png
+        - image/jpeg
+        - image/gif
+        - image/webp
+        - image/svg+xml
+
+    Note:
+        Image documents are automatically encoded for vision-capable
+        LLM models in the AIMessages.document_to_prompt() method.
+
+    Example:
+        >>> is_image_mime_type('image/png')
+        True
+        >>> is_image_mime_type('application/pdf')
+        False
+    """
     return mime_type.startswith("image/")