cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__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,14 +17,23 @@ 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, default=None)
+    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)  # File size in bytes
+    data_size = Column(Integer, nullable=True)  # File size in bytes
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=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/data/processing/document_types/UnstructuredDocument.py

@@ -18,11 +18,8 @@ class UnstructuredDocument(Document):
             except ModuleNotFoundError:
                 raise UnstructuredLibraryImportError
 
-            if self.raw_data_location.startswith("s3://"):
-                async with open_data_file(self.raw_data_location, mode="rb") as f:
-                    elements = partition(file=f, content_type=self.mime_type)
-            else:
-                elements = partition(self.raw_data_location, content_type=self.mime_type)
+            async with open_data_file(self.raw_data_location, mode="rb") as f:
+                elements = partition(file=f, content_type=self.mime_type)
 
             in_memory_file = StringIO("\n\n".join([str(el) for el in elements]))
             in_memory_file.seek(0)

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

@@ -1,3 +1,4 @@
+import time
 from cognee.shared.logging_utils import get_logger
 from typing import List, Dict, Union, Optional, Type
 
@@ -8,7 +9,7 @@ from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
 from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph
 import heapq
 
-logger = get_logger()
+logger = get_logger("CogneeGraph")
 
 
 class CogneeGraph(CogneeAbstractGraph):
@@ -66,7 +67,13 @@ class CogneeGraph(CogneeAbstractGraph):
     ) -> None:
         if node_dimension < 1 or edge_dimension < 1:
             raise InvalidValueError(message="Dimensions must be positive integers")
+
         try:
+            import time
+
+            start_time = time.time()
+
+            # Determine projection strategy
             if node_type is not None and node_name is not None:
                 nodes_data, edges_data = await adapter.get_nodeset_subgraph(
                     node_type=node_type, node_name=node_name
@@ -83,16 +90,17 @@ class CogneeGraph(CogneeAbstractGraph):
                 nodes_data, edges_data = await adapter.get_filtered_graph_data(
                     attribute_filters=memory_fragment_filter
                 )
-
                 if not nodes_data or not edges_data:
                     raise EntityNotFoundError(
                         message="Empty filtered graph projected from the database."
                     )
 
+            # Process nodes
             for node_id, properties in nodes_data:
                 node_attributes = {key: properties.get(key) for key in node_properties_to_project}
                 self.add_node(Node(str(node_id), node_attributes, dimension=node_dimension))
 
+            # Process edges
             for source_id, target_id, relationship_type, properties in edges_data:
                 source_node = self.get_node(str(source_id))
                 target_node = self.get_node(str(target_id))
@@ -113,17 +121,23 @@ class CogneeGraph(CogneeAbstractGraph):
 
                     source_node.add_skeleton_edge(edge)
                     target_node.add_skeleton_edge(edge)
-
                 else:
                     raise EntityNotFoundError(
                         message=f"Edge references nonexistent nodes: {source_id} -> {target_id}"
                     )
 
-        except (ValueError, TypeError) as e:
-            print(f"Error projecting graph: {e}")
-            raise e
+            # Final statistics
+            projection_time = time.time() - start_time
+            logger.info(
+                f"Graph projection completed: {len(self.nodes)} nodes, {len(self.edges)} edges in {projection_time:.2f}s"
+            )
+
+        except Exception as e:
+            logger.error(f"Error during graph projection: {str(e)}")
+            raise
 
     async def map_vector_distances_to_graph_nodes(self, node_distances) -> None:
+        mapped_nodes = 0
         for category, scored_results in node_distances.items():
             for scored_result in scored_results:
                 node_id = str(scored_result.id)
@@ -131,48 +145,44 @@ class CogneeGraph(CogneeAbstractGraph):
                 node = self.get_node(node_id)
                 if node:
                     node.add_attribute("vector_distance", score)
+                    mapped_nodes += 1
 
-    async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> None:
+    async def map_vector_distances_to_graph_edges(
+        self, vector_engine, query_vector, edge_distances
+    ) -> None:
         try:
-            query_vector = await vector_engine.embed_data([query])
-            query_vector = query_vector[0]
             if query_vector is None or len(query_vector) == 0:
                 raise ValueError("Failed to generate query embedding.")
 
-            edge_distances = await vector_engine.search(
-                collection_name="EdgeType_relationship_name",
-                query_text=query,
-                limit=0,
-            )
+            if edge_distances is None:
+                start_time = time.time()
+                edge_distances = await vector_engine.search(
+                    collection_name="EdgeType_relationship_name",
+                    query_vector=query_vector,
+                    limit=0,
+                )
+                projection_time = time.time() - start_time
+                logger.info(
+                    f"Edge collection distances were calculated separately from nodes in {projection_time:.2f}s"
+                )
 
             embedding_map = {result.payload["text"]: result.score for result in edge_distances}
 
             for edge in self.edges:
                 relationship_type = edge.attributes.get("relationship_type")
-                if not relationship_type or relationship_type not in embedding_map:
-                    print(f"Edge {edge} has an unknown or missing relationship type.")
-                    continue
-
-                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:
-            print(f"Error mapping vector distances to edges: {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 = []
-        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)
+        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
 
-        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()
@@ -33,7 +39,7 @@ async def get_formatted_graph_data(dataset_id: UUID, user_id: UUID):
                 lambda edge: {
                     "source": str(edge[0]),
                     "target": str(edge[1]),
-                    "label": edge[2],
+                    "label": str(edge[2]),
                 },
                 edges,
             )