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

cognee/modules/users/permissions/methods/get_document_ids_for_user.py

@@ -8,6 +8,16 @@ from ...models import ACL, Permission
 
 
 async def get_document_ids_for_user(user_id: UUID, datasets: list[str] = None) -> list[str]:
+    """
+        Return a list of documents ids for which the user has read permission.
+        If datasets are specified, return only documents from those datasets.
+    Args:
+        user_id: Id of the user
+        datasets: List of datasets
+
+    Returns:
+        list[str]: List of documents for which the user has read permission
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_principal.py

@@ -6,6 +6,15 @@ from ...models.Principal import Principal
 
 
 async def get_principal(principal_id: UUID):
+    """
+        Return information about a user based on their id
+    Args:
+        principal_id: Id of the user
+
+    Returns:
+        principal: Information about the user (principal)
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_principal_datasets.py

@@ -9,6 +9,17 @@ from ...models.ACL import ACL
 
 
 async def get_principal_datasets(principal: Principal, permission_type: str) -> list[Dataset]:
+    """
+        Return a list of datasets for which the user (principal) has a certain permission.
+    Args:
+        principal: Information about the user
+        permission_type: Type of permission
+
+    Returns:
+        list[Dataset]: List of datasets for which the user (principal)
+        has the permission (permission_type).
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_role.py

@@ -9,6 +9,16 @@ from ...models.Role import Role
 
 
 async def get_role(tenant_id: UUID, role_name: str):
+    """
+        Return the role with the name role_name of the given tenant.
+    Args:
+        tenant_id: Id of the given tenant
+        role_name: Name of the role
+
+    Returns
+        The role for the given tenant.
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_specific_user_permission_datasets.py

@@ -15,9 +15,9 @@ async def get_specific_user_permission_datasets(
         Return a list of datasets user has given permission for. If a list of datasets is provided,
         verify for which datasets user has appropriate permission for and return list of datasets he has permission for.
     Args:
-        user_id:
-        permission_type:
-        dataset_ids:
+        user_id: Id of the user.
+        permission_type: Type of the permission.
+        dataset_ids: Ids of the provided datasets
 
     Returns:
         list[Dataset]: List of datasets user has permission for

cognee/modules/users/permissions/methods/get_tenant.py

@@ -8,6 +8,15 @@ from ...models.Tenant import Tenant
 
 
 async def get_tenant(tenant_id: UUID):
+    """
+        Return information about the tenant based on the given id.
+    Args:
+        tenant_id: Id of the given tenant
+
+    Returns
+        Information about the given tenant.
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/give_default_permission_to_role.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_role(role_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the role with the given id as a default permission.
+    Args:
+        role_id: Id of the role
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/give_default_permission_to_tenant.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_tenant(tenant_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the tenant with the given id as a default permission.
+    Args:
+        tenant_id: Id of the tenant
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         tenant = (

cognee/modules/users/permissions/methods/give_default_permission_to_user.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_user(user_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the user with the given id as a default permission.
+    Args:
+        user_id: Id of the tenant
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = (await session.execute(select(User).where(User.id == user_id))).scalars().first()

cognee/modules/users/permissions/methods/give_permission_on_dataset.py

@@ -24,6 +24,16 @@ async def give_permission_on_dataset(
     dataset_id: UUID,
     permission_name: str,
 ):
+    """
+        Give a specific permission on a dataset to a user.
+    Args:
+        principal: User who is being given the permission on the dataset
+        dataset_id: Id of the dataset
+        permission_name: Name of permission to give
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/roles/methods/add_user_to_role.py

@@ -21,6 +21,17 @@ from cognee.modules.users.models import (
 
 
 async def add_user_to_role(user_id: UUID, role_id: UUID, owner_id: UUID):
+    """
+        Add a user with the given id to the role with the given id.
+    Args:
+        user_id: Id of the user.
+        role_id: Id of the role.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = (await session.execute(select(User).where(User.id == user_id))).scalars().first()

cognee/modules/users/roles/methods/create_role.py

@@ -16,6 +16,16 @@ async def create_role(
     role_name: str,
     owner_id: UUID,
 ):
+    """
+        Create a new role with the given name, if the request owner with the given id
+        has the necessary permission.
+    Args:
+        role_name: Name of the new role.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(owner_id)

cognee/modules/users/tenants/methods/add_user_to_tenant.py

@@ -13,6 +13,18 @@ from cognee.modules.users.exceptions import (
 
 
 async def add_user_to_tenant(user_id: UUID, tenant_id: UUID, owner_id: UUID):
+    """
+        Add a user with the given id to the tenant with the given id.
+        This can only be successful if the request owner with the given id is the tenant owner.
+    Args:
+        user_id: Id of the user.
+        tenant_id: Id of the tenant.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(user_id)

cognee/modules/users/tenants/methods/create_tenant.py

@@ -8,6 +8,16 @@ from cognee.modules.users.methods import get_user
 
 
 async def create_tenant(tenant_name: str, user_id: UUID):
+    """
+        Create a new tenant with the given name, for the user with the given id.
+        This user is the owner of the tenant.
+    Args:
+        tenant_name: Name of the new tenant.
+        user_id: Id of the user.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         try:

cognee/root_dir.py

@@ -20,6 +20,11 @@ def ensure_absolute_path(path: str) -> str:
     """
     if path is None:
         raise ValueError("Path cannot be None")
+
+    # Check if it's an S3 URL - S3 URLs are absolute by definition
+    if path.startswith("s3://"):
+        return path
+
     path_obj = Path(path).expanduser()
     if path_obj.is_absolute():
         return str(path_obj.resolve())

cognee/shared/cache.py

@@ -0,0 +1,346 @@
+"""
+Storage-aware cache management utilities for Cognee.
+
+This module provides cache functionality that works with both local and cloud storage
+backends (like S3) through the StorageManager abstraction.
+"""
+
+import hashlib
+import zipfile
+import asyncio
+from typing import Optional, Tuple
+import aiohttp
+import logging
+from io import BytesIO
+
+from cognee.base_config import get_base_config
+from cognee.infrastructure.files.storage.get_file_storage import get_file_storage
+from cognee.infrastructure.files.storage.StorageManager import StorageManager
+from cognee.shared.utils import create_secure_ssl_context
+
+logger = logging.getLogger(__name__)
+
+
+class StorageAwareCache:
+    """
+    A cache manager that works with different storage backends (local, S3, etc.)
+    """
+
+    def __init__(self, cache_subdir: str = "cache"):
+        """
+        Initialize the cache manager.
+
+        Args:
+            cache_subdir: Subdirectory name within the system root for caching
+        """
+        self.base_config = get_base_config()
+        # Since we're using cache_root_directory, don't add extra cache prefix
+        self.cache_base_path = ""
+        self.storage_manager: StorageManager = get_file_storage(
+            self.base_config.cache_root_directory
+        )
+
+        # Print absolute path
+        storage_path = self.storage_manager.storage.storage_path
+        if storage_path.startswith("s3://"):
+            absolute_path = storage_path  # S3 paths are already absolute
+        else:
+            import os
+
+            absolute_path = os.path.abspath(storage_path)
+        logger.info(f"Storage manager absolute path: {absolute_path}")
+
+    async def get_cache_dir(self) -> str:
+        """Get the base cache directory path."""
+        cache_path = self.cache_base_path or "."  # Use "." for root when cache_base_path is empty
+        await self.storage_manager.ensure_directory_exists(cache_path)
+        return cache_path
+
+    async def get_cache_subdir(self, name: str) -> str:
+        """Get a specific cache subdirectory."""
+        if self.cache_base_path:
+            cache_path = f"{self.cache_base_path}/{name}"
+        else:
+            cache_path = name
+        await self.storage_manager.ensure_directory_exists(cache_path)
+
+        # Return the absolute path based on storage system
+        if self.storage_manager.storage.storage_path.startswith("s3://"):
+            return cache_path
+        elif hasattr(self.storage_manager.storage, "storage_path"):
+            return f"{self.storage_manager.storage.storage_path}/{cache_path}"
+        else:
+            # Fallback for other storage types
+            return cache_path
+
+    async def delete_cache(self):
+        """Delete the entire cache directory."""
+        logger.info("Deleting cache...")
+        try:
+            await self.storage_manager.remove_all(self.cache_base_path)
+            logger.info("✓ Cache deleted successfully!")
+        except Exception as e:
+            logger.error(f"Error deleting cache: {e}")
+            raise
+
+    async def _is_cache_valid(self, cache_dir: str, version_or_hash: str) -> bool:
+        """Check if cached content is valid for the given version/hash."""
+        version_file = f"{cache_dir}/version.txt"
+
+        if not await self.storage_manager.file_exists(version_file):
+            return False
+
+        try:
+            async with self.storage_manager.open(version_file, "r") as f:
+                cached_version = (await asyncio.to_thread(f.read)).strip()
+                return cached_version == version_or_hash
+        except Exception as e:
+            logger.debug(f"Error checking cache validity: {e}")
+            return False
+
+    async def _clear_cache(self, cache_dir: str) -> None:
+        """Clear a cache directory."""
+        try:
+            await self.storage_manager.remove_all(cache_dir)
+        except Exception as e:
+            logger.debug(f"Error clearing cache directory {cache_dir}: {e}")
+
+    async def _check_remote_content_freshness(
+        self, url: str, cache_dir: str
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Check if remote content is fresher than cached version using HTTP headers.
+
+        Returns:
+            Tuple of (is_fresh: bool, new_identifier: Optional[str])
+        """
+        try:
+            # Make a HEAD request to check headers without downloading
+            ssl_context = create_secure_ssl_context()
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.head(url, timeout=aiohttp.ClientTimeout(total=30)) as response:
+                    response.raise_for_status()
+
+                    # Try ETag first (most reliable)
+                    etag = response.headers.get("ETag", "").strip('"')
+                    last_modified = response.headers.get("Last-Modified", "")
+
+            # Use ETag if available, otherwise Last-Modified
+            remote_identifier = etag if etag else last_modified
+
+            if not remote_identifier:
+                logger.debug("No freshness headers available, cannot check for updates")
+                return True, None  # Assume fresh if no headers
+
+            # Check cached identifier
+            identifier_file = f"{cache_dir}/content_id.txt"
+            if await self.storage_manager.file_exists(identifier_file):
+                async with self.storage_manager.open(identifier_file, "r") as f:
+                    cached_identifier = (await asyncio.to_thread(f.read)).strip()
+                    if cached_identifier == remote_identifier:
+                        logger.debug(f"Content is fresh (identifier: {remote_identifier[:20]}...)")
+                        return True, None
+                    else:
+                        logger.info(
+                            f"Content has changed (old: {cached_identifier[:20]}..., new: {remote_identifier[:20]}...)"
+                        )
+                        return False, remote_identifier
+            else:
+                # No cached identifier, treat as stale
+                return False, remote_identifier
+
+        except Exception as e:
+            logger.debug(f"Could not check remote freshness: {e}")
+            return True, None  # Assume fresh if we can't check
+
+    async def download_and_extract_zip(
+        self, url: str, cache_subdir_name: str, version_or_hash: str, force: bool = False
+    ) -> str:
+        """
+        Download a zip file and extract it to cache directory with content freshness checking.
+
+        Args:
+            url: URL to download zip file from
+            cache_subdir_name: Name of the cache subdirectory
+            version_or_hash: Version string or content hash for cache validation
+            force: If True, re-download even if already cached
+
+        Returns:
+            Path to the cached directory
+        """
+        cache_dir = await self.get_cache_subdir(cache_subdir_name)
+
+        # Check if already cached and valid
+        if not force and await self._is_cache_valid(cache_dir, version_or_hash):
+            # Also check if remote content has changed
+            is_fresh, new_identifier = await self._check_remote_content_freshness(url, cache_dir)
+            if is_fresh:
+                logger.debug(f"Content already cached and fresh for version {version_or_hash}")
+                return cache_dir
+            else:
+                logger.info("Cached content is stale, updating...")
+
+        # Clear old cache if it exists
+        await self._clear_cache(cache_dir)
+
+        logger.info(f"Downloading content from {url}...")
+
+        # Download the zip file
+        zip_content = BytesIO()
+        etag = ""
+        last_modified = ""
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.get(url, timeout=aiohttp.ClientTimeout(total=60)) as response:
+                response.raise_for_status()
+
+                # Extract headers before consuming response
+                etag = response.headers.get("ETag", "").strip('"')
+                last_modified = response.headers.get("Last-Modified", "")
+
+                # Read the response content
+                async for chunk in response.content.iter_chunked(8192):
+                    zip_content.write(chunk)
+        zip_content.seek(0)
+
+        # Extract the archive
+        await self.storage_manager.ensure_directory_exists(cache_dir)
+
+        # Extract files and store them using StorageManager
+        with zipfile.ZipFile(zip_content, "r") as zip_file:
+            for file_info in zip_file.infolist():
+                if file_info.is_dir():
+                    # Create directory
+                    dir_path = f"{cache_dir}/{file_info.filename}"
+                    await self.storage_manager.ensure_directory_exists(dir_path)
+                else:
+                    # Extract and store file
+                    file_data = zip_file.read(file_info.filename)
+                    file_path = f"{cache_dir}/{file_info.filename}"
+                    await self.storage_manager.store(file_path, BytesIO(file_data), overwrite=True)
+
+        # Write version info for future cache validation
+        version_file = f"{cache_dir}/version.txt"
+        await self.storage_manager.store(version_file, version_or_hash, overwrite=True)
+
+        # Store content identifier from response headers for freshness checking
+        content_identifier = etag if etag else last_modified
+
+        if content_identifier:
+            identifier_file = f"{cache_dir}/content_id.txt"
+            await self.storage_manager.store(identifier_file, content_identifier, overwrite=True)
+            logger.debug(f"Stored content identifier: {content_identifier[:20]}...")
+
+        logger.info("✓ Content downloaded and cached successfully!")
+        return cache_dir
+
+    async def file_exists(self, file_path: str) -> bool:
+        """Check if a file exists in cache storage."""
+        return await self.storage_manager.file_exists(file_path)
+
+    async def read_file(self, file_path: str, encoding: str = "utf-8"):
+        """Read a file from cache storage."""
+        return self.storage_manager.open(file_path, encoding=encoding)
+
+    async def list_files(self, directory_path: str):
+        """List files in a cache directory."""
+        try:
+            file_list = await self.storage_manager.list_files(directory_path)
+
+            # For S3 storage, convert relative paths to full S3 URLs
+            if self.storage_manager.storage.storage_path.startswith("s3://"):
+                full_paths = []
+                for file_path in file_list:
+                    full_s3_path = f"{self.storage_manager.storage.storage_path}/{file_path}"
+                    full_paths.append(full_s3_path)
+                return full_paths
+            else:
+                # For local storage, return absolute paths
+                storage_path = self.storage_manager.storage.storage_path
+                if not storage_path.startswith("/"):
+                    import os
+
+                    storage_path = os.path.abspath(storage_path)
+
+                full_paths = []
+                for file_path in file_list:
+                    if file_path.startswith("/"):
+                        full_paths.append(file_path)  # Already absolute
+                    else:
+                        full_paths.append(f"{storage_path}/{file_path}")
+                return full_paths
+
+        except Exception as e:
+            logger.debug(f"Error listing files in {directory_path}: {e}")
+            return []
+
+
+# Convenience functions that maintain API compatibility
+_cache_manager = None
+
+
+def get_cache_manager() -> StorageAwareCache:
+    """Get a singleton cache manager instance."""
+    global _cache_manager
+    if _cache_manager is None:
+        _cache_manager = StorageAwareCache()
+    return _cache_manager
+
+
+def generate_content_hash(url: str, additional_data: str = "") -> str:
+    """Generate a content hash from URL and optional additional data."""
+    content = f"{url}:{additional_data}"
+    return hashlib.md5(content.encode()).hexdigest()[:12]  # Short hash for readability
+
+
+# Async wrapper functions for backward compatibility
+async def delete_cache():
+    """Delete the Cognee cache directory."""
+    cache_manager = get_cache_manager()
+    await cache_manager.delete_cache()
+
+
+async def get_cognee_cache_dir() -> str:
+    """Get the base Cognee cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.get_cache_dir()
+
+
+async def get_cache_subdir(name: str) -> str:
+    """Get a specific cache subdirectory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.get_cache_subdir(name)
+
+
+async def download_and_extract_zip(
+    url: str, cache_dir_name: str, version_or_hash: str, force: bool = False
+) -> str:
+    """Download a zip file and extract it to cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.download_and_extract_zip(url, cache_dir_name, version_or_hash, force)
+
+
+async def get_tutorial_data_dir() -> str:
+    """Get the tutorial data cache directory."""
+    return await get_cache_subdir("tutorial_data")
+
+
+# Cache file operations
+async def cache_file_exists(file_path: str) -> bool:
+    """Check if a file exists in cache storage."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.file_exists(file_path)
+
+
+async def read_cache_file(file_path: str, encoding: str = "utf-8"):
+    """Read a file from cache storage."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.read_file(file_path, encoding)
+
+
+async def list_cache_files(directory_path: str):
+    """List files in a cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.list_files(directory_path)

cognee/shared/utils.py

@@ -1,6 +1,7 @@
 """This module contains utility functions for the cognee."""
 
 import os
+import ssl
 import requests
 from datetime import datetime, timezone
 import matplotlib.pyplot as plt
@@ -18,6 +19,17 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 proxy_url = "https://test.prometh.ai"
 
 
+def create_secure_ssl_context() -> ssl.SSLContext:
+    """
+    Create a secure SSL context.
+
+    By default, use the system's certificate store.
+    If users report SSL issues, I'm keeping this open in case we need to switch to:
+        ssl.create_default_context(cafile=certifi.where())
+    """
+    return ssl.create_default_context()
+
+
 def get_entities(tagged_tokens):
     import nltk