cognee 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

cognee/cli/commands/add_command.py

@@ -70,11 +70,11 @@ After adding data, use `cognee cognify` to process it into knowledge graphs.
                     await cognee.add(data=data_to_add, dataset_name=args.dataset_name)
                     fmt.success(f"Successfully added data to dataset '{args.dataset_name}'")
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to add data: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to add data: {str(e)}") from e
 
             asyncio.run(run_add())
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error adding data: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Failed to add data: {str(e)}", error_code=1) from e

cognee/cli/commands/cognify_command.py

@@ -107,7 +107,7 @@ After successful cognify processing, use `cognee search` to query the knowledge
                     )
                     return result
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to cognify: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to cognify: {str(e)}") from e
 
             result = asyncio.run(run_cognify())
 
@@ -124,5 +124,5 @@ After successful cognify processing, use `cognee search` to query the knowledge
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error during cognification: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error during cognification: {str(e)}", error_code=1) from e

cognee/cli/commands/config_command.py

@@ -79,8 +79,10 @@ Configuration changes will affect how cognee processes and stores data.
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error managing configuration: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(
+                f"Error managing configuration: {str(e)}", error_code=1
+            ) from e
 
     def _handle_get(self, args: argparse.Namespace) -> None:
         try:
@@ -122,7 +124,7 @@ Configuration changes will affect how cognee processes and stores data.
                     fmt.note("Configuration viewing not fully implemented yet")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to get configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to get configuration: {str(e)}") from e
 
     def _handle_set(self, args: argparse.Namespace) -> None:
         try:
@@ -141,7 +143,7 @@ Configuration changes will affect how cognee processes and stores data.
                 fmt.error(f"Failed to set configuration key '{args.key}'")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to set configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to set configuration: {str(e)}") from e
 
     def _handle_unset(self, args: argparse.Namespace) -> None:
         try:
@@ -189,7 +191,7 @@ Configuration changes will affect how cognee processes and stores data.
                 fmt.note("Use 'cognee config list' to see all available configuration options")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to unset configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to unset configuration: {str(e)}") from e
 
     def _handle_list(self, args: argparse.Namespace) -> None:
         try:
@@ -209,7 +211,7 @@ Configuration changes will affect how cognee processes and stores data.
             fmt.echo("  cognee config reset         - Reset all to defaults")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to list configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to list configuration: {str(e)}") from e
 
     def _handle_reset(self, args: argparse.Namespace) -> None:
         try:
@@ -222,4 +224,4 @@ Configuration changes will affect how cognee processes and stores data.
             fmt.echo("This would reset all settings to their default values")
 
         except Exception as e:
-            raise CliCommandInnerException(f"Failed to reset configuration: {str(e)}")
+            raise CliCommandInnerException(f"Failed to reset configuration: {str(e)}") from e

cognee/cli/commands/delete_command.py

@@ -100,7 +100,7 @@ Be careful with deletion operations as they are irreversible.
                     else:
                         await cognee.delete(dataset_name=args.dataset_name, user_id=args.user_id)
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to delete: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to delete: {str(e)}") from e
 
             asyncio.run(run_delete())
             # This success message may be inaccurate due to the underlying bug, but we leave it for now.
@@ -108,5 +108,5 @@ Be careful with deletion operations as they are irreversible.
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e

cognee/cli/commands/search_command.py

@@ -31,10 +31,6 @@ Search Types & Use Cases:
     Traditional RAG using document chunks without graph structure.
     Best for: Direct document retrieval, specific fact-finding.
 
-**INSIGHTS**:
-    Structured entity relationships and semantic connections.
-    Best for: Understanding concept relationships, knowledge mapping.
-
 **CHUNKS**:
     Raw text segments that match the query semantically.
     Best for: Finding specific passages, citations, exact content.
@@ -108,7 +104,7 @@ Search Types & Use Cases:
                     )
                     return results
                 except Exception as e:
-                    raise CliCommandInnerException(f"Failed to search: {str(e)}")
+                    raise CliCommandInnerException(f"Failed to search: {str(e)}") from e
 
             results = asyncio.run(run_search())
 
@@ -145,5 +141,5 @@ Search Types & Use Cases:
 
         except Exception as e:
             if isinstance(e, CliCommandInnerException):
-                raise CliCommandException(str(e), error_code=1)
-            raise CliCommandException(f"Error searching: {str(e)}", error_code=1)
+                raise CliCommandException(str(e), error_code=1) from e
+            raise CliCommandException(f"Error searching: {str(e)}", error_code=1) from e

cognee/cli/config.py

@@ -19,7 +19,6 @@ COMMAND_DESCRIPTIONS = {
 SEARCH_TYPE_CHOICES = [
     "GRAPH_COMPLETION",
     "RAG_COMPLETION",
-    "INSIGHTS",
     "CHUNKS",
     "SUMMARIES",
     "CODE",

cognee/context_global_variables.py

@@ -12,6 +12,11 @@ from cognee.modules.users.methods import get_user
 #       for different async tasks, threads and processes
 vector_db_config = ContextVar("vector_db_config", default=None)
 graph_db_config = ContextVar("graph_db_config", default=None)
+session_user = ContextVar("session_user", default=None)
+
+
+async def set_session_user_context_variable(user):
+    session_user.set(user)
 
 
 async def set_database_global_context_variables(dataset: Union[str, UUID], user_id: UUID):

cognee/exceptions/exceptions.py

@@ -56,7 +56,7 @@ class CogneeValidationError(CogneeApiError):
         self,
         message: str = "A validation error occurred.",
         name: str = "CogneeValidationError",
-        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
         log=True,
         log_level="ERROR",
     ):

cognee/infrastructure/databases/cache/__init__.py

@@ -0,0 +1,2 @@
+from .get_cache_engine import get_cache_engine
+from .config import get_cache_config

cognee/infrastructure/databases/cache/cache_db_interface.py

@@ -0,0 +1,79 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+
+
+class CacheDBInterface(ABC):
+    """
+    Abstract base class for distributed cache coordination systems (e.g., Redis, Memcached).
+    Provides a common interface for lock acquisition, release, and context-managed locking.
+    """
+
+    def __init__(self, host: str, port: int, lock_key: str):
+        self.host = host
+        self.port = port
+        self.lock_key = lock_key
+        self.lock = None
+
+    @abstractmethod
+    def acquire_lock(self):
+        """
+        Acquire a lock on the given key.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    @abstractmethod
+    def release_lock(self):
+        """
+        Release the lock if it is held.
+        Must be implemented by subclasses.
+        """
+        pass
+
+    @contextmanager
+    def hold_lock(self):
+        """
+        Context manager for safely acquiring and releasing the lock.
+        """
+        self.acquire()
+        try:
+            yield
+        finally:
+            self.release()
+
+    @abstractmethod
+    async def add_qa(
+        self,
+        user_id: str,
+        session_id: str,
+        question: str,
+        context: str,
+        answer: str,
+        ttl: int | None = 86400,
+    ):
+        """
+        Add a Q/A/context triplet to a cache session.
+        """
+
+        pass
+
+    @abstractmethod
+    async def get_latest_qa(self, user_id: str, session_id: str, last_n: int = 5):
+        """
+        Retrieve the most recent Q/A/context triplets for a session.
+        """
+        pass
+
+    @abstractmethod
+    async def get_all_qas(self, user_id: str, session_id: str):
+        """
+        Retrieve all Q/A/context triplets for the given session.
+        """
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """
+        Gracefully close any async connections.
+        """
+        pass

cognee/infrastructure/databases/cache/config.py

@@ -0,0 +1,44 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from functools import lru_cache
+from typing import Optional
+
+
+class CacheConfig(BaseSettings):
+    """
+    Configuration for distributed cache systems (e.g., Redis), used for locking or coordination.
+
+    Attributes:
+    - shared_kuzu_lock: Shared kuzu lock logic on/off.
+    - cache_host: Hostname of the cache service.
+    - cache_port: Port number for the cache service.
+    - agentic_lock_expire: Automatic lock expiration time (in seconds).
+    - agentic_lock_timeout: Maximum time (in seconds) to wait for the lock release.
+    """
+
+    caching: bool = False
+    shared_kuzu_lock: bool = False
+    cache_host: str = "localhost"
+    cache_port: int = 6379
+    cache_username: Optional[str] = None
+    cache_password: Optional[str] = None
+    agentic_lock_expire: int = 240
+    agentic_lock_timeout: int = 300
+
+    model_config = SettingsConfigDict(env_file=".env", extra="allow")
+
+    def to_dict(self) -> dict:
+        return {
+            "caching": self.caching,
+            "shared_kuzu_lock": self.shared_kuzu_lock,
+            "cache_host": self.cache_host,
+            "cache_port": self.cache_port,
+            "cache_username": self.cache_username,
+            "cache_password": self.cache_password,
+            "agentic_lock_expire": self.agentic_lock_expire,
+            "agentic_lock_timeout": self.agentic_lock_timeout,
+        }
+
+
+@lru_cache
+def get_cache_config():
+    return CacheConfig()

cognee/infrastructure/databases/cache/get_cache_engine.py

@@ -0,0 +1,67 @@
+"""Factory to get the appropriate cache coordination engine (e.g., Redis)."""
+
+from functools import lru_cache
+from typing import Optional
+from cognee.infrastructure.databases.cache.config import get_cache_config
+from cognee.infrastructure.databases.cache.cache_db_interface import CacheDBInterface
+
+config = get_cache_config()
+
+
+@lru_cache
+def create_cache_engine(
+    cache_host: str,
+    cache_port: int,
+    cache_username: str,
+    cache_password: str,
+    lock_key: str,
+    agentic_lock_expire: int = 240,
+    agentic_lock_timeout: int = 300,
+):
+    """
+    Factory function to instantiate a cache coordination backend (currently Redis).
+
+    Parameters:
+    -----------
+    - cache_host: Hostname or IP of the cache server.
+    - cache_port: Port number to connect to.
+    - cache_username: Username to authenticate with.
+    - cache_password: Password to authenticate with.
+    - lock_key: Identifier used for the locking resource.
+    - agentic_lock_expire: Duration to hold the lock after acquisition.
+    - agentic_lock_timeout: Max time to wait for the lock before failing.
+
+    Returns:
+    --------
+    - CacheDBInterface: An instance of the appropriate cache adapter. :TODO: Now we support only Redis. later if we add more here we can split the logic
+    """
+    if config.caching:
+        from cognee.infrastructure.databases.cache.redis.RedisAdapter import RedisAdapter
+
+        return RedisAdapter(
+            host=cache_host,
+            port=cache_port,
+            username=cache_username,
+            password=cache_password,
+            lock_name=lock_key,
+            timeout=agentic_lock_expire,
+            blocking_timeout=agentic_lock_timeout,
+        )
+    else:
+        return None
+
+
+def get_cache_engine(lock_key: Optional[str] = None) -> CacheDBInterface:
+    """
+    Returns a cache adapter instance using current context configuration.
+    """
+
+    return create_cache_engine(
+        cache_host=config.cache_host,
+        cache_port=config.cache_port,
+        cache_username=config.cache_username,
+        cache_password=config.cache_password,
+        lock_key=lock_key,
+        agentic_lock_expire=config.agentic_lock_expire,
+        agentic_lock_timeout=config.agentic_lock_timeout,
+    )

cognee/infrastructure/databases/cache/redis/RedisAdapter.py

@@ -0,0 +1,243 @@
+import asyncio
+import redis
+import redis.asyncio as aioredis
+from contextlib import contextmanager
+from cognee.infrastructure.databases.cache.cache_db_interface import CacheDBInterface
+from cognee.infrastructure.databases.exceptions import CacheConnectionError
+from cognee.shared.logging_utils import get_logger
+from datetime import datetime
+import json
+
+logger = get_logger("RedisAdapter")
+
+
+class RedisAdapter(CacheDBInterface):
+    def __init__(
+        self,
+        host,
+        port,
+        lock_name="default_lock",
+        username=None,
+        password=None,
+        timeout=240,
+        blocking_timeout=300,
+        connection_timeout=30,
+    ):
+        super().__init__(host, port, lock_name)
+
+        self.host = host
+        self.port = port
+        self.connection_timeout = connection_timeout
+
+        try:
+            self.sync_redis = redis.Redis(
+                host=host,
+                port=port,
+                username=username,
+                password=password,
+                socket_connect_timeout=connection_timeout,
+                socket_timeout=connection_timeout,
+            )
+            self.async_redis = aioredis.Redis(
+                host=host,
+                port=port,
+                username=username,
+                password=password,
+                decode_responses=True,
+                socket_connect_timeout=connection_timeout,
+            )
+            self.timeout = timeout
+            self.blocking_timeout = blocking_timeout
+
+            # Validate connection on initialization
+            self._validate_connection()
+            logger.info(f"Successfully connected to Redis at {host}:{port}")
+
+        except (redis.ConnectionError, redis.TimeoutError) as e:
+            error_msg = f"Failed to connect to Redis at {host}:{port}: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+        except Exception as e:
+            error_msg = f"Unexpected error initializing Redis adapter: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+
+    def _validate_connection(self):
+        """Validate Redis connection is available."""
+        try:
+            self.sync_redis.ping()
+        except (redis.ConnectionError, redis.TimeoutError) as e:
+            raise CacheConnectionError(
+                f"Cannot connect to Redis at {self.host}:{self.port}: {str(e)}"
+            ) from e
+
+    def acquire_lock(self):
+        """
+        Acquire the Redis lock manually. Raises if acquisition fails. (Sync because of Kuzu)
+        """
+        self.lock = self.sync_redis.lock(
+            name=self.lock_key,
+            timeout=self.timeout,
+            blocking_timeout=self.blocking_timeout,
+        )
+
+        acquired = self.lock.acquire()
+        if not acquired:
+            raise RuntimeError(f"Could not acquire Redis lock: {self.lock_key}")
+
+        return self.lock
+
+    def release_lock(self):
+        """
+        Release the Redis lock manually, if held. (Sync because of Kuzu)
+        """
+        if self.lock:
+            try:
+                self.lock.release()
+                self.lock = None
+            except redis.exceptions.LockError:
+                pass
+
+    @contextmanager
+    def hold_lock(self):
+        """
+        Context manager for acquiring and releasing the Redis lock automatically. (Sync because of Kuzu)
+        """
+        self.acquire()
+        try:
+            yield
+        finally:
+            self.release()
+
+    async def add_qa(
+        self,
+        user_id: str,
+        session_id: str,
+        question: str,
+        context: str,
+        answer: str,
+        ttl: int | None = 86400,
+    ):
+        """
+        Add a Q/A/context triplet to a Redis list for this session.
+        Creates the session if it doesn't exist.
+
+        Args:
+            user_id (str): The user ID.
+            session_id: Unique identifier for the session.
+            question: User question text.
+            context: Context used to answer.
+            answer: Assistant answer text.
+            ttl: Optional time-to-live (seconds). If provided, the session expires after this time.
+
+        Raises:
+            CacheConnectionError: If Redis connection fails or times out.
+        """
+        try:
+            session_key = f"agent_sessions:{user_id}:{session_id}"
+
+            qa_entry = {
+                "time": datetime.utcnow().isoformat(),
+                "question": question,
+                "context": context,
+                "answer": answer,
+            }
+
+            await self.async_redis.rpush(session_key, json.dumps(qa_entry))
+
+            if ttl is not None:
+                await self.async_redis.expire(session_key, ttl)
+
+        except (redis.ConnectionError, redis.TimeoutError) as e:
+            error_msg = f"Redis connection error while adding Q&A: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+        except Exception as e:
+            error_msg = f"Unexpected error while adding Q&A to Redis: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+
+    async def get_latest_qa(self, user_id: str, session_id: str, last_n: int = 5):
+        """
+        Retrieve the most recent Q/A/context triplet(s) for the given session.
+        """
+        session_key = f"agent_sessions:{user_id}:{session_id}"
+        if last_n == 1:
+            data = await self.async_redis.lindex(session_key, -1)
+            return [json.loads(data)] if data else None
+        else:
+            data = await self.async_redis.lrange(session_key, -last_n, -1)
+            return [json.loads(d) for d in data] if data else []
+
+    async def get_all_qas(self, user_id: str, session_id: str):
+        """
+        Retrieve all Q/A/context triplets for the given session.
+        """
+        session_key = f"agent_sessions:{user_id}:{session_id}"
+        entries = await self.async_redis.lrange(session_key, 0, -1)
+        return [json.loads(e) for e in entries]
+
+    async def close(self):
+        """
+        Gracefully close the async Redis connection.
+        """
+        await self.async_redis.aclose()
+
+
+async def main():
+    HOST = "localhost"
+    PORT = 6379
+
+    adapter = RedisAdapter(host=HOST, port=PORT)
+    session_id = "demo_session"
+    user_id = "demo_user_id"
+
+    print("\nAdding sample Q/A pairs...")
+    await adapter.add_qa(
+        user_id,
+        session_id,
+        "What is Redis?",
+        "Basic DB context",
+        "Redis is an in-memory data store.",
+    )
+    await adapter.add_qa(
+        user_id,
+        session_id,
+        "Who created Redis?",
+        "Historical context",
+        "Salvatore Sanfilippo (antirez).",
+    )
+
+    print("\nLatest QA:")
+    latest = await adapter.get_latest_qa(user_id, session_id)
+    print(json.dumps(latest, indent=2))
+
+    print("\nLast 2 QAs:")
+    last_two = await adapter.get_latest_qa(user_id, session_id, last_n=2)
+    print(json.dumps(last_two, indent=2))
+
+    session_id = "session_expire_demo"
+
+    await adapter.add_qa(
+        user_id,
+        session_id,
+        "What is Redis?",
+        "Database context",
+        "Redis is an in-memory data store.",
+    )
+
+    await adapter.add_qa(
+        user_id,
+        session_id,
+        "Who created Redis?",
+        "History context",
+        "Salvatore Sanfilippo (antirez).",
+    )
+
+    print(await adapter.get_all_qas(user_id, session_id))
+
+    await adapter.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cognee/infrastructure/databases/exceptions/__init__.py

@@ -11,4 +11,5 @@ from .exceptions import (
     EmbeddingException,
     MissingQueryParameterError,
     MutuallyExclusiveQueryParametersError,
+    CacheConnectionError,
 )

cognee/infrastructure/databases/exceptions/exceptions.py

@@ -15,7 +15,7 @@ class DatabaseNotCreatedError(CogneeSystemError):
         self,
         message: str = "The database has not been created yet. Please call `await setup()` first.",
         name: str = "DatabaseNotCreatedError",
-        status_code: int = status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code: int = status.HTTP_422_UNPROCESSABLE_CONTENT,
     ):
         super().__init__(message, name, status_code)
 
@@ -99,7 +99,7 @@ class EmbeddingException(CogneeConfigurationError):
         self,
         message: str = "Embedding Exception.",
         name: str = "EmbeddingException",
-        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
     ):
         super().__init__(message, name, status_code)
 
@@ -132,3 +132,19 @@ class MutuallyExclusiveQueryParametersError(CogneeValidationError):
     ):
         message = "The search function accepts either text or embedding as input, but not both."
         super().__init__(message, name, status_code)
+
+
+class CacheConnectionError(CogneeConfigurationError):
+    """
+    Raised when connection to the cache database (e.g., Redis) fails.
+
+    This error indicates that the cache service is unavailable or misconfigured.
+    """
+
+    def __init__(
+        self,
+        message: str = "Failed to connect to cache database. Please check your cache configuration.",
+        name: str = "CacheConnectionError",
+        status_code: int = status.HTTP_503_SERVICE_UNAVAILABLE,
+    ):
+        super().__init__(message, name, status_code)

cognee/infrastructure/databases/graph/get_graph_engine.py

@@ -162,5 +162,5 @@ def create_graph_engine(
 
     raise EnvironmentError(
         f"Unsupported graph database provider: {graph_database_provider}. "
-        f"Supported providers are: {', '.join(list(supported_databases.keys()) + ['neo4j', 'kuzu', 'kuzu-remote', 'memgraph', 'neptune', 'neptune_analytics'])}"
+        f"Supported providers are: {', '.join(list(supported_databases.keys()) + ['neo4j', 'kuzu', 'kuzu-remote', 'neptune', 'neptune_analytics'])}"
     )

cognee/infrastructure/databases/graph/graph_db_interface.py

@@ -159,6 +159,11 @@ class GraphDBInterface(ABC):
     - get_connections
     """
 
+    @abstractmethod
+    async def is_empty(self) -> bool:
+        logger.warning("is_empty() is not implemented")
+        return True
+
     @abstractmethod
     async def query(self, query: str, params: dict) -> List[Any]:
         """