cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

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]:
         """

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

@@ -4,7 +4,7 @@ import os
 import json
 import asyncio
 import tempfile
-from uuid import UUID
+from uuid import UUID, uuid5, NAMESPACE_OID
 from kuzu import Connection
 from kuzu.database import Database
 from datetime import datetime, timezone
@@ -23,9 +23,14 @@ from cognee.infrastructure.engine import DataPoint
 from cognee.modules.storage.utils import JSONEncoder
 from cognee.modules.engine.utils.generate_timestamp_datapoint import date_to_int
 from cognee.tasks.temporal_graph.models import Timestamp
+from cognee.infrastructure.databases.cache.config import get_cache_config
 
 logger = get_logger()
 
+cache_config = get_cache_config()
+if cache_config.shared_kuzu_lock:
+    from cognee.infrastructure.databases.cache.get_cache_engine import get_cache_engine
+
 
 class KuzuAdapter(GraphDBInterface):
     """
@@ -39,12 +44,20 @@ class KuzuAdapter(GraphDBInterface):
 
     def __init__(self, db_path: str):
         """Initialize Kuzu database connection and schema."""
+        self.open_connections = 0
+        self._is_closed = False
         self.db_path = db_path  # Path for the database directory
         self.db: Optional[Database] = None
         self.connection: Optional[Connection] = None
-        self.executor = ThreadPoolExecutor()
-        self._initialize_connection()
+        if cache_config.shared_kuzu_lock:
+            self.redis_lock = get_cache_engine(
+                lock_key="kuzu-lock-" + str(uuid5(NAMESPACE_OID, db_path))
+            )
+        else:
+            self.executor = ThreadPoolExecutor()
+            self._initialize_connection()
         self.KUZU_ASYNC_LOCK = asyncio.Lock()
+        self._connection_change_lock = asyncio.Lock()
 
     def _initialize_connection(self) -> None:
         """Initialize the Kuzu database connection and schema."""
@@ -185,6 +198,15 @@ class KuzuAdapter(GraphDBInterface):
         except FileNotFoundError:
             logger.warning(f"Kuzu S3 storage file not found: {self.db_path}")
 
+    async def is_empty(self) -> bool:
+        query = """
+        MATCH (n)
+        RETURN true
+        LIMIT 1;
+        """
+        query_result = await self.query(query)
+        return len(query_result) == 0
+
     async def query(self, query: str, params: Optional[dict] = None) -> List[Tuple]:
         """
         Execute a Kuzu query asynchronously with automatic reconnection.
@@ -209,9 +231,13 @@ class KuzuAdapter(GraphDBInterface):
         params = params or {}
 
         def blocking_query():
+            lock_acquired = False
             try:
+                if cache_config.shared_kuzu_lock:
+                    self.redis_lock.acquire_lock()
+                    lock_acquired = True
                 if not self.connection:
-                    logger.debug("Reconnecting to Kuzu database...")
+                    logger.info("Reconnecting to Kuzu database...")
                     self._initialize_connection()
 
                 result = self.connection.execute(query, params)
@@ -225,12 +251,47 @@ class KuzuAdapter(GraphDBInterface):
                             val = val.as_py()
                         processed_rows.append(val)
                     rows.append(tuple(processed_rows))
+
                 return rows
             except Exception as e:
                 logger.error(f"Query execution failed: {str(e)}")
                 raise
-
-        return await loop.run_in_executor(self.executor, blocking_query)
+            finally:
+                if cache_config.shared_kuzu_lock and lock_acquired:
+                    try:
+                        self.close()
+                    finally:
+                        self.redis_lock.release_lock()
+
+        if cache_config.shared_kuzu_lock:
+            async with self._connection_change_lock:
+                self.open_connections += 1
+                logger.info(f"Open connections after open: {self.open_connections}")
+                try:
+                    result = blocking_query()
+                finally:
+                    self.open_connections -= 1
+                    logger.info(f"Open connections after close: {self.open_connections}")
+                return result
+        else:
+            result = await loop.run_in_executor(self.executor, blocking_query)
+            return result
+
+    def close(self):
+        if self.connection:
+            del self.connection
+            self.connection = None
+        if self.db:
+            del self.db
+            self.db = None
+        self._is_closed = True
+        logger.info("Kuzu database closed successfully")
+
+    def reopen(self):
+        if self._is_closed:
+            self._is_closed = False
+            self._initialize_connection()
+            logger.info("Kuzu database re-opened successfully")
 
     @asynccontextmanager
     async def get_session(self):
@@ -1305,9 +1366,15 @@ class KuzuAdapter(GraphDBInterface):
                 params[param_name] = values
 
         where_clause = " AND ".join(where_clauses)
-        nodes_query = (
-            f"MATCH (n:Node) WHERE {where_clause} RETURN n.id, {{properties: n.properties}}"
-        )
+        nodes_query = f"""
+        MATCH (n:Node)
+        WHERE {where_clause}
+        RETURN n.id, {{
+            name: n.name,
+            type: n.type,
+            properties: n.properties
+        }}
+        """
         edges_query = f"""
         MATCH (n1:Node)-[r:EDGE]->(n2:Node)
         WHERE {where_clause.replace("n.", "n1.")} AND {where_clause.replace("n.", "n2.")}
@@ -1557,44 +1624,6 @@ class KuzuAdapter(GraphDBInterface):
             logger.error(f"Failed to delete graph data: {e}")
             raise
 
-    async def clear_database(self) -> None:
-        """
-        Clear all data from the database by deleting the database files and reinitializing.
-
-        This method removes all files associated with the database and reinitializes the Kuzu
-        database structure, ensuring a completely empty state. It handles exceptions that might
-        occur during file deletions or initializations carefully.
-        """
-        try:
-            if self.connection:
-                self.connection = None
-            if self.db:
-                self.db.close()
-                self.db = None
-
-            db_dir = os.path.dirname(self.db_path)
-            db_name = os.path.basename(self.db_path)
-            file_storage = get_file_storage(db_dir)
-
-            if await file_storage.file_exists(db_name):
-                await file_storage.remove_all()
-                logger.info(f"Deleted Kuzu database files at {self.db_path}")
-
-            # Reinitialize the database
-            self._initialize_connection()
-            # Verify the database is empty
-            result = self.connection.execute("MATCH (n:Node) RETURN COUNT(n)")
-            count = result.get_next()[0] if result.has_next() else 0
-            if count > 0:
-                logger.warning(
-                    f"Database still contains {count} nodes after clearing, forcing deletion"
-                )
-                self.connection.execute("MATCH (n:Node) DETACH DELETE n")
-            logger.info("Database cleared successfully")
-        except Exception as e:
-            logger.error(f"Error during database clearing: {e}")
-            raise
-
     async def get_document_subgraph(self, data_id: str):
         """
         Get all nodes that should be deleted when removing a document.

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -68,6 +68,7 @@ class Neo4jAdapter(GraphDBInterface):
             auth=auth,
             max_connection_lifetime=120,
             notifications_min_severity="OFF",
+            keep_alive=True,
         )
 
     async def initialize(self) -> None:
@@ -86,6 +87,15 @@ class Neo4jAdapter(GraphDBInterface):
         async with self.driver.session(database=self.graph_database_name) as session:
             yield session
 
+    async def is_empty(self) -> bool:
+        query = """
+        RETURN EXISTS {
+        MATCH (n)
+        } AS node_exists;
+        """
+        query_result = await self.query(query)
+        return not query_result[0]["node_exists"]
+
     @deadlock_retry()
     async def query(
         self,
@@ -205,7 +215,7 @@ class Neo4jAdapter(GraphDBInterface):
             {
                 "node_id": str(node.id),
                 "label": type(node).__name__,
-                "properties": self.serialize_properties(node.model_dump()),
+                "properties": self.serialize_properties(dict(node)),
             }
             for node in nodes
         ]
@@ -1066,7 +1076,7 @@ class Neo4jAdapter(GraphDBInterface):
         query_nodes = f"""
         MATCH (n)
         WHERE {where_clause}
-        RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties
+        RETURN n.id AS id, labels(n) AS labels, properties(n) AS properties
         """
         result_nodes = await self.query(query_nodes)
 
@@ -1081,7 +1091,7 @@ class Neo4jAdapter(GraphDBInterface):
         query_edges = f"""
         MATCH (n)-[r]->(m)
         WHERE {where_clause} AND {where_clause.replace("n.", "m.")}
-        RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
+        RETURN n.id AS source, n.id AS target, TYPE(r) AS type, properties(r) AS properties
         """
         result_edges = await self.query(query_edges)
 

cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py

@@ -8,7 +8,7 @@ from cognee.infrastructure.utils.calculate_backoff import calculate_backoff
 logger = get_logger("deadlock_retry")
 
 
-def deadlock_retry(max_retries=5):
+def deadlock_retry(max_retries=10):
     """
     Decorator that automatically retries an asynchronous function when rate limit errors occur.
 

cognee/infrastructure/databases/graph/neptune_driver/neptune_utils.py

@@ -53,7 +53,7 @@ def parse_neptune_url(url: str) -> Tuple[str, str]:
         return graph_id, region
 
     except Exception as e:
-        raise ValueError(f"Failed to parse Neptune Analytics URL '{url}': {str(e)}")
+        raise ValueError(f"Failed to parse Neptune Analytics URL '{url}': {str(e)}") from e
 
 
 def validate_graph_id(graph_id: str) -> bool:

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -283,7 +283,7 @@ class SQLAlchemyAdapter:
             try:
                 data_entity = (await session.scalars(select(Data).where(Data.id == data_id))).one()
             except (ValueError, NoResultFound) as e:
-                raise EntityNotFoundError(message=f"Entity not found: {str(e)}")
+                raise EntityNotFoundError(message=f"Entity not found: {str(e)}") from e
 
             # Check if other data objects point to the same raw data location
             raw_data_location_entities = (

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

@@ -1,8 +1,17 @@
-from cognee.shared.logging_utils import get_logger
+import os
+import logging
 from typing import List, Optional
 from fastembed import TextEmbedding
 import litellm
-import os
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
+
+from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions import EmbeddingException
 from cognee.infrastructure.llm.tokenizer.TikToken import (
@@ -57,6 +66,13 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def embed_text(self, text: List[str]) -> List[List[float]]:
         """
         Embed the given text into numerical vectors.
@@ -90,7 +106,9 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
 
         except Exception as error:
             logger.error(f"Embedding error in FastembedEmbeddingEngine: {str(error)}")
-            raise EmbeddingException(f"Failed to index data points using model {self.model}")
+            raise EmbeddingException(
+                f"Failed to index data points using model {self.model}"
+            ) from error
 
     def get_vector_size(self) -> int:
         """

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

@@ -1,15 +1,21 @@
 import asyncio
+import logging
+
 from cognee.shared.logging_utils import get_logger
 from typing import List, Optional
 import numpy as np
 import math
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
 import litellm
 import os
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions import EmbeddingException
-from cognee.infrastructure.llm.tokenizer.Gemini import (
-    GeminiTokenizer,
-)
 from cognee.infrastructure.llm.tokenizer.HuggingFace import (
     HuggingFaceTokenizer,
 )
@@ -19,10 +25,6 @@ from cognee.infrastructure.llm.tokenizer.Mistral import (
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )
-from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
-    embedding_rate_limit_async,
-    embedding_sleep_and_retry_async,
-)
 
 litellm.set_verbose = False
 logger = get_logger("LiteLLMEmbeddingEngine")
@@ -76,8 +78,13 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
-    @embedding_sleep_and_retry_async()
-    @embedding_rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def embed_text(self, text: List[str]) -> List[List[float]]:
         """
         Embed a list of text strings into vector representations.
@@ -150,7 +157,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
             litellm.exceptions.NotFoundError,
         ) as e:
             logger.error(f"Embedding error with model {self.model}: {str(e)}")
-            raise EmbeddingException(f"Failed to index data points using model {self.model}")
+            raise EmbeddingException(f"Failed to index data points using model {self.model}") from e
 
         except Exception as error:
             logger.error("Error embedding text: %s", str(error))