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

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
 

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -41,6 +41,7 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
             abs_path.is_file()
         except (OSError, ValueError):
             # In case file path is too long it's most likely not a relative path
+            abs_path = data_item
             logger.debug(f"Data item was too long to be a possible file path: {abs_path}")
             abs_path = Path("")