cognee 0.3.4.dev1__py3-none-any.whl → 0.3.4.dev3__py3-none-any.whl

cognee/api/v1/cloud/routers/get_checks_router.py

@@ -16,7 +16,7 @@ def get_checks_router():
         api_token = request.headers.get("X-Api-Key")
 
         if api_token is None:
-            return CloudApiKeyMissingError()
+            raise CloudApiKeyMissingError()
 
         return await check_api_key(api_token)
 

cognee/api/v1/prune/prune.py

@@ -7,8 +7,8 @@ class prune:
         await _prune_data()
 
     @staticmethod
-    async def prune_system(graph=True, vector=True, metadata=False):
-        await _prune_system(graph, vector, metadata)
+    async def prune_system(graph=True, vector=True, metadata=False, cache=True):
+        await _prune_system(graph, vector, metadata, cache)
 
 
 if __name__ == "__main__":

cognee/api/v1/sync/sync.py

@@ -23,6 +23,7 @@ from cognee.modules.sync.methods import (
     mark_sync_completed,
     mark_sync_failed,
 )
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("sync")
 
@@ -583,7 +584,9 @@ async def _check_hashes_diff(
     logger.info(f"Checking missing hashes on cloud for dataset {dataset.id}")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, json=payload.dict(), headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()
@@ -630,7 +633,9 @@ async def _download_missing_files(
 
     headers = {"X-Api-Key": auth_token}
 
-    async with aiohttp.ClientSession() as session:
+    ssl_context = create_secure_ssl_context()
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    async with aiohttp.ClientSession(connector=connector) as session:
         for file_hash in hashes_missing_on_local:
             try:
                 # Download file from cloud by hash
@@ -749,7 +754,9 @@ async def _upload_missing_files(
 
     headers = {"X-Api-Key": auth_token}
 
-    async with aiohttp.ClientSession() as session:
+    ssl_context = create_secure_ssl_context()
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    async with aiohttp.ClientSession(connector=connector) as session:
         for file_info in files_to_upload:
             try:
                 file_dir = os.path.dirname(file_info.raw_data_location)
@@ -809,7 +816,9 @@ async def _prune_cloud_dataset(
     logger.info("Pruning cloud dataset to match local state")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.put(url, json=payload.dict(), headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()
@@ -852,7 +861,9 @@ async def _trigger_remote_cognify(
     logger.info(f"Triggering cognify processing for dataset {dataset_id}")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, json=payload, headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()

cognee/base_config.py

@@ -10,13 +10,27 @@ import pydantic
 class BaseConfig(BaseSettings):
     data_root_directory: str = get_absolute_path(".data_storage")
     system_root_directory: str = get_absolute_path(".cognee_system")
+    cache_root_directory: str = get_absolute_path(".cognee_cache")
     monitoring_tool: object = Observer.LANGFUSE
 
     @pydantic.model_validator(mode="after")
     def validate_paths(self):
+        # Adding this here temporarily to ensure that the cache root directory is set correctly for S3 storage automatically
+        # I'll remove this after we update documentation for S3 storage
+        # Auto-configure cache root directory for S3 storage if not explicitly set
+        storage_backend = os.getenv("STORAGE_BACKEND", "").lower()
+        cache_root_env = os.getenv("CACHE_ROOT_DIRECTORY")
+
+        if storage_backend == "s3" and not cache_root_env:
+            # Auto-generate S3 cache path when using S3 storage
+            bucket_name = os.getenv("STORAGE_BUCKET_NAME")
+            if bucket_name:
+                self.cache_root_directory = f"s3://{bucket_name}/cognee/cache"
+
         # Require absolute paths for root directories
         self.data_root_directory = ensure_absolute_path(self.data_root_directory)
         self.system_root_directory = ensure_absolute_path(self.system_root_directory)
+        self.cache_root_directory = ensure_absolute_path(self.cache_root_directory)
         return self
 
     langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY")
@@ -31,6 +45,7 @@ class BaseConfig(BaseSettings):
             "data_root_directory": self.data_root_directory,
             "system_root_directory": self.system_root_directory,
             "monitoring_tool": self.monitoring_tool,
+            "cache_root_directory": self.cache_root_directory,
         }
 
 

cognee/infrastructure/databases/graph/kuzu/remote_kuzu_adapter.py

@@ -7,6 +7,7 @@ import aiohttp
 from uuid import UUID
 
 from cognee.infrastructure.databases.graph.kuzu.adapter import KuzuAdapter
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger()
 
@@ -42,7 +43,9 @@ class RemoteKuzuAdapter(KuzuAdapter):
     async def _get_session(self) -> aiohttp.ClientSession:
         """Get or create an aiohttp session."""
         if self._session is None or self._session.closed:
-            self._session = aiohttp.ClientSession()
+            ssl_context = create_secure_ssl_context()
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+            self._session = aiohttp.ClientSession(connector=connector)
         return self._session
 
     async def close(self):

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -14,6 +14,7 @@ from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter im
     embedding_rate_limit_async,
     embedding_sleep_and_retry_async,
 )
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("OllamaEmbeddingEngine")
 
@@ -101,7 +102,9 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
 
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(
                 self.endpoint, json=payload, headers=headers, timeout=60.0
             ) as response:

cognee/infrastructure/files/storage/LocalFileStorage.py

@@ -253,6 +253,56 @@ class LocalFileStorage(Storage):
         if os.path.exists(full_file_path):
             os.remove(full_file_path)
 
+    def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+        from pathlib import Path
+
+        parsed_storage_path = get_parsed_path(self.storage_path)
+
+        if directory_path:
+            full_directory_path = os.path.join(parsed_storage_path, directory_path)
+        else:
+            full_directory_path = parsed_storage_path
+
+        directory_pathlib = Path(full_directory_path)
+
+        if not directory_pathlib.exists() or not directory_pathlib.is_dir():
+            return []
+
+        files = []
+
+        if recursive:
+            # Use rglob for recursive search
+            for file_path in directory_pathlib.rglob("*"):
+                if file_path.is_file():
+                    # Get relative path from storage root
+                    relative_path = os.path.relpath(str(file_path), parsed_storage_path)
+                    # Normalize path separators for consistency
+                    relative_path = relative_path.replace(os.sep, "/")
+                    files.append(relative_path)
+        else:
+            # Use iterdir for just immediate directory
+            for file_path in directory_pathlib.iterdir():
+                if file_path.is_file():
+                    # Get relative path from storage root
+                    relative_path = os.path.relpath(str(file_path), parsed_storage_path)
+                    # Normalize path separators for consistency
+                    relative_path = relative_path.replace(os.sep, "/")
+                    files.append(relative_path)
+
+        return files
+
     def remove_all(self, tree_path: str = None):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/infrastructure/files/storage/S3FileStorage.py

@@ -155,21 +155,19 @@ class S3FileStorage(Storage):
         """
         Ensure that the specified directory exists, creating it if necessary.
 
-        If the directory already exists, no action is taken.
+        For S3 storage, this is a no-op since directories are created implicitly
+        when files are written to paths. S3 doesn't have actual directories,
+        just object keys with prefixes that appear as directories.
 
         Parameters:
         -----------
 
             - directory_path (str): The path of the directory to check or create.
         """
-        if not directory_path.strip():
-            directory_path = self.storage_path.replace("s3://", "")
-
-        def ensure_directory():
-            if not self.s3.exists(directory_path):
-                self.s3.makedirs(directory_path, exist_ok=True)
-
-        await run_async(ensure_directory)
+        # In S3, directories don't exist as separate entities - they're just prefixes
+        # When you write a file to s3://bucket/path/to/file.txt, the "directories"
+        # path/ and path/to/ are implicitly created. No explicit action needed.
+        pass
 
     async def copy_file(self, source_file_path: str, destination_file_path: str):
         """
@@ -213,6 +211,55 @@ class S3FileStorage(Storage):
 
         await run_async(remove_file)
 
+    async def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+
+        def list_files_sync():
+            if directory_path:
+                # Combine storage path with directory path
+                full_path = os.path.join(self.storage_path.replace("s3://", ""), directory_path)
+            else:
+                full_path = self.storage_path.replace("s3://", "")
+
+            if recursive:
+                # Use ** for recursive search
+                pattern = f"{full_path}/**"
+            else:
+                # Just files in the immediate directory
+                pattern = f"{full_path}/*"
+
+            # Use s3fs glob to find files
+            try:
+                all_paths = self.s3.glob(pattern)
+                # Filter to only files (not directories)
+                files = [path for path in all_paths if self.s3.isfile(path)]
+
+                # Convert back to relative paths from storage root
+                storage_prefix = self.storage_path.replace("s3://", "")
+                relative_files = []
+                for file_path in files:
+                    if file_path.startswith(storage_prefix):
+                        relative_path = file_path[len(storage_prefix) :].lstrip("/")
+                        relative_files.append(relative_path)
+
+                return relative_files
+            except Exception:
+                # If directory doesn't exist or other error, return empty list
+                return []
+
+        return await run_async(list_files_sync)
+
     async def remove_all(self, tree_path: str):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/infrastructure/files/storage/StorageManager.py

@@ -135,6 +135,24 @@ class StorageManager:
         else:
             return self.storage.remove(file_path)
 
+    async def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+        if inspect.iscoroutinefunction(self.storage.list_files):
+            return await self.storage.list_files(directory_path, recursive)
+        else:
+            return self.storage.list_files(directory_path, recursive)
+
     async def remove_all(self, tree_path: str = None):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/modules/cloud/operations/check_api_key.py

@@ -1,6 +1,7 @@
 import aiohttp
 
 from cognee.modules.cloud.exceptions import CloudConnectionError
+from cognee.shared.utils import create_secure_ssl_context
 
 
 async def check_api_key(auth_token: str):
@@ -10,7 +11,9 @@ async def check_api_key(auth_token: str):
     headers = {"X-Api-Key": auth_token}
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, headers=headers) as response:
                 if response.status == 200:
                     return

cognee/modules/data/deletion/prune_system.py

@@ -1,9 +1,10 @@
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.shared.cache import delete_cache
 
 
-async def prune_system(graph=True, vector=True, metadata=True):
+async def prune_system(graph=True, vector=True, metadata=True, cache=True):
     if graph:
         graph_engine = await get_graph_engine()
         await graph_engine.delete_graph()
@@ -15,3 +16,6 @@ async def prune_system(graph=True, vector=True, metadata=True):
     if metadata:
         db_engine = get_relational_engine()
         await db_engine.delete_database()
+
+    if cache:
+        await delete_cache()

cognee/modules/notebooks/methods/create_notebook.py

@@ -7,6 +7,38 @@ from cognee.infrastructure.databases.relational import with_async_session
 from ..models.Notebook import Notebook, NotebookCell
 
 
+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="Python Development with Cognee Tutorial 🧠",
+            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(
     user_id: UUID,

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/retrieval/temporal_retriever.py

@@ -113,7 +113,7 @@ class TemporalRetriever(GraphCompletionRetriever):
             logger.info(
                 "No timestamps identified based on the query, performing retrieval using triplet search on events and entities."
             )
-            triplets = await self.get_context(query)
+            triplets = await self.get_triplets(query)
             return await self.resolve_edges_to_text(triplets)
 
         if ids:
@@ -122,7 +122,7 @@ class TemporalRetriever(GraphCompletionRetriever):
             logger.info(
                 "No events identified based on timestamp filtering, performing retrieval using triplet search on events and entities."
             )
-            triplets = await self.get_context(query)
+            triplets = await self.get_triplets(query)
             return await self.resolve_edges_to_text(triplets)
 
         vector_engine = get_vector_engine()