cognee 0.2.2.dev0__py3-none-any.whl → 0.2.3__py3-none-any.whl

cognee/infrastructure/loaders/external/pypdf_loader.py

@@ -0,0 +1,96 @@
+from typing import List
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+
+logger = get_logger(__name__)
+
+
+class PyPdfLoader(LoaderInterface):
+    """
+    PDF loader using pypdf library.
+
+    Extracts text content from PDF files page by page, providing
+    structured page information and handling PDF-specific errors.
+    """
+
+    @property
+    def supported_extensions(self) -> List[str]:
+        return ["pdf"]
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        return ["application/pdf"]
+
+    @property
+    def loader_name(self) -> str:
+        return "pypdf_loader"
+
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """Check if file can be handled by this loader."""
+        # Check file extension
+        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, strict: bool = False, **kwargs) -> str:
+        """
+        Load PDF file and extract text content.
+
+        Args:
+            file_path: Path to the PDF file
+            strict: Whether to use strict mode for PDF reading
+            **kwargs: Additional arguments
+
+        Returns:
+            LoaderResult with extracted text content and metadata
+
+        Raises:
+            ImportError: If pypdf is not installed
+            Exception: If PDF processing fails
+        """
+        try:
+            from pypdf import PdfReader
+        except ImportError as e:
+            raise ImportError(
+                "pypdf is required for PDF processing. Install with: pip install pypdf"
+            ) from e
+
+        try:
+            with open(file_path, "rb") as file:
+                file_metadata = await get_file_metadata(file)
+                # Name ingested file of current loader based on original file content hash
+                storage_file_name = "text_" + file_metadata["content_hash"] + ".txt"
+
+                logger.info(f"Reading PDF: {file_path}")
+                reader = PdfReader(file, strict=strict)
+
+                content_parts = []
+                page_texts = []
+
+                for page_num, page in enumerate(reader.pages, 1):
+                    try:
+                        page_text = page.extract_text()
+                        if page_text.strip():  # Only add non-empty pages
+                            page_texts.append(page_text)
+                            content_parts.append(f"Page {page_num}:\n{page_text}\n")
+                    except Exception as e:
+                        logger.warning(f"Failed to extract text from page {page_num}: {e}")
+                        continue
+
+                # Combine all content
+                full_content = "\n".join(content_parts)
+
+                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, full_content)
+
+                return full_file_path
+
+        except Exception as e:
+            logger.error(f"Failed to process PDF {file_path}: {e}")
+            raise Exception(f"PDF processing failed: {e}") from e

cognee/infrastructure/loaders/external/unstructured_loader.py

@@ -0,0 +1,127 @@
+from typing import List
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+
+logger = get_logger(__name__)
+
+
+class UnstructuredLoader(LoaderInterface):
+    """
+    Document loader using the unstructured library.
+
+    Handles various document formats including docx, pptx, xlsx, odt, etc.
+    Uses the unstructured library's auto-partition functionality.
+    """
+
+    @property
+    def supported_extensions(self) -> List[str]:
+        return [
+            "docx",
+            "doc",
+            "odt",  # Word documents
+            "xlsx",
+            "xls",
+            "ods",  # Spreadsheets
+            "pptx",
+            "ppt",
+            "odp",  # Presentations
+            "rtf",
+            "html",
+            "htm",  # Rich text and HTML
+            "eml",
+            "msg",  # Email formats
+            "epub",  # eBooks
+        ]
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        return [
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",  # docx
+            "application/msword",  # doc
+            "application/vnd.oasis.opendocument.text",  # odt
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # xlsx
+            "application/vnd.ms-excel",  # xls
+            "application/vnd.oasis.opendocument.spreadsheet",  # ods
+            "application/vnd.openxmlformats-officedocument.presentationml.presentation",  # pptx
+            "application/vnd.ms-powerpoint",  # ppt
+            "application/vnd.oasis.opendocument.presentation",  # odp
+            "application/rtf",  # rtf
+            "text/html",  # html
+            "message/rfc822",  # eml
+            "application/epub+zip",  # epub
+        ]
+
+    @property
+    def loader_name(self) -> str:
+        return "unstructured_loader"
+
+    def can_handle(self, extension: str, mime_type: str) -> bool:
+        """Check if file can be handled by this loader."""
+        # Check file extension
+        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, strategy: str = "auto", **kwargs):
+        """
+        Load document using unstructured library.
+
+        Args:
+            file_path: Path to the document file
+            strategy: Partitioning strategy ("auto", "fast", "hi_res", "ocr_only")
+            **kwargs: Additional arguments passed to unstructured partition
+
+        Returns:
+            LoaderResult with extracted text content and metadata
+
+        Raises:
+            ImportError: If unstructured is not installed
+            Exception: If document processing fails
+        """
+        try:
+            from unstructured.partition.auto import partition
+        except ImportError as e:
+            raise ImportError(
+                "unstructured is required for document processing. "
+                "Install with: pip install unstructured"
+            ) from e
+
+        try:
+            logger.info(f"Processing document: {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"
+
+            # Set partitioning parameters
+            partition_kwargs = {"filename": file_path, "strategy": strategy, **kwargs}
+
+            # Use partition to extract elements
+            elements = partition(**partition_kwargs)
+
+            # Process elements into text content
+            text_parts = []
+
+            for element in elements:
+                element_text = str(element).strip()
+                if element_text:
+                    text_parts.append(element_text)
+
+            # Combine all text content
+            full_content = "\n\n".join(text_parts)
+
+            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, full_content)
+
+            return full_file_path
+
+        except Exception as e:
+            logger.error(f"Failed to process document {file_path}: {e}")
+            raise Exception(f"Document processing failed: {e}") from e

cognee/infrastructure/loaders/get_loader_engine.py

@@ -0,0 +1,18 @@
+from functools import lru_cache
+from .LoaderEngine import LoaderEngine
+from .create_loader_engine import create_loader_engine
+
+
+@lru_cache
+def get_loader_engine() -> LoaderEngine:
+    """
+    Factory function to get loader engine.
+
+    Follows cognee's pattern with @lru_cache for efficient reuse
+    of engine instances. Configuration is loaded from environment
+    variables and settings.
+
+    Returns:
+        Cached LoaderEngine instance configured with current settings
+    """
+    return create_loader_engine()

cognee/infrastructure/loaders/supported_loaders.py

@@ -0,0 +1,18 @@
+from cognee.infrastructure.loaders.external import PyPdfLoader
+from cognee.infrastructure.loaders.core import TextLoader, AudioLoader, ImageLoader
+
+# Registry for loader implementations
+supported_loaders = {
+    PyPdfLoader.loader_name: PyPdfLoader,
+    TextLoader.loader_name: TextLoader,
+    ImageLoader.loader_name: ImageLoader,
+    AudioLoader.loader_name: AudioLoader,
+}
+
+# Try adding optional loaders
+try:
+    from cognee.infrastructure.loaders.external import UnstructuredLoader
+
+    supported_loaders[UnstructuredLoader.loader_name] = UnstructuredLoader
+except ImportError:
+    pass

cognee/infrastructure/loaders/use_loader.py

@@ -0,0 +1,21 @@
+from .supported_loaders import supported_loaders
+
+
+def use_loader(loader_name: str, loader_class):
+    """
+    Register a loader at runtime.
+
+    This allows external packages and custom loaders to be registered
+    into the loader system.
+
+    Args:
+        loader_name: Unique name for the loader
+        loader_class: Loader class implementing LoaderInterface
+
+    Example:
+        from cognee.infrastructure.loaders import use_loader
+        from my_package import MyCustomLoader
+
+        use_loader("my_custom_loader", MyCustomLoader)
+    """
+    supported_loaders[loader_name] = loader_class

cognee/modules/data/methods/__init__.py

@@ -6,6 +6,7 @@ from .get_dataset import get_dataset
 from .get_datasets import get_datasets
 from .get_datasets_by_name import get_datasets_by_name
 from .get_dataset_data import get_dataset_data
+from .get_authorized_dataset import get_authorized_dataset
 from .get_data import get_data
 from .get_unique_dataset_id import get_unique_dataset_id
 from .get_authorized_existing_datasets import get_authorized_existing_datasets

cognee/modules/data/methods/get_authorized_dataset.py

@@ -0,0 +1,23 @@
+from typing import Optional
+from uuid import UUID
+from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
+from ..models import Dataset
+
+
+async def get_authorized_dataset(
+    user_id: UUID, dataset_id: UUID, permission_type="read"
+) -> Optional[Dataset]:
+    """
+    Get a specific dataset with permissions for a user.
+
+    Args:
+        user_id (UUID): user id
+        dataset_id (UUID): dataset id
+        permission_type (str): permission type(read, write, delete, share), default is read
+
+    Returns:
+        Optional[Dataset]: dataset with permissions
+    """
+    datasets = await get_specific_user_permission_datasets(user_id, permission_type, [dataset_id])
+
+    return datasets[0] if datasets else None

cognee/modules/data/models/Data.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timezone
 from uuid import uuid4
 from sqlalchemy import UUID, Column, DateTime, String, JSON, Integer
+from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.orm import relationship
 
 from cognee.infrastructure.databases.relational import Base
@@ -16,12 +17,21 @@ class Data(Base):
     name = Column(String)
     extension = Column(String)
     mime_type = Column(String)
+    original_extension = Column(String, nullable=True)
+    original_mime_type = Column(String, nullable=True)
+    loader_engine = Column(String)
     raw_data_location = Column(String)
+    original_data_location = Column(String)
     owner_id = Column(UUID, index=True)
     tenant_id = Column(UUID, index=True, nullable=True)
     content_hash = Column(String)
+    raw_content_hash = Column(String)
     external_metadata = Column(JSON)
-    node_set = Column(JSON, nullable=True)  # Store NodeSet as JSON list of strings
+    # Store NodeSet as JSON list of strings
+    node_set = Column(JSON, nullable=True)
+    # MutableDict allows SQLAlchemy to notice key-value pair changes, without it changing a value for a key
+    # wouldn't be noticed when commiting a database session
+    pipeline_status = Column(MutableDict.as_mutable(JSON))
     token_count = Column(Integer)
     data_size = Column(Integer, nullable=True)  # File size in bytes
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))

cognee/modules/data/processing/document_types/AudioDocument.py

@@ -1,5 +1,5 @@
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
 from cognee.modules.chunking.Chunker import Chunker
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 from .Document import Document
 
@@ -8,7 +8,7 @@ class AudioDocument(Document):
     type: str = "audio"
 
     async def create_transcript(self):
-        result = await get_llm_client().create_transcript(self.raw_data_location)
+        result = await LLMGateway.create_transcript(self.raw_data_location)
         return result.text
 
     async def read(self, chunker_cls: Chunker, max_chunk_size: int):

cognee/modules/data/processing/document_types/ImageDocument.py

@@ -1,4 +1,4 @@
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.chunking.Chunker import Chunker
 
 from .Document import Document
@@ -8,7 +8,7 @@ class ImageDocument(Document):
     type: str = "image"
 
     async def transcribe_image(self):
-        result = await get_llm_client().transcribe_image(self.raw_data_location)
+        result = await LLMGateway.transcribe_image(self.raw_data_location)
         return result.choices[0].message.content
 
     async def read(self, chunker_cls: Chunker, max_chunk_size: int):

cognee/modules/data/processing/document_types/PdfDocument.py

@@ -5,7 +5,6 @@ from cognee.modules.chunking.Chunker import Chunker
 from cognee.infrastructure.files.utils.open_data_file import open_data_file
 
 from .Document import Document
-from .exceptions.exceptions import PyPdfInternalError
 
 logger = get_logger("PDFDocument")
 
@@ -17,18 +16,12 @@ class PdfDocument(Document):
         async with open_data_file(self.raw_data_location, mode="rb") as stream:
             logger.info(f"Reading PDF: {self.raw_data_location}")
 
-            try:
-                file = PdfReader(stream, strict=False)
-            except Exception:
-                raise PyPdfInternalError()
+            file = PdfReader(stream, strict=False)
 
             async def get_text():
-                try:
-                    for page in file.pages:
-                        page_text = page.extract_text()
-                        yield page_text
-                except Exception:
-                    raise PyPdfInternalError()
+                for page in file.pages:
+                    page_text = page.extract_text()
+                    yield page_text
 
             chunker = chunker_cls(self, get_text=get_text, max_chunk_size=max_chunk_size)
 

cognee/modules/engine/utils/generate_edge_id.py

@@ -0,0 +1,5 @@
+from uuid import NAMESPACE_OID, uuid5
+
+
+def generate_edge_id(edge_id: str) -> str:
+    return uuid5(NAMESPACE_OID, edge_id.lower().replace(" ", "_").replace("'", ""))

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -170,28 +170,19 @@ class CogneeGraph(CogneeAbstractGraph):
 
             for edge in self.edges:
                 relationship_type = edge.attributes.get("relationship_type")
-                if relationship_type and relationship_type in embedding_map:
-                    edge.attributes["vector_distance"] = embedding_map[relationship_type]
+                distance = embedding_map.get(relationship_type, None)
+                if distance is not None:
+                    edge.attributes["vector_distance"] = distance
 
         except Exception as ex:
             logger.error(f"Error mapping vector distances to edges: {str(ex)}")
             raise ex
 
     async def calculate_top_triplet_importances(self, k: int) -> List:
-        min_heap = []
+        def score(edge):
+            n1 = edge.node1.attributes.get("vector_distance", 1)
+            n2 = edge.node2.attributes.get("vector_distance", 1)
+            e = edge.attributes.get("vector_distance", 1)
+            return n1 + n2 + e
 
-        for i, edge in enumerate(self.edges):
-            source_node = self.get_node(edge.node1.id)
-            target_node = self.get_node(edge.node2.id)
-
-            source_distance = source_node.attributes.get("vector_distance", 1) if source_node else 1
-            target_distance = target_node.attributes.get("vector_distance", 1) if target_node else 1
-            edge_distance = edge.attributes.get("vector_distance", 1)
-
-            total_distance = source_distance + target_distance + edge_distance
-
-            heapq.heappush(min_heap, (-total_distance, i, edge))
-            if len(min_heap) > k:
-                heapq.heappop(min_heap)
-
-        return [edge for _, _, edge in sorted(min_heap)]
+        return heapq.nsmallest(k, self.edges, key=score)

cognee/modules/graph/methods/get_formatted_graph_data.py

@@ -1,10 +1,16 @@
 from uuid import UUID
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.context_global_variables import set_database_global_context_variables
+from cognee.modules.data.exceptions.exceptions import DatasetNotFoundError
+from cognee.modules.data.methods import get_authorized_dataset
 
 
 async def get_formatted_graph_data(dataset_id: UUID, user_id: UUID):
-    await set_database_global_context_variables(dataset_id, user_id)
+    dataset = await get_authorized_dataset(user_id, dataset_id)
+    if not dataset:
+        raise DatasetNotFoundError(message="Dataset not found.")
+
+    await set_database_global_context_variables(dataset_id, dataset.owner_id)
 
     graph_client = await get_graph_engine()
     (nodes, edges) = await graph_client.get_graph_data()