cognee 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

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/notebooks/methods/create_notebook.py

@@ -6,6 +6,40 @@ from cognee.infrastructure.databases.relational import with_async_session
 
 from ..models.Notebook import Notebook, NotebookCell
 
+TUTORIAL_NOTEBOOK_NAME = "Python Development with Cognee Tutorial 🧠"
+
+
+async def _create_tutorial_notebook(
+    user_id: UUID, session: AsyncSession, force_refresh: bool = False
+) -> None:
+    """
+    Create the default tutorial notebook for new users.
+    Dynamically fetches from: https://github.com/topoteretes/cognee/blob/notebook_tutorial/notebooks/starter_tutorial.zip
+    """
+    TUTORIAL_ZIP_URL = (
+        "https://github.com/topoteretes/cognee/raw/notebook_tutorial/notebooks/starter_tutorial.zip"
+    )
+
+    try:
+        # Create notebook from remote zip file (includes notebook + data files)
+        notebook = await Notebook.from_ipynb_zip_url(
+            zip_url=TUTORIAL_ZIP_URL,
+            owner_id=user_id,
+            notebook_filename="tutorial.ipynb",
+            name=TUTORIAL_NOTEBOOK_NAME,
+            deletable=False,
+            force=force_refresh,
+        )
+
+        # Add to session and commit
+        session.add(notebook)
+        await session.commit()
+
+    except Exception as e:
+        print(f"Failed to fetch tutorial notebook from {TUTORIAL_ZIP_URL}: {e}")
+
+        raise e
+
 
 @with_async_session
 async def create_notebook(

cognee/modules/notebooks/methods/get_notebooks.py

@@ -1,11 +1,16 @@
 from uuid import UUID
 from typing import List
-from sqlalchemy import select
+from sqlalchemy import select, and_
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import with_async_session
 
 from ..models.Notebook import Notebook
+from .create_notebook import _create_tutorial_notebook, TUTORIAL_NOTEBOOK_NAME
+
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger()
 
 
 @with_async_session
@@ -13,6 +18,27 @@ async def get_notebooks(
     user_id: UUID,
     session: AsyncSession,
 ) -> List[Notebook]:
+    # Check if tutorial notebook already exists for this user
+    tutorial_query = select(Notebook).where(
+        and_(
+            Notebook.owner_id == user_id,
+            Notebook.name == TUTORIAL_NOTEBOOK_NAME,
+            ~Notebook.deletable,
+        )
+    )
+    tutorial_result = await session.execute(tutorial_query)
+    tutorial_notebook = tutorial_result.scalar_one_or_none()
+
+    # If tutorial notebook doesn't exist, create it
+    if tutorial_notebook is None:
+        logger.info(f"Tutorial notebook not found for user {user_id}, creating it")
+        try:
+            await _create_tutorial_notebook(user_id, session, force_refresh=False)
+        except Exception as e:
+            # Log the error but continue to return existing notebooks
+            logger.error(f"Failed to create tutorial notebook for user {user_id}: {e}")
+
+    # Get all notebooks for the user
     result = await session.execute(select(Notebook).where(Notebook.owner_id == user_id))
 
     return list(result.scalars().all())

cognee/modules/notebooks/models/Notebook.py

@@ -1,13 +1,24 @@
 import json
-from typing import List, Literal
+import nbformat
+import asyncio
+from nbformat.notebooknode import NotebookNode
+from typing import List, Literal, Optional, cast, Tuple
 from uuid import uuid4, UUID as UUID_t
 from pydantic import BaseModel, ConfigDict
 from datetime import datetime, timezone
 from fastapi.encoders import jsonable_encoder
 from sqlalchemy import Boolean, Column, DateTime, JSON, UUID, String, TypeDecorator
 from sqlalchemy.orm import mapped_column, Mapped
+from pathlib import Path
 
 from cognee.infrastructure.databases.relational import Base
+from cognee.shared.cache import (
+    download_and_extract_zip,
+    get_tutorial_data_dir,
+    generate_content_hash,
+)
+from cognee.infrastructure.files.storage.get_file_storage import get_file_storage
+from cognee.base_config import get_base_config
 
 
 class NotebookCell(BaseModel):
@@ -51,3 +62,197 @@ class Notebook(Base):
     deletable: Mapped[bool] = mapped_column(Boolean, default=True)
 
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+
+    @classmethod
+    async def from_ipynb_zip_url(
+        cls,
+        zip_url: str,
+        owner_id: UUID_t,
+        notebook_filename: str = "tutorial.ipynb",
+        name: Optional[str] = None,
+        deletable: bool = True,
+        force: bool = False,
+    ) -> "Notebook":
+        """
+        Create a Notebook instance from a remote zip file containing notebook + data files.
+
+        Args:
+            zip_url: Remote URL to fetch the .zip file from
+            owner_id: UUID of the notebook owner
+            notebook_filename: Name of the .ipynb file within the zip
+            name: Optional custom name for the notebook
+            deletable: Whether the notebook can be deleted
+            force: If True, re-download even if already cached
+
+        Returns:
+            Notebook instance
+        """
+        # Generate a cache key based on the zip URL
+        content_hash = generate_content_hash(zip_url, notebook_filename)
+
+        # Download and extract the zip file to tutorial_data/{content_hash}
+        try:
+            extracted_cache_dir = await download_and_extract_zip(
+                url=zip_url,
+                cache_dir_name=f"tutorial_data/{content_hash}",
+                version_or_hash=content_hash,
+                force=force,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to download tutorial zip from {zip_url}") from e
+
+        # Use cache system to access the notebook file
+        from cognee.shared.cache import cache_file_exists, read_cache_file
+
+        notebook_file_path = f"{extracted_cache_dir}/{notebook_filename}"
+
+        # Check if the notebook file exists in cache
+        if not await cache_file_exists(notebook_file_path):
+            raise FileNotFoundError(f"Notebook file '{notebook_filename}' not found in zip")
+
+        # Read and parse the notebook using cache system
+        async with await read_cache_file(notebook_file_path, encoding="utf-8") as f:
+            notebook_content = await asyncio.to_thread(f.read)
+        notebook = cls.from_ipynb_string(notebook_content, owner_id, name, deletable)
+
+        # Update file paths in notebook cells to point to actual cached data files
+        await cls._update_file_paths_in_cells(notebook, extracted_cache_dir)
+
+        return notebook
+
+    @staticmethod
+    async def _update_file_paths_in_cells(notebook: "Notebook", cache_dir: str) -> None:
+        """
+        Update file paths in code cells to use actual cached data files.
+        Works with both local filesystem and S3 storage.
+
+        Args:
+            notebook: Parsed Notebook instance with cells to update
+            cache_dir: Path to the cached tutorial directory containing data files
+        """
+        import re
+        from cognee.shared.cache import list_cache_files, cache_file_exists
+        from cognee.shared.logging_utils import get_logger
+
+        logger = get_logger()
+
+        # Look for data files in the data subdirectory
+        data_dir = f"{cache_dir}/data"
+
+        try:
+            # Get all data files in the cache directory using cache system
+            data_files = {}
+            if await cache_file_exists(data_dir):
+                file_list = await list_cache_files(data_dir)
+            else:
+                file_list = []
+
+            for file_path in file_list:
+                # Extract just the filename
+                filename = file_path.split("/")[-1]
+                # Use the file path as provided by cache system
+                data_files[filename] = file_path
+
+        except Exception as e:
+            # If we can't list files, skip updating paths
+            logger.error(f"Error listing data files in {data_dir}: {e}")
+            return
+
+        # Pattern to match file://data/filename patterns in code cells
+        file_pattern = r'"file://data/([^"]+)"'
+
+        def replace_path(match):
+            filename = match.group(1)
+            if filename in data_files:
+                file_path = data_files[filename]
+                # For local filesystem, preserve file:// prefix
+                if not file_path.startswith("s3://"):
+                    return f'"file://{file_path}"'
+                else:
+                    # For S3, return the S3 URL as-is
+                    return f'"{file_path}"'
+            return match.group(0)  # Keep original if file not found
+
+        # Update only code cells
+        updated_cells = 0
+        for cell in notebook.cells:
+            if cell.type == "code":
+                original_content = cell.content
+                # Update file paths in the cell content
+                cell.content = re.sub(file_pattern, replace_path, cell.content)
+                if original_content != cell.content:
+                    updated_cells += 1
+
+        # Log summary of updates (useful for monitoring)
+        if updated_cells > 0:
+            logger.info(f"Updated file paths in {updated_cells} notebook cells")
+
+    @classmethod
+    def from_ipynb_string(
+        cls,
+        notebook_content: str,
+        owner_id: UUID_t,
+        name: Optional[str] = None,
+        deletable: bool = True,
+    ) -> "Notebook":
+        """
+        Create a Notebook instance from Jupyter notebook string content.
+
+        Args:
+            notebook_content: Raw Jupyter notebook content as string
+            owner_id: UUID of the notebook owner
+            name: Optional custom name for the notebook
+            deletable: Whether the notebook can be deleted
+
+        Returns:
+            Notebook instance ready to be saved to database
+        """
+        # Parse and validate the Jupyter notebook using nbformat
+        # Note: nbformat.reads() has loose typing, so we cast to NotebookNode
+        jupyter_nb = cast(
+            NotebookNode, nbformat.reads(notebook_content, as_version=nbformat.NO_CONVERT)
+        )
+
+        # Convert Jupyter cells to NotebookCell objects
+        cells = []
+        for jupyter_cell in jupyter_nb.cells:
+            # Each cell is also a NotebookNode with dynamic attributes
+            cell = cast(NotebookNode, jupyter_cell)
+            # Skip raw cells as they're not supported in our model
+            if cell.cell_type == "raw":
+                continue
+
+            # Get the source content
+            content = cell.source
+
+            # Generate a name based on content or cell index
+            cell_name = cls._generate_cell_name(cell)
+
+            # Map cell types (jupyter uses "code"/"markdown", we use same)
+            cell_type = "code" if cell.cell_type == "code" else "markdown"
+
+            cells.append(NotebookCell(id=uuid4(), type=cell_type, name=cell_name, content=content))
+
+        # Extract notebook name from metadata if not provided
+        if name is None:
+            kernelspec = jupyter_nb.metadata.get("kernelspec", {})
+            name = kernelspec.get("display_name") or kernelspec.get("name", "Imported Notebook")
+
+        return cls(id=uuid4(), owner_id=owner_id, name=name, cells=cells, deletable=deletable)
+
+    @staticmethod
+    def _generate_cell_name(jupyter_cell: NotebookNode) -> str:
+        """Generate a meaningful name for a notebook cell using nbformat cell."""
+        if jupyter_cell.cell_type == "markdown":
+            # Try to extract a title from markdown headers
+            content = jupyter_cell.source
+
+            lines = content.strip().split("\n")
+            if lines and lines[0].startswith("#"):
+                # Extract header text, clean it up
+                header = lines[0].lstrip("#").strip()
+                return header[:50] if len(header) > 50 else header
+            else:
+                return "Markdown Cell"
+        else:
+            return "Code Cell"

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

cognee/modules/ontology/models.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+
+class AttachedOntologyNode:
+    """Lightweight wrapper to be able to parse any ontology solution and generalize cognee interface."""
+
+    def __init__(self, uri: Any, category: str):
+        self.uri = uri
+        self.name = self._extract_name(uri)
+        self.category = category
+
+    @staticmethod
+    def _extract_name(uri: Any) -> str:
+        uri_str = str(uri)
+        if "#" in uri_str:
+            return uri_str.split("#")[-1]
+        return uri_str.rstrip("/").split("/")[-1]
+
+    def __repr__(self):
+        return f"AttachedOntologyNode(name={self.name}, category={self.category})"

cognee/modules/ontology/ontology_config.py

@@ -0,0 +1,24 @@
+from typing import TypedDict, Optional
+
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+from cognee.modules.ontology.matching_strategies import MatchingStrategy
+
+
+class OntologyConfig(TypedDict, total=False):
+    """Configuration containing ontology resolver.
+
+    Attributes:
+        ontology_resolver: The ontology resolver instance to use
+    """
+
+    ontology_resolver: Optional[BaseOntologyResolver]
+
+
+class Config(TypedDict, total=False):
+    """Top-level configuration dictionary.
+
+    Attributes:
+        ontology_config: Configuration containing ontology resolver
+    """
+
+    ontology_config: Optional[OntologyConfig]

cognee/modules/ontology/ontology_env_config.py

@@ -0,0 +1,45 @@
+"""This module contains the configuration for ontology handling."""
+
+from functools import lru_cache
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class OntologyEnvConfig(BaseSettings):
+    """
+    Represents the configuration for ontology handling, including parameters for
+    ontology file storage and resolution/matching strategies.
+
+    Public methods:
+    - to_dict
+
+    Instance variables:
+    - ontology_resolver
+    - ontology_matching
+    - ontology_file_path
+    - model_config
+    """
+
+    ontology_resolver: str = "rdflib"
+    matching_strategy: str = "fuzzy"
+    ontology_file_path: str = ""
+
+    model_config = SettingsConfigDict(env_file=".env", extra="allow", populate_by_name=True)
+
+    def to_dict(self) -> dict:
+        """
+        Return the configuration as a dictionary.
+        """
+        return {
+            "ontology_resolver": self.ontology_resolver,
+            "matching_strategy": self.matching_strategy,
+            "ontology_file_path": self.ontology_file_path,
+        }
+
+
+@lru_cache
+def get_ontology_env_config():
+    """
+    Retrieve the ontology configuration. This function utilizes caching to return a
+    singleton instance of the OntologyConfig class for efficiency.
+    """
+    return OntologyEnvConfig()