cognee 0.3.4.dev2__tar.gz → 0.3.4.dev4__tar.gz

{cognee-0.3.4.dev2 → cognee-0.3.4.dev4}/.env.template

@@ -47,6 +47,28 @@ BAML_LLM_API_VERSION=""
 # DATA_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_data/'
 # SYSTEM_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_system/'
 
+################################################################################
+# ☁️ Storage Backend Settings
+################################################################################
+# Configure storage backend (local filesystem or S3)
+# STORAGE_BACKEND="local"  # Default: uses local filesystem
+# 
+# -- To switch to S3 storage, uncomment and fill these: ---------------------
+# STORAGE_BACKEND="s3"
+# STORAGE_BUCKET_NAME="your-bucket-name"
+# AWS_REGION="us-east-1"
+# AWS_ACCESS_KEY_ID="your-access-key"
+# AWS_SECRET_ACCESS_KEY="your-secret-key"
+# 
+# -- S3 Root Directories (optional) -----------------------------------------
+# DATA_ROOT_DIRECTORY="s3://your-bucket/cognee/data"
+# SYSTEM_ROOT_DIRECTORY="s3://your-bucket/cognee/system"
+# 
+# -- Cache Directory (auto-configured for S3) -------------------------------
+# When STORAGE_BACKEND=s3, cache automatically uses S3: s3://BUCKET/cognee/cache
+# To override the automatic S3 cache location, uncomment:
+# CACHE_ROOT_DIRECTORY="s3://your-bucket/cognee/cache"
+
 ################################################################################
 # 🗄️ Relational database settings
 ################################################################################

{cognee-0.3.4.dev2 → cognee-0.3.4.dev4}/.gitignore

@@ -186,6 +186,7 @@ cognee/cache/
 # Default cognee system directory, used in development
 .cognee_system/
 .data_storage/
+.cognee_cache/
 .artifacts/
 .anon_id
 

{cognee-0.3.4.dev2 → cognee-0.3.4.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.3.4.dev2
+Version: 0.3.4.dev4
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -33,6 +33,7 @@ Requires-Dist: langfuse<3,>=2.32.0
 Requires-Dist: limits<5,>=4.4.1
 Requires-Dist: litellm<2.0.0,>=1.71.0
 Requires-Dist: matplotlib<4,>=3.8.3
+Requires-Dist: nbformat<6.0.0,>=5.7.0
 Requires-Dist: networkx<4,>=3.4.2
 Requires-Dist: nltk<4.0.0,>=3.9.1
 Requires-Dist: numpy<=4.0.0,>=1.26.4

{cognee-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev2 → cognee-0.3.4.dev4}/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-0.3.4.dev4/cognee/modules/notebooks/methods/create_notebook.py

@@ -0,0 +1,60 @@
+from uuid import UUID
+from typing import List, Optional
+from sqlalchemy.ext.asyncio import AsyncSession
+
+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(
+    user_id: UUID,
+    notebook_name: str,
+    cells: Optional[List[NotebookCell]],
+    deletable: Optional[bool],
+    session: AsyncSession,
+) -> Notebook:
+    notebook = Notebook(
+        name=notebook_name, owner_id=user_id, cells=cells, deletable=deletable or True
+    )
+
+    session.add(notebook)
+
+    await session.commit()
+
+    return notebook

cognee-0.3.4.dev4/cognee/modules/notebooks/methods/get_notebooks.py

@@ -0,0 +1,44 @@
+from uuid import UUID
+from typing import List
+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
+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())