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

ai_pipeline_core/__init__.py

@@ -16,7 +16,7 @@ from .prompt_manager import PromptManager
 from .settings import settings
 from .tracing import trace
 
-__version__ = "0.1.4"
+__version__ = "0.1.6"
 
 __all__ = [
     "Document",

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/documents/mime_type.py

@@ -6,53 +6,75 @@ 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 using python-magic"""
+    """Detect MIME type from content and filename
 
-    try:
-        if name.endswith(".md") and content.decode("utf-8"):
-            return "text/markdown"
-    except UnicodeDecodeError:
-        pass
+    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
+    """
 
-    if len(content) <= 4:
+    # 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)
-        return mime
+        # 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}, falling back to extension")
-        return mime_type_from_extension(name)
+        logger.warning(f"MIME detection failed for {name}: {e}")
     except Exception as e:
         logger.error(f"Unexpected error in MIME detection for {name}: {e}")
-        return mime_type_from_extension(name)
+
+    # 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"""
-    ext = name.lower().split(".")[-1] if "." in name else ""
+    """Get MIME type based on file 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",
-    }
-
-    return mime_map.get(ext, "application/octet-stream")
+    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:
@@ -68,6 +90,16 @@ def is_text_mime_type(mime_type: str) -> bool:
     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"

{ai_pipeline_core-0.1.5.dist-info → ai_pipeline_core-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.1.5
+Version: 0.1.6
 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
@@ -471,7 +471,7 @@ Built with:
 
 ## Stability Notice
 
-**Current Version**: 0.1.5
+**Current Version**: 0.1.6
 **Status**: Internal Preview
 **API Stability**: Unstable - Breaking changes expected
 **Recommended Use**: Learning and reference only

{ai_pipeline_core-0.1.5.dist-info → ai_pipeline_core-0.1.6.dist-info}/RECORD

@@ -1,14 +1,14 @@
-ai_pipeline_core/__init__.py,sha256=xjZh6D4fkepTm3LK42qaTpNHgN2meK4uthrW9pjwLfE,779
+ai_pipeline_core/__init__.py,sha256=2zVmBkQNdYcy2mQmw8qO1et3a5pv6KMTXfxrcTfA7EM,779
 ai_pipeline_core/exceptions.py,sha256=_vW0Hbw2LGb5tcVvH0YzTKMff7QOPfCRr3w-w_zPyCE,968
 ai_pipeline_core/prompt_manager.py,sha256=XmNUdMIC0WrE9fF0LIcfozAKOGrlYwj8AfXvCndIH-o,4693
 ai_pipeline_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_pipeline_core/settings.py,sha256=Zl2BPa6IHzh-B5V7cg5mtySr1dhWZQYYKxXz3BwrHlQ,615
 ai_pipeline_core/tracing.py,sha256=_bijptKWXh7V_xENFQGF11-B70rOGwV6g0qdBoF-VCw,7890
 ai_pipeline_core/documents/__init__.py,sha256=rEnKj-sSlZ9WnFlZAmSGVi1P8vnsHmU9O9_YwtP40ms,242
-ai_pipeline_core/documents/document.py,sha256=cKa9v0GpDYh48Aqj0u_tCwh4_uCcn5YsKzC537f0v0I,10089
+ai_pipeline_core/documents/document.py,sha256=e3IBr0TThucBAaOHvdqv0X--iCcBrqh2jzFTyaOp7O0,12418
 ai_pipeline_core/documents/document_list.py,sha256=HOG_uZDazA9CJB7Lr_tNcDFzb5Ff9RUt0ELWQK_eYNM,4940
 ai_pipeline_core/documents/flow_document.py,sha256=qsV-2JYOMhkvAj7lW54ZNH_4QUclld9h06CoU59tWww,815
-ai_pipeline_core/documents/mime_type.py,sha256=tMWGH9PVmHe6a_IzdaJUqIHf4qnwQOwOCBhsgW2AyTE,2244
+ai_pipeline_core/documents/mime_type.py,sha256=sBhNRoBJQ35JoHWhJzBGpp00WFDfMdEX0JZKKkR7QH0,3371
 ai_pipeline_core/documents/task_document.py,sha256=WjHqtl1d60XFBBqewNRdz1OqBErGI0jRx15oQYCTHo8,907
 ai_pipeline_core/documents/utils.py,sha256=BdE4taSl1vrBhxnFbOP5nDA7lXIcvY__AMRTHoaNb5M,2764
 ai_pipeline_core/flow/__init__.py,sha256=_Sji2yY1ICkvVX6QiiGWKzqIXtg9UAiuvhjHSK_gdO8,57
@@ -23,7 +23,7 @@ ai_pipeline_core/logging/__init__.py,sha256=DOO6ckgnMVXl29Sy7q6jhO-iW96h54pCHQDz
 ai_pipeline_core/logging/logging.yml,sha256=YTW48keO_K5bkkb-KXGM7ZuaYKiquLsjsURei8Ql0V4,1353
 ai_pipeline_core/logging/logging_config.py,sha256=6MBz9nnVNvqiLDoyy9-R3sWkn6927Re5hdz4hwTptpI,4903
 ai_pipeline_core/logging/logging_mixin.py,sha256=RDaR2ju2-vKTJRzXGa0DquGPT8_UxahWjvKJnaD0IV8,7810
-ai_pipeline_core-0.1.5.dist-info/METADATA,sha256=U1OIPjGwAGsuyJ3QnhUtJQWMzj-OqkXDuyH6cW8Dq70,15869
-ai_pipeline_core-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ai_pipeline_core-0.1.5.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.1.5.dist-info/RECORD,,
+ai_pipeline_core-0.1.6.dist-info/METADATA,sha256=BSLr818JTSrsTGPjOEB7bQoCv4q3ep-0YX55UgJRH4s,15869
+ai_pipeline_core-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ai_pipeline_core-0.1.6.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.1.6.dist-info/RECORD,,