cognee 0.3.4.dev4__py3-none-any.whl → 0.3.6__py3-none-any.whl

cognee/infrastructure/loaders/external/advanced_pdf_loader.py

@@ -0,0 +1,244 @@
+"""Advanced PDF loader leveraging unstructured for layout-aware extraction."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+import asyncio
+from cognee.infrastructure.files.storage import get_file_storage, get_storage_config
+from cognee.infrastructure.files.utils.get_file_metadata import get_file_metadata
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.shared.logging_utils import get_logger
+
+from cognee.infrastructure.loaders.external.pypdf_loader import PyPdfLoader
+
+logger = get_logger(__name__)
+
+try:
+    from unstructured.partition.pdf import partition_pdf
+except ImportError as e:
+    logger.info(
+        "unstructured[pdf] not installed, can't use AdvancedPdfLoader, will use PyPdfLoader instead."
+    )
+    raise ImportError from e
+
+
+@dataclass
+class _PageBuffer:
+    page_num: Optional[int]
+    segments: List[str]
+
+
+class AdvancedPdfLoader(LoaderInterface):
+    """
+    PDF loader using unstructured library.
+
+    Extracts text content, images, tables 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 "advanced_pdf_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: Any) -> str:
+        """Load PDF file using unstructured library. If Exception occurs, fallback to PyPDFLoader.
+
+        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
+
+        """
+        try:
+            logger.info(f"Processing PDF: {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: Dict[str, Any] = {
+                "filename": file_path,
+                "strategy": strategy,
+                "infer_table_structure": True,
+                "include_page_breaks": False,
+                "include_metadata": True,
+                **kwargs,
+            }
+            # Use partition to extract elements
+            elements = partition_pdf(**partition_kwargs)
+
+            # Process elements into text content
+            page_contents = self._format_elements_by_page(elements)
+
+            # Check if there is any content
+            if not page_contents:
+                logger.warning(
+                    "AdvancedPdfLoader returned no content. Falling back to PyPDF loader."
+                )
+                return await self._fallback(file_path, **kwargs)
+
+            # Combine all page outputs
+            full_content = "\n".join(page_contents)
+
+            # Store the content
+            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 exc:
+            logger.warning("Failed to process PDF with AdvancedPdfLoader: %s", exc)
+            return await self._fallback(file_path, **kwargs)
+
+    async def _fallback(self, file_path: str, **kwargs: Any) -> str:
+        logger.info("Falling back to PyPDF loader for %s", file_path)
+        fallback_loader = PyPdfLoader()
+        return await fallback_loader.load(file_path, **kwargs)
+
+    def _format_elements_by_page(self, elements: List[Any]) -> List[str]:
+        """Format elements by page."""
+        page_buffers: List[_PageBuffer] = []
+        current_buffer = _PageBuffer(page_num=None, segments=[])
+
+        for element in elements:
+            element_dict = self._safe_to_dict(element)
+            metadata = element_dict.get("metadata", {})
+            page_num = metadata.get("page_number")
+
+            if current_buffer.page_num != page_num:
+                if current_buffer.segments:
+                    page_buffers.append(current_buffer)
+                current_buffer = _PageBuffer(page_num=page_num, segments=[])
+
+            formatted = self._format_element(element_dict)
+
+            if formatted:
+                current_buffer.segments.append(formatted)
+
+        if current_buffer.segments:
+            page_buffers.append(current_buffer)
+
+        page_contents: List[str] = []
+        for buffer in page_buffers:
+            header = f"Page {buffer.page_num}:\n" if buffer.page_num is not None else "Page:"
+            content = header + "\n\n".join(buffer.segments) + "\n"
+            page_contents.append(str(content))
+        return page_contents
+
+    def _format_element(
+        self,
+        element: Dict[str, Any],
+    ) -> str:
+        """Format element."""
+        element_type = element.get("type")
+        text = self._clean_text(element.get("text", ""))
+        metadata = element.get("metadata", {})
+
+        if element_type.lower() == "table":
+            return self._format_table_element(element) or text
+
+        if element_type.lower() == "image":
+            description = text or self._format_image_element(metadata)
+            return description
+
+        # Ignore header and footer
+        if element_type.lower() in ["header", "footer"]:
+            pass
+
+        return text
+
+    def _format_table_element(self, element: Dict[str, Any]) -> str:
+        """Format table element."""
+        metadata = element.get("metadata", {})
+        text = self._clean_text(element.get("text", ""))
+        table_html = metadata.get("text_as_html")
+
+        if table_html:
+            return table_html.strip()
+
+        return text
+
+    def _format_image_element(self, metadata: Dict[str, Any]) -> str:
+        """Format image."""
+        placeholder = "[Image omitted]"
+        image_text = placeholder
+        coordinates = metadata.get("coordinates", {})
+        points = coordinates.get("points") if isinstance(coordinates, dict) else None
+        if points and isinstance(points, tuple) and len(points) == 4:
+            leftup = points[0]
+            rightdown = points[3]
+            if (
+                isinstance(leftup, tuple)
+                and isinstance(rightdown, tuple)
+                and len(leftup) == 2
+                and len(rightdown) == 2
+            ):
+                image_text = f"{placeholder} (bbox=({leftup[0]}, {leftup[1]}, {rightdown[0]}, {rightdown[1]}))"
+
+            layout_width = coordinates.get("layout_width")
+            layout_height = coordinates.get("layout_height")
+            system = coordinates.get("system")
+            if layout_width and layout_height and system:
+                image_text = (
+                    image_text
+                    + f", system={system}, layout_width={layout_width}, layout_height={layout_height}))"
+                )
+
+        return image_text
+
+    def _safe_to_dict(self, element: Any) -> Dict[str, Any]:
+        """Safe to dict."""
+        try:
+            if hasattr(element, "to_dict"):
+                return element.to_dict()
+        except Exception:
+            pass
+        fallback_type = getattr(element, "category", None)
+        if not fallback_type:
+            fallback_type = getattr(element, "__class__", type("", (), {})).__name__
+
+        return {
+            "type": fallback_type,
+            "text": getattr(element, "text", ""),
+            "metadata": getattr(element, "metadata", {}),
+        }
+
+    def _clean_text(self, value: Any) -> str:
+        if value is None:
+            return ""
+        return str(value).replace("\xa0", " ").strip()
+
+
+if __name__ == "__main__":
+    loader = AdvancedPdfLoader()
+    asyncio.run(
+        loader.load(
+            "/Users/xiaotao/work/cognee/cognee/infrastructure/loaders/external/attention_is_all_you_need.pdf"
+        )
+    )

cognee/infrastructure/loaders/supported_loaders.py

@@ -16,3 +16,10 @@ try:
     supported_loaders[UnstructuredLoader.loader_name] = UnstructuredLoader
 except ImportError:
     pass
+
+try:
+    from cognee.infrastructure.loaders.external import AdvancedPdfLoader
+
+    supported_loaders[AdvancedPdfLoader.loader_name] = AdvancedPdfLoader
+except ImportError:
+    pass

cognee/modules/data/methods/create_authorized_dataset.py

@@ -6,6 +6,15 @@ from .create_dataset import create_dataset
 
 
 async def create_authorized_dataset(dataset_name: str, user: User) -> Dataset:
+    """
+        Create a new dataset and give all permissions on this dataset to the given user.
+    Args:
+        dataset_name: Name of the dataset.
+        user: The user object.
+
+    Returns:
+        Dataset: The new authorized dataset.
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/data/methods/get_authorized_dataset.py

@@ -15,7 +15,7 @@ async def get_authorized_dataset(
     Get a specific dataset with permissions for a user.
 
     Args:
-        user_id (UUID): user id
+        user: User object
         dataset_id (UUID): dataset id
         permission_type (str): permission type(read, write, delete, share), default is read
 

cognee/modules/data/methods/get_authorized_dataset_by_name.py

@@ -11,6 +11,17 @@ from ..models import Dataset
 async def get_authorized_dataset_by_name(
     dataset_name: str, user: User, permission_type: str
 ) -> Optional[Dataset]:
+    """
+    Get a specific dataset with the given name, with permissions for a given user.
+
+    Args:
+        dataset_name: Name of the dataset.
+        user: User object.
+        permission_type (str): permission type(read, write, delete, share), default is read
+
+    Returns:
+        Optional[Dataset]: dataset with permissions
+    """
     authorized_datasets = await get_authorized_existing_datasets([], permission_type, user)
 
     return next((dataset for dataset in authorized_datasets if dataset.name == dataset_name), None)

cognee/modules/data/methods/get_deletion_counts.py

@@ -0,0 +1,92 @@
+from uuid import UUID
+from cognee.cli.exceptions import CliCommandException
+from cognee.infrastructure.databases.exceptions.exceptions import EntityNotFoundError
+from sqlalchemy import select
+from sqlalchemy.sql import func
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Dataset, Data, DatasetData
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_user
+from dataclasses import dataclass
+
+
+@dataclass
+class DeletionCountsPreview:
+    datasets: int = 0
+    data_entries: int = 0
+    users: int = 0
+
+
+async def get_deletion_counts(
+    dataset_name: str = None, user_id: str = None, all_data: bool = False
+) -> DeletionCountsPreview:
+    """
+    Calculates the number of items that will be deleted based on the provided arguments.
+    """
+    counts = DeletionCountsPreview()
+    relational_engine = get_relational_engine()
+    async with relational_engine.get_async_session() as session:
+        if dataset_name:
+            # Find the dataset by name
+            dataset_result = await session.execute(
+                select(Dataset).where(Dataset.name == dataset_name)
+            )
+            dataset = dataset_result.scalar_one_or_none()
+
+            if dataset is None:
+                raise CliCommandException(
+                    f"No Dataset exists with the name {dataset_name}", error_code=1
+                )
+
+            # Count data entries linked to this dataset
+            count_query = (
+                select(func.count())
+                .select_from(DatasetData)
+                .where(DatasetData.dataset_id == dataset.id)
+            )
+            data_entry_count = (await session.execute(count_query)).scalar_one()
+            counts.users = 1
+            counts.datasets = 1
+            counts.entries = data_entry_count
+            return counts
+
+        elif all_data:
+            # Simplified logic: Get total counts directly from the tables.
+            counts.datasets = (
+                await session.execute(select(func.count()).select_from(Dataset))
+            ).scalar_one()
+            counts.entries = (
+                await session.execute(select(func.count()).select_from(Data))
+            ).scalar_one()
+            counts.users = (
+                await session.execute(select(func.count()).select_from(User))
+            ).scalar_one()
+            return counts
+
+        # Placeholder for user_id logic
+        elif user_id:
+            user = None
+            try:
+                user_uuid = UUID(user_id)
+                user = await get_user(user_uuid)
+            except (ValueError, EntityNotFoundError):
+                raise CliCommandException(f"No User exists with ID {user_id}", error_code=1)
+            counts.users = 1
+            # Find all datasets owned by this user
+            datasets_query = select(Dataset).where(Dataset.owner_id == user.id)
+            user_datasets = (await session.execute(datasets_query)).scalars().all()
+            dataset_count = len(user_datasets)
+            counts.datasets = dataset_count
+            if dataset_count > 0:
+                dataset_ids = [d.id for d in user_datasets]
+                # Count all data entries across all of the user's datasets
+                data_count_query = (
+                    select(func.count())
+                    .select_from(DatasetData)
+                    .where(DatasetData.dataset_id.in_(dataset_ids))
+                )
+                data_entry_count = (await session.execute(data_count_query)).scalar_one()
+                counts.entries = data_entry_count
+            else:
+                counts.entries = 0
+            return counts

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -161,7 +161,7 @@ class CogneeGraph(CogneeAbstractGraph):
                 edge_distances = await vector_engine.search(
                     collection_name="EdgeType_relationship_name",
                     query_vector=query_vector,
-                    limit=0,
+                    limit=None,
                 )
                 projection_time = time.time() - start_time
                 logger.info(

cognee/modules/graph/utils/expand_with_nodes_and_edges.py

@@ -7,8 +7,14 @@ from cognee.modules.engine.utils import (
     generate_node_id,
     generate_node_name,
 )
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+from cognee.modules.ontology.ontology_env_config import get_ontology_env_config
 from cognee.shared.data_models import KnowledgeGraph
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+from cognee.modules.ontology.get_default_ontology_resolver import (
+    get_default_ontology_resolver,
+    get_ontology_resolver_from_env,
+)
 
 
 def _create_node_key(node_id: str, category: str) -> str:
@@ -83,7 +89,7 @@ def _process_ontology_edges(
 
 def _create_type_node(
     node_type: str,
-    ontology_resolver: OntologyResolver,
+    ontology_resolver: RDFLibOntologyResolver,
     added_nodes_map: dict,
     added_ontology_nodes_map: dict,
     name_mapping: dict,
@@ -141,7 +147,7 @@ def _create_entity_node(
     node_name: str,
     node_description: str,
     type_node: EntityType,
-    ontology_resolver: OntologyResolver,
+    ontology_resolver: RDFLibOntologyResolver,
     added_nodes_map: dict,
     added_ontology_nodes_map: dict,
     name_mapping: dict,
@@ -198,7 +204,7 @@ def _create_entity_node(
 def _process_graph_nodes(
     data_chunk: DocumentChunk,
     graph: KnowledgeGraph,
-    ontology_resolver: OntologyResolver,
+    ontology_resolver: RDFLibOntologyResolver,
     added_nodes_map: dict,
     added_ontology_nodes_map: dict,
     name_mapping: dict,
@@ -277,7 +283,7 @@ def _process_graph_edges(
 def expand_with_nodes_and_edges(
     data_chunks: list[DocumentChunk],
     chunk_graphs: list[KnowledgeGraph],
-    ontology_resolver: OntologyResolver = None,
+    ontology_resolver: BaseOntologyResolver = None,
     existing_edges_map: Optional[dict[str, bool]] = None,
 ):
     """
@@ -296,8 +302,8 @@ def expand_with_nodes_and_edges(
         chunk_graphs (list[KnowledgeGraph]): List of knowledge graphs corresponding to each
             data chunk. Each graph contains nodes (entities) and edges (relationships) extracted
             from the chunk content.
-        ontology_resolver (OntologyResolver, optional): Resolver for validating entities and
-            types against an ontology. If None, a default OntologyResolver is created.
+        ontology_resolver (BaseOntologyResolver, optional): Resolver for validating entities and
+            types against an ontology. If None, a default RDFLibOntologyResolver is created.
             Defaults to None.
         existing_edges_map (dict[str, bool], optional): Mapping of existing edge keys to prevent
             duplicate edge creation. Keys are formatted as "{source_id}_{target_id}_{relation}".
@@ -320,7 +326,15 @@ def expand_with_nodes_and_edges(
         existing_edges_map = {}
 
     if ontology_resolver is None:
-        ontology_resolver = OntologyResolver()
+        ontology_config = get_ontology_env_config()
+        if (
+            ontology_config.ontology_file_path
+            and ontology_config.ontology_resolver
+            and ontology_config.matching_strategy
+        ):
+            ontology_resolver = get_ontology_resolver_from_env(**ontology_config.to_dict())
+        else:
+            ontology_resolver = get_default_ontology_resolver()
 
     added_nodes_map = {}
     added_ontology_nodes_map = {}

cognee/modules/graph/utils/retrieve_existing_edges.py

@@ -23,8 +23,6 @@ async def retrieve_existing_edges(
         chunk_graphs (list[KnowledgeGraph]): List of knowledge graphs corresponding to each
             data chunk. Each graph contains nodes (entities) and edges (relationships) that
             were extracted from the chunk content.
-        graph_engine (GraphDBInterface): Interface to the graph database that will be queried
-            to check for existing edges. Must implement the has_edges() method.
 
     Returns:
         dict[str, bool]: A mapping of edge keys to boolean values indicating existence.

cognee/modules/ingestion/data_types/TextData.py

@@ -1,7 +1,6 @@
 from typing import BinaryIO
 from contextlib import asynccontextmanager
 import hashlib
-from cognee.infrastructure.data.utils.extract_keywords import extract_keywords
 from .IngestionData import IngestionData
 
 

cognee/modules/observability/get_observe.py

@@ -9,3 +9,17 @@ def get_observe():
         from langfuse.decorators import observe
 
         return observe
+    elif monitoring == Observer.NONE:
+        # Return a no-op decorator that handles keyword arguments
+        def no_op_decorator(*args, **kwargs):
+            if len(args) == 1 and callable(args[0]) and not kwargs:
+                # Direct decoration: @observe
+                return args[0]
+            else:
+                # Parameterized decoration: @observe(as_type="generation")
+                def decorator(func):
+                    return func
+
+                return decorator
+
+        return no_op_decorator

cognee/modules/observability/observers.py

@@ -4,6 +4,7 @@ from enum import Enum
 class Observer(str, Enum):
     """Monitoring tools"""
 
+    NONE = "none"
     LANGFUSE = "langfuse"
     LLMLITE = "llmlite"
     LANGSMITH = "langsmith"

cognee/modules/ontology/base_ontology_resolver.py

@@ -0,0 +1,42 @@
+from abc import ABC, abstractmethod
+from typing import List, Tuple, Optional
+
+from cognee.modules.ontology.models import AttachedOntologyNode
+from cognee.modules.ontology.matching_strategies import MatchingStrategy, FuzzyMatchingStrategy
+
+
+class BaseOntologyResolver(ABC):
+    """Abstract base class for ontology resolvers."""
+
+    def __init__(self, matching_strategy: Optional[MatchingStrategy] = None):
+        """Initialize the ontology resolver with a matching strategy.
+
+        Args:
+            matching_strategy: The strategy to use for entity matching.
+                              Defaults to FuzzyMatchingStrategy if None.
+        """
+        self.matching_strategy = matching_strategy or FuzzyMatchingStrategy()
+
+    @abstractmethod
+    def build_lookup(self) -> None:
+        """Build the lookup dictionary for ontology entities."""
+        pass
+
+    @abstractmethod
+    def refresh_lookup(self) -> None:
+        """Refresh the lookup dictionary."""
+        pass
+
+    @abstractmethod
+    def find_closest_match(self, name: str, category: str) -> Optional[str]:
+        """Find the closest match for a given name in the specified category."""
+        pass
+
+    @abstractmethod
+    def get_subgraph(
+        self, node_name: str, node_type: str = "individuals", directed: bool = True
+    ) -> Tuple[
+        List[AttachedOntologyNode], List[Tuple[str, str, str]], Optional[AttachedOntologyNode]
+    ]:
+        """Get a subgraph for the given node."""
+        pass

cognee/modules/ontology/get_default_ontology_resolver.py

@@ -0,0 +1,41 @@
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+
+def get_default_ontology_resolver() -> BaseOntologyResolver:
+    return RDFLibOntologyResolver(ontology_file=None, matching_strategy=FuzzyMatchingStrategy())
+
+
+def get_ontology_resolver_from_env(
+    ontology_resolver: str = "", matching_strategy: str = "", ontology_file_path: str = ""
+) -> BaseOntologyResolver:
+    """
+    Create and return an ontology resolver instance based on environment parameters.
+
+    Currently, this function supports only the RDFLib-based ontology resolver
+    with a fuzzy matching strategy.
+
+    Args:
+        ontology_resolver (str): The ontology resolver type to use.
+            Supported value: "rdflib".
+        matching_strategy (str): The matching strategy to apply.
+            Supported value: "fuzzy".
+        ontology_file_path (str): Path to the ontology file required for the resolver.
+
+    Returns:
+        BaseOntologyResolver: An instance of the requested ontology resolver.
+
+    Raises:
+        EnvironmentError: If the provided resolver or strategy is unsupported,
+            or if required parameters are missing.
+    """
+    if ontology_resolver == "rdflib" and matching_strategy == "fuzzy" and ontology_file_path:
+        return RDFLibOntologyResolver(
+            matching_strategy=FuzzyMatchingStrategy(), ontology_file=ontology_file_path
+        )
+    else:
+        raise EnvironmentError(
+            f"Unsupported ontology resolver: {ontology_resolver}. "
+            f"Supported resolvers are: RdfLib with FuzzyMatchingStrategy."
+        )

cognee/modules/ontology/matching_strategies.py

@@ -0,0 +1,53 @@
+import difflib
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+
+class MatchingStrategy(ABC):
+    """Abstract base class for ontology entity matching strategies."""
+
+    @abstractmethod
+    def find_match(self, name: str, candidates: List[str]) -> Optional[str]:
+        """Find the best match for a given name from a list of candidates.
+
+        Args:
+            name: The name to match
+            candidates: List of candidate names to match against
+
+        Returns:
+            The best matching candidate name, or None if no match found
+        """
+        pass
+
+
+class FuzzyMatchingStrategy(MatchingStrategy):
+    """Fuzzy matching strategy using difflib for approximate string matching."""
+
+    def __init__(self, cutoff: float = 0.8):
+        """Initialize fuzzy matching strategy.
+
+        Args:
+            cutoff: Minimum similarity score (0.0 to 1.0) for a match to be considered valid
+        """
+        self.cutoff = cutoff
+
+    def find_match(self, name: str, candidates: List[str]) -> Optional[str]:
+        """Find the closest fuzzy match for a given name.
+
+        Args:
+            name: The normalized name to match
+            candidates: List of normalized candidate names
+
+        Returns:
+            The best matching candidate name, or None if no match meets the cutoff
+        """
+        if not candidates:
+            return None
+
+        # Check for exact match first
+        if name in candidates:
+            return name
+
+        # Find fuzzy match
+        best_match = difflib.get_close_matches(name, candidates, n=1, cutoff=self.cutoff)
+        return best_match[0] if best_match else None