cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/infrastructure/loaders/LoaderEngine.py

@@ -0,0 +1,156 @@
+import filetype
+from typing import Dict, List, Optional, Any
+from .LoaderInterface import LoaderInterface
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class LoaderEngine:
+    """
+    Main loader engine for managing file loaders.
+
+    Follows cognee's adapter pattern similar to database engines,
+    providing a centralized system for file loading operations.
+    """
+
+    def __init__(self):
+        """
+        Initialize the loader engine.
+
+        Args:
+            default_loader_priority: Priority order for loader selection
+        """
+        self._loaders: Dict[str, LoaderInterface] = {}
+        self._extension_map: Dict[str, List[LoaderInterface]] = {}
+        self._mime_type_map: Dict[str, List[LoaderInterface]] = {}
+
+        self.default_loader_priority = [
+            "text_loader",
+            "pypdf_loader",
+            "image_loader",
+            "audio_loader",
+            "unstructured_loader",
+        ]
+
+    def register_loader(self, loader: LoaderInterface) -> bool:
+        """
+        Register a loader with the engine.
+
+        Args:
+            loader: LoaderInterface implementation to register
+
+        Returns:
+            True if loader was registered successfully, False otherwise
+        """
+
+        self._loaders[loader.loader_name] = loader
+
+        # Map extensions to loaders
+        for ext in loader.supported_extensions:
+            ext_lower = ext.lower()
+            if ext_lower not in self._extension_map:
+                self._extension_map[ext_lower] = []
+            self._extension_map[ext_lower].append(loader)
+
+        # Map mime types to loaders
+        for mime_type in loader.supported_mime_types:
+            if mime_type not in self._mime_type_map:
+                self._mime_type_map[mime_type] = []
+            self._mime_type_map[mime_type].append(loader)
+
+        logger.info(f"Registered loader: {loader.loader_name}")
+        return True
+
+    def get_loader(
+        self, file_path: str, preferred_loaders: List[str] = None
+    ) -> Optional[LoaderInterface]:
+        """
+        Get appropriate loader for a file.
+
+        Args:
+            file_path: Path to the file to be processed
+            preferred_loaders: List of preferred loader names to try first
+
+        Returns:
+            LoaderInterface that can handle the file, or None if not found
+        """
+
+        file_info = filetype.guess(file_path)
+
+        # Try preferred loaders first
+        if preferred_loaders:
+            for loader_name in preferred_loaders:
+                if loader_name in self._loaders:
+                    loader = self._loaders[loader_name]
+                    if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
+                        return loader
+                else:
+                    raise ValueError(f"Loader does not exist: {loader_name}")
+
+        # Try default priority order
+        for loader_name in self.default_loader_priority:
+            if loader_name in self._loaders:
+                loader = self._loaders[loader_name]
+                if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
+                    return loader
+            else:
+                raise ValueError(f"Loader does not exist: {loader_name}")
+
+        return None
+
+    async def load_file(
+        self,
+        file_path: str,
+        file_stream: Optional[Any],
+        preferred_loaders: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        """
+        Load file using appropriate loader.
+
+        Args:
+            file_path: Path to the file to be processed
+            preferred_loaders: List of preferred loader names to try first
+            **kwargs: Additional loader-specific configuration
+
+        Raises:
+            ValueError: If no suitable loader is found
+            Exception: If file processing fails
+        """
+        loader = self.get_loader(file_path, preferred_loaders)
+        if not loader:
+            raise ValueError(f"No loader found for file: {file_path}")
+
+        logger.debug(f"Loading {file_path} with {loader.loader_name}")
+        # TODO: loading needs to be reworked to work with both file streams and file locations
+        return await loader.load(file_path, **kwargs)
+
+    def get_available_loaders(self) -> List[str]:
+        """
+        Get list of available loader names.
+
+        Returns:
+            List of registered loader names
+        """
+        return list(self._loaders.keys())
+
+    def get_loader_info(self, loader_name: str) -> Dict[str, any]:
+        """
+        Get information about a specific loader.
+
+        Args:
+            loader_name: Name of the loader to inspect
+
+        Returns:
+            Dictionary containing loader information
+        """
+        if loader_name not in self._loaders:
+            return {}
+
+        loader = self._loaders[loader_name]
+        return {
+            "name": loader.loader_name,
+            "extensions": loader.supported_extensions,
+            "mime_types": loader.supported_mime_types,
+        }

cognee/infrastructure/loaders/LoaderInterface.py

@@ -0,0 +1,73 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Any
+
+
+class LoaderInterface(ABC):
+    """
+    Base interface for all file loaders in cognee.
+
+    This interface follows cognee's established pattern for database adapters,
+    ensuring consistent behavior across all loader implementations.
+    """
+
+    @property
+    @abstractmethod
+    def supported_extensions(self) -> List[str]:
+        """
+        List of file extensions this loader supports.
+
+        Returns:
+            List of extensions including the dot (e.g., ['.txt', '.md'])
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def supported_mime_types(self) -> List[str]:
+        """
+        List of MIME types this loader supports.
+
+        Returns:
+            List of MIME type strings (e.g., ['text/plain', 'application/pdf'])
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def loader_name(self) -> str:
+        """
+        Unique name identifier for this loader.
+
+        Returns:
+            String identifier used for registration and configuration
+        """
+        pass
+
+    @abstractmethod
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """
+        Check if this loader can handle the given file.
+
+        Args:
+            extension: File extension
+            mime_type: MIME type of the file
+
+        Returns:
+            True if this loader can process the file, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    async def load(self, file_path: str, file_stream: Optional[Any] = None, **kwargs):
+        """
+        Load and process the file, returning standardized result.
+
+        Args:
+            file_path: Path to the file to be processed
+            file_stream: If file stream is provided it will be used to process file instead
+            **kwargs: Additional loader-specific configuration
+
+        Raises:
+            Exception: If file cannot be processed
+        """
+        pass

cognee/infrastructure/loaders/__init__.py

@@ -0,0 +1,18 @@
+"""
+File loader infrastructure for cognee.
+
+This package provides a plugin-based system for loading different file formats
+into cognee, following the same patterns as database adapters.
+
+Main exports:
+- get_loader_engine(): Factory function to get configured loader engine
+- use_loader(): Register custom loaders at runtime
+- LoaderInterface: Base interface for implementing loaders
+- LoaderResult, ContentType: Data models for loader results
+"""
+
+from .get_loader_engine import get_loader_engine
+from .use_loader import use_loader
+from .LoaderInterface import LoaderInterface
+
+__all__ = ["get_loader_engine", "use_loader", "LoaderInterface"]

cognee/infrastructure/loaders/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core loader implementations that are always available."""
+
+from .text_loader import TextLoader
+from .audio_loader import AudioLoader
+from .image_loader import ImageLoader
+
+__all__ = ["TextLoader", "AudioLoader", "ImageLoader"]

cognee/infrastructure/loaders/core/audio_loader.py

@@ -0,0 +1,98 @@
+import os
+from typing import List
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+
+
+class AudioLoader(LoaderInterface):
+    """
+    Core text file loader that handles basic text file formats.
+
+    This loader is always available and serves as the fallback for
+    text-based files when no specialized loader is available.
+    """
+
+    @property
+    def supported_extensions(self) -> List[str]:
+        """Supported text file extensions."""
+        return [
+            "aac",  # Audio documents
+            "mid",
+            "mp3",
+            "m4a",
+            "ogg",
+            "flac",
+            "wav",
+            "amr",
+            "aiff",
+        ]
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        """Supported MIME types for text content."""
+        return [
+            "audio/aac",
+            "audio/midi",
+            "audio/mpeg",
+            "audio/mp4",
+            "audio/ogg",
+            "audio/flac",
+            "audio/wav",
+            "audio/amr",
+            "audio/aiff",
+        ]
+
+    @property
+    def loader_name(self) -> str:
+        """Unique identifier for this loader."""
+        return "audio_loader"
+
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """
+        Check if this loader can handle the given file.
+
+        Args:
+            extension: File extension
+            mime_type: Optional MIME type
+
+        Returns:
+            True if file can be handled, False otherwise
+        """
+        if extension in self.supported_extensions and mime_type in self.supported_mime_types:
+            return True
+        return False
+
+    async def load(self, file_path: str, **kwargs):
+        """
+        Load and process the audio file.
+
+        Args:
+            file_path: Path to the file to load
+            **kwargs: Additional configuration (unused)
+
+        Returns:
+            LoaderResult containing the file content and metadata
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            OSError: If file cannot be read
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        with open(file_path, "rb") as f:
+            file_metadata = await get_file_metadata(f)
+        # Name ingested file of current loader based on original file content hash
+        storage_file_name = "text_" + file_metadata["content_hash"] + ".txt"
+
+        result = await LLMGateway.create_transcript(file_path)
+
+        storage_config = get_storage_config()
+        data_root_directory = storage_config["data_root_directory"]
+        storage = get_file_storage(data_root_directory)
+
+        full_file_path = await storage.store(storage_file_name, result.text)
+
+        return full_file_path

cognee/infrastructure/loaders/core/image_loader.py

@@ -0,0 +1,114 @@
+import os
+from typing import List
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+
+
+class ImageLoader(LoaderInterface):
+    """
+    Core image file loader that handles basic image file formats.
+    """
+
+    @property
+    def supported_extensions(self) -> List[str]:
+        """Supported text file extensions."""
+        return [
+            "png",
+            "dwg",
+            "xcf",
+            "jpg",
+            ".jpe",
+            ".jpeg",
+            "jpx",
+            "apng",
+            "gif",
+            "webp",
+            "cr2",
+            "tif",
+            "tiff",
+            "bmp",
+            "jxr",
+            "psd",
+            "ico",
+            "heic",
+            "avif",
+        ]
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        """Supported MIME types for text content."""
+        return [
+            "image/png",
+            "image/vnd.dwg",
+            "image/x-xcf",
+            "image/jpeg",
+            "image/jpx",
+            "image/apng",
+            "image/gif",
+            "image/webp",
+            "image/x-canon-cr2",
+            "image/tiff",
+            "image/bmp",
+            "image/jxr",
+            "image/vnd.adobe.photoshop",
+            "image/vnd.microsoft.icon",
+            "image/heic",
+            "image/avif",
+        ]
+
+    @property
+    def loader_name(self) -> str:
+        """Unique identifier for this loader."""
+        return "image_loader"
+
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """
+        Check if this loader can handle the given file.
+
+        Args:
+            extension: File extension
+            mime_type: Optional MIME type
+
+        Returns:
+            True if file can be handled, False otherwise
+        """
+        if extension in self.supported_extensions and mime_type in self.supported_mime_types:
+            return True
+
+        return False
+
+    async def load(self, file_path: str, **kwargs):
+        """
+        Load and process the image file.
+
+        Args:
+            file_path: Path to the file to load
+            **kwargs: Additional configuration (unused)
+
+        Returns:
+            LoaderResult containing the file content and metadata
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            UnicodeDecodeError: If file cannot be decoded with specified encoding
+            OSError: If file cannot be read
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        with open(file_path, "rb") as f:
+            file_metadata = await get_file_metadata(f)
+        # Name ingested file of current loader based on original file content hash
+        storage_file_name = "text_" + file_metadata["content_hash"] + ".txt"
+
+        result = await LLMGateway.transcribe_image(file_path)
+
+        storage_config = get_storage_config()
+        data_root_directory = storage_config["data_root_directory"]
+        storage = get_file_storage(data_root_directory)
+
+        full_file_path = await storage.store(storage_file_name, result.choices[0].message.content)
+
+        return full_file_path

cognee/infrastructure/loaders/core/text_loader.py

@@ -0,0 +1,90 @@
+import os
+from typing import List
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+
+
+class TextLoader(LoaderInterface):
+    """
+    Core text file loader that handles basic text file formats.
+
+    This loader is always available and serves as the fallback for
+    text-based files when no specialized loader is available.
+    """
+
+    @property
+    def supported_extensions(self) -> List[str]:
+        """Supported text file extensions."""
+        return ["txt", "md", "csv", "json", "xml", "yaml", "yml", "log"]
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        """Supported MIME types for text content."""
+        return [
+            "text/plain",
+            "text/markdown",
+            "text/csv",
+            "application/json",
+            "text/xml",
+            "application/xml",
+            "text/yaml",
+            "application/yaml",
+        ]
+
+    @property
+    def loader_name(self) -> str:
+        """Unique identifier for this loader."""
+        return "text_loader"
+
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """
+        Check if this loader can handle the given file.
+
+        Args:
+            extension: File extension
+            mime_type: Optional MIME type
+
+        Returns:
+            True if file can be handled, False otherwise
+        """
+        if extension in self.supported_extensions and mime_type in self.supported_mime_types:
+            return True
+
+        return False
+
+    async def load(self, file_path: str, encoding: str = "utf-8", **kwargs):
+        """
+        Load and process the text file.
+
+        Args:
+            file_path: Path to the file to load
+            encoding: Text encoding to use (default: utf-8)
+            **kwargs: Additional configuration (unused)
+
+        Returns:
+            LoaderResult containing the file content and metadata
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            UnicodeDecodeError: If file cannot be decoded with specified encoding
+            OSError: If file cannot be read
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        with open(file_path, "rb") as f:
+            file_metadata = await get_file_metadata(f)
+        # Name ingested file of current loader based on original file content hash
+        storage_file_name = "text_" + file_metadata["content_hash"] + ".txt"
+
+        with open(file_path, "r", encoding=encoding) as f:
+            content = f.read()
+
+        storage_config = get_storage_config()
+        data_root_directory = storage_config["data_root_directory"]
+        storage = get_file_storage(data_root_directory)
+
+        full_file_path = await storage.store(storage_file_name, content)
+
+        return full_file_path

cognee/infrastructure/loaders/create_loader_engine.py

@@ -0,0 +1,32 @@
+from .LoaderEngine import LoaderEngine
+from .supported_loaders import supported_loaders
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+def create_loader_engine() -> LoaderEngine:
+    """
+    Create loader engine with given configuration.
+
+    Follows cognee's pattern for engine creation functions used
+    in database adapters.
+
+    Args:
+        default_loader_priority: Priority order for loader selection
+
+    Returns:
+        Configured LoaderEngine instance
+    """
+    engine = LoaderEngine()
+
+    # Register supported loaders from registry
+    for loader_name, loader_class in supported_loaders.items():
+        try:
+            loader_instance = loader_class()
+            engine.register_loader(loader_instance)
+        except Exception as e:
+            # Log but don't fail - allow engine to continue with other loaders
+            logger.warning(f"Failed to register loader {loader_name}: {e}")
+
+    return engine

cognee/infrastructure/loaders/external/__init__.py

@@ -0,0 +1,22 @@
+"""
+External loader implementations for cognee.
+
+This module contains loaders that depend on external libraries:
+- pypdf_loader: PDF processing using pypdf
+- unstructured_loader: Document processing using unstructured
+- dlt_loader: Data lake/warehouse integration using DLT
+
+These loaders are optional and only available if their dependencies are installed.
+"""
+
+from .pypdf_loader import PyPdfLoader
+
+__all__ = ["PyPdfLoader"]
+
+# Conditional imports based on dependency availability
+try:
+    from .unstructured_loader import UnstructuredLoader
+
+    __all__.append("UnstructuredLoader")
+except ImportError:
+    pass