cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2.dev0__py3-none-any.whl

cognee/infrastructure/databases/cache/cache_db_interface.py

@@ -8,10 +8,13 @@ class CacheDBInterface(ABC):
     Provides a common interface for lock acquisition, release, and context-managed locking.
     """
 
-    def __init__(self, host: str, port: int, lock_key: str):
+    def __init__(
+        self, host: str, port: int, lock_key: str = "default_lock", log_key: str = "usage_logs"
+    ):
         self.host = host
         self.port = port
         self.lock_key = lock_key
+        self.log_key = log_key
         self.lock = None
 
     @abstractmethod
@@ -77,3 +80,37 @@ class CacheDBInterface(ABC):
         Gracefully close any async connections.
         """
         pass
+
+    @abstractmethod
+    async def log_usage(
+        self,
+        user_id: str,
+        log_entry: dict,
+        ttl: int | None = 604800,
+    ):
+        """
+        Log usage information (API endpoint calls, MCP tool invocations) to cache.
+
+        Args:
+            user_id: The user ID.
+            log_entry: Dictionary containing usage log information.
+            ttl: Optional time-to-live (seconds). If provided, the log list expires after this time.
+
+        Raises:
+            CacheConnectionError: If cache connection fails or times out.
+        """
+        pass
+
+    @abstractmethod
+    async def get_usage_logs(self, user_id: str, limit: int = 100):
+        """
+        Retrieve usage logs for a given user.
+
+        Args:
+            user_id: The user ID.
+            limit: Maximum number of logs to retrieve (default: 100).
+
+        Returns:
+            List of usage log entries, most recent first.
+        """
+        pass

cognee/infrastructure/databases/cache/config.py

@@ -13,6 +13,8 @@ class CacheConfig(BaseSettings):
     - 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.
+    - usage_logging: Enable/disable usage logging for API endpoints and MCP tools.
+    - usage_logging_ttl: Time-to-live for usage logs in seconds (default: 7 days).
     """
 
     cache_backend: Literal["redis", "fs"] = "fs"
@@ -24,6 +26,8 @@ class CacheConfig(BaseSettings):
     cache_password: Optional[str] = None
     agentic_lock_expire: int = 240
     agentic_lock_timeout: int = 300
+    usage_logging: bool = False
+    usage_logging_ttl: int = 604800
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -38,6 +42,8 @@ class CacheConfig(BaseSettings):
             "cache_password": self.cache_password,
             "agentic_lock_expire": self.agentic_lock_expire,
             "agentic_lock_timeout": self.agentic_lock_timeout,
+            "usage_logging": self.usage_logging,
+            "usage_logging_ttl": self.usage_logging_ttl,
         }
 
 

cognee/infrastructure/databases/cache/fscache/FsCacheAdapter.py

@@ -89,6 +89,27 @@ class FSCacheAdapter(CacheDBInterface):
             return None
         return json.loads(value)
 
+    async def log_usage(
+        self,
+        user_id: str,
+        log_entry: dict,
+        ttl: int | None = 604800,
+    ):
+        """
+        Usage logging is not supported in filesystem cache backend.
+        This method is a no-op to satisfy the interface.
+        """
+        logger.warning("Usage logging not supported in FSCacheAdapter, skipping")
+        pass
+
+    async def get_usage_logs(self, user_id: str, limit: int = 100):
+        """
+        Usage logging is not supported in filesystem cache backend.
+        This method returns an empty list to satisfy the interface.
+        """
+        logger.warning("Usage logging not supported in FSCacheAdapter, returning empty list")
+        return []
+
     async def close(self):
         if self.cache is not None:
             self.cache.expire()

cognee/infrastructure/databases/cache/get_cache_engine.py

@@ -1,7 +1,6 @@
 """Factory to get the appropriate cache coordination engine (e.g., Redis)."""
 
 from functools import lru_cache
-import os
 from typing import Optional
 from cognee.infrastructure.databases.cache.config import get_cache_config
 from cognee.infrastructure.databases.cache.cache_db_interface import CacheDBInterface
@@ -17,6 +16,7 @@ def create_cache_engine(
     cache_username: str,
     cache_password: str,
     lock_key: str,
+    log_key: str,
     agentic_lock_expire: int = 240,
     agentic_lock_timeout: int = 300,
 ):
@@ -30,6 +30,7 @@ def create_cache_engine(
     - cache_username: Username to authenticate with.
     - cache_password: Password to authenticate with.
     - lock_key: Identifier used for the locking resource.
+    - log_key: Identifier used for usage logging.
     - agentic_lock_expire: Duration to hold the lock after acquisition.
     - agentic_lock_timeout: Max time to wait for the lock before failing.
 
@@ -37,7 +38,7 @@ def create_cache_engine(
     --------
     - CacheDBInterface: An instance of the appropriate cache adapter.
     """
-    if config.caching:
+    if config.caching or config.usage_logging:
         from cognee.infrastructure.databases.cache.redis.RedisAdapter import RedisAdapter
 
         if config.cache_backend == "redis":
@@ -47,6 +48,7 @@ def create_cache_engine(
                 username=cache_username,
                 password=cache_password,
                 lock_name=lock_key,
+                log_key=log_key,
                 timeout=agentic_lock_expire,
                 blocking_timeout=agentic_lock_timeout,
             )
@@ -61,7 +63,10 @@ def create_cache_engine(
         return None
 
 
-def get_cache_engine(lock_key: Optional[str] = None) -> CacheDBInterface:
+def get_cache_engine(
+    lock_key: Optional[str] = "default_lock",
+    log_key: Optional[str] = "usage_logs",
+) -> Optional[CacheDBInterface]:
     """
     Returns a cache adapter instance using current context configuration.
     """
@@ -72,6 +77,7 @@ def get_cache_engine(lock_key: Optional[str] = None) -> CacheDBInterface:
         cache_username=config.cache_username,
         cache_password=config.cache_password,
         lock_key=lock_key,
+        log_key=log_key,
         agentic_lock_expire=config.agentic_lock_expire,
         agentic_lock_timeout=config.agentic_lock_timeout,
     )

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

@@ -17,13 +17,14 @@ class RedisAdapter(CacheDBInterface):
         host,
         port,
         lock_name="default_lock",
+        log_key="usage_logs",
         username=None,
         password=None,
         timeout=240,
         blocking_timeout=300,
         connection_timeout=30,
     ):
-        super().__init__(host, port, lock_name)
+        super().__init__(host, port, lock_name, log_key)
 
         self.host = host
         self.port = port
@@ -177,6 +178,64 @@ class RedisAdapter(CacheDBInterface):
         entries = await self.async_redis.lrange(session_key, 0, -1)
         return [json.loads(e) for e in entries]
 
+    async def log_usage(
+        self,
+        user_id: str,
+        log_entry: dict,
+        ttl: int | None = 604800,
+    ):
+        """
+        Log usage information (API endpoint calls, MCP tool invocations) to Redis.
+
+        Args:
+            user_id: The user ID.
+            log_entry: Dictionary containing usage log information.
+            ttl: Optional time-to-live (seconds). If provided, the log list expires after this time.
+
+        Raises:
+            CacheConnectionError: If Redis connection fails or times out.
+        """
+        try:
+            usage_logs_key = f"{self.log_key}:{user_id}"
+
+            await self.async_redis.rpush(usage_logs_key, json.dumps(log_entry))
+
+            if ttl is not None:
+                await self.async_redis.expire(usage_logs_key, ttl)
+
+        except (redis.ConnectionError, redis.TimeoutError) as e:
+            error_msg = f"Redis connection error while logging usage: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+        except Exception as e:
+            error_msg = f"Unexpected error while logging usage to Redis: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+
+    async def get_usage_logs(self, user_id: str, limit: int = 100):
+        """
+        Retrieve usage logs for a given user.
+
+        Args:
+            user_id: The user ID.
+            limit: Maximum number of logs to retrieve (default: 100).
+
+        Returns:
+            List of usage log entries, most recent first.
+        """
+        try:
+            usage_logs_key = f"{self.log_key}:{user_id}"
+            entries = await self.async_redis.lrange(usage_logs_key, -limit, -1)
+            return [json.loads(e) for e in reversed(entries)] if entries else []
+        except (redis.ConnectionError, redis.TimeoutError) as e:
+            error_msg = f"Redis connection error while retrieving usage logs: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+        except Exception as e:
+            error_msg = f"Unexpected error while retrieving usage logs from Redis: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+
     async def close(self):
         """
         Gracefully close the async Redis connection.

cognee/infrastructure/databases/dataset_database_handler/supported_dataset_database_handlers.py

@@ -7,6 +7,9 @@ from cognee.infrastructure.databases.vector.lancedb.LanceDBDatasetDatabaseHandle
 from cognee.infrastructure.databases.graph.kuzu.KuzuDatasetDatabaseHandler import (
     KuzuDatasetDatabaseHandler,
 )
+from cognee.infrastructure.databases.vector.pgvector.PGVectorDatasetDatabaseHandler import (
+    PGVectorDatasetDatabaseHandler,
+)
 
 supported_dataset_database_handlers = {
     "neo4j_aura_dev": {
@@ -14,5 +17,9 @@ supported_dataset_database_handlers = {
         "handler_provider": "neo4j",
     },
     "lancedb": {"handler_instance": LanceDBDatasetDatabaseHandler, "handler_provider": "lancedb"},
+    "pgvector": {
+        "handler_instance": PGVectorDatasetDatabaseHandler,
+        "handler_provider": "pgvector",
+    },
     "kuzu": {"handler_instance": KuzuDatasetDatabaseHandler, "handler_provider": "kuzu"},
 }

cognee/infrastructure/databases/graph/get_graph_engine.py

@@ -24,7 +24,6 @@ async def get_graph_engine() -> GraphDBInterface:
     return graph_client
 
 
-@lru_cache
 def create_graph_engine(
     graph_database_provider,
     graph_file_path,
@@ -35,6 +34,35 @@ def create_graph_engine(
     graph_database_port="",
     graph_database_key="",
     graph_dataset_database_handler="",
+):
+    """
+    Wrapper function to call create graph engine with caching.
+    For a detailed description, see _create_graph_engine.
+    """
+    return _create_graph_engine(
+        graph_database_provider,
+        graph_file_path,
+        graph_database_url,
+        graph_database_name,
+        graph_database_username,
+        graph_database_password,
+        graph_database_port,
+        graph_database_key,
+        graph_dataset_database_handler,
+    )
+
+
+@lru_cache
+def _create_graph_engine(
+    graph_database_provider,
+    graph_file_path,
+    graph_database_url="",
+    graph_database_name="",
+    graph_database_username="",
+    graph_database_password="",
+    graph_database_port="",
+    graph_database_key="",
+    graph_dataset_database_handler="",
 ):
     """
     Create a graph engine based on the specified provider type.

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

@@ -1,11 +1,13 @@
 import os
+import aiohttp
 import asyncio
-import requests
 import base64
 import hashlib
 from uuid import UUID
 from typing import Optional
+from urllib.parse import urlparse
 from cryptography.fernet import Fernet
+from aiohttp import BasicAuth
 
 from cognee.infrastructure.databases.graph import get_graph_config
 from cognee.modules.users.models import User, DatasetDatabase
@@ -23,7 +25,6 @@ class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
 
     Quality of life improvements:
     - Allow configuration of different Neo4j Aura plans and regions.
-    - Requests should be made async, currently a blocking requests library is used.
     """
 
     @classmethod
@@ -49,6 +50,7 @@ class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
         graph_db_name = f"{dataset_id}"
 
         # Client credentials and encryption
+        # Note: Should not be used as class variables so that they are not persisted in memory longer than needed
         client_id = os.environ.get("NEO4J_CLIENT_ID", None)
         client_secret = os.environ.get("NEO4J_CLIENT_SECRET", None)
         tenant_id = os.environ.get("NEO4J_TENANT_ID", None)
@@ -63,22 +65,13 @@ class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
                 "NEO4J_CLIENT_ID, NEO4J_CLIENT_SECRET, and NEO4J_TENANT_ID environment variables must be set to use Neo4j Aura DatasetDatabase Handling."
             )
 
-        # Make the request with HTTP Basic Auth
-        def get_aura_token(client_id: str, client_secret: str) -> dict:
-            url = "https://api.neo4j.io/oauth/token"
-            data = {"grant_type": "client_credentials"}  # sent as application/x-www-form-urlencoded
-
-            resp = requests.post(url, data=data, auth=(client_id, client_secret))
-            resp.raise_for_status()  # raises if the request failed
-            return resp.json()
-
-        resp = get_aura_token(client_id, client_secret)
+        resp_token = await cls._get_aura_token(client_id, client_secret)
 
         url = "https://api.neo4j.io/v1/instances"
 
         headers = {
             "accept": "application/json",
-            "Authorization": f"Bearer {resp['access_token']}",
+            "Authorization": f"Bearer {resp_token['access_token']}",
             "Content-Type": "application/json",
         }
 
@@ -96,31 +89,38 @@ class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
             "cloud_provider": "gcp",
         }
 
-        response = requests.post(url, headers=headers, json=payload)
+        async def _create_database_instance_request():
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url, headers=headers, json=payload) as resp:
+                    resp.raise_for_status()
+                    return await resp.json()
+
+        resp_create = await _create_database_instance_request()
 
         graph_db_name = "neo4j"  # Has to be 'neo4j' for Aura
-        graph_db_url = response.json()["data"]["connection_url"]
-        graph_db_key = resp["access_token"]
-        graph_db_username = response.json()["data"]["username"]
-        graph_db_password = response.json()["data"]["password"]
+        graph_db_url = resp_create["data"]["connection_url"]
+        graph_db_key = resp_token["access_token"]
+        graph_db_username = resp_create["data"]["username"]
+        graph_db_password = resp_create["data"]["password"]
 
         async def _wait_for_neo4j_instance_provisioning(instance_id: str, headers: dict):
             # Poll until the instance is running
             status_url = f"https://api.neo4j.io/v1/instances/{instance_id}"
             status = ""
             for attempt in range(30):  # Try for up to ~5 minutes
-                status_resp = requests.get(
-                    status_url, headers=headers
-                )  # TODO: Use async requests with httpx
-                status = status_resp.json()["data"]["status"]
-                if status.lower() == "running":
-                    return
-                await asyncio.sleep(10)
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(status_url, headers=headers) as resp:
+                        resp.raise_for_status()
+                        status_resp = await resp.json()
+                        status = status_resp["data"]["status"]
+                        if status.lower() == "running":
+                            return
+                        await asyncio.sleep(10)
             raise TimeoutError(
                 f"Neo4j instance '{graph_db_name}' did not become ready within 5 minutes. Status: {status}"
             )
 
-        instance_id = response.json()["data"]["id"]
+        instance_id = resp_create["data"]["id"]
         await _wait_for_neo4j_instance_provisioning(instance_id, headers)
 
         encrypted_db_password_bytes = cipher.encrypt(graph_db_password.encode())
@@ -165,4 +165,39 @@ class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
 
     @classmethod
     async def delete_dataset(cls, dataset_database: DatasetDatabase):
-        pass
+        # Get dataset database information and credentials
+        dataset_database = await cls.resolve_dataset_connection_info(dataset_database)
+
+        parsed_url = urlparse(dataset_database.graph_database_url)
+        instance_id = parsed_url.hostname.split(".")[0]
+
+        url = f"https://api.neo4j.io/v1/instances/{instance_id}"
+
+        # Get access token for Neo4j Aura API
+        # Client credentials
+        client_id = os.environ.get("NEO4J_CLIENT_ID", None)
+        client_secret = os.environ.get("NEO4J_CLIENT_SECRET", None)
+        resp = await cls._get_aura_token(client_id, client_secret)
+
+        headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {resp['access_token']}",
+            "Content-Type": "application/json",
+        }
+
+        async with aiohttp.ClientSession() as session:
+            async with session.delete(url, headers=headers) as resp:
+                resp.raise_for_status()
+                return await resp.json()
+
+    @classmethod
+    async def _get_aura_token(cls, client_id: str, client_secret: str) -> dict:
+        url = "https://api.neo4j.io/oauth/token"
+        data = {"grant_type": "client_credentials"}  # sent as application/x-www-form-urlencoded
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                url, data=data, auth=BasicAuth(client_id, client_secret)
+            ) as resp:
+                resp.raise_for_status()
+                return await resp.json()

cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py

@@ -236,6 +236,7 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
         query_vector: Optional[List[float]] = None,
         limit: Optional[int] = None,
         with_vector: bool = False,
+        include_payload: bool = False,  # TODO: Add support for this parameter
     ):
         """
         Perform a search in the specified collection using either a text query or a vector
@@ -290,7 +291,7 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
         query_string = f"""
         CALL neptune.algo.vectors.topKByEmbeddingWithFiltering({{
                 topK: {limit},
-                embedding: {embedding}, 
+                embedding: {embedding},
                 nodeFilter: {{ equals: {{property: '{self._COLLECTION_PREFIX}', value: '{collection_name}'}} }}
               }}
             )
@@ -299,7 +300,7 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
 
         if with_vector:
             query_string += """
-        WITH node, score, id(node) as node_id 
+        WITH node, score, id(node) as node_id
         MATCH (n)
         WHERE id(n) = id(node)
         CALL neptune.algo.vectors.get(n)
@@ -319,7 +320,12 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
             self._na_exception_handler(e, query_string)
 
     async def batch_search(
-        self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
+        self,
+        collection_name: str,
+        query_texts: List[str],
+        limit: int,
+        with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         """
         Perform a batch search using multiple text queries against a collection.
@@ -342,7 +348,14 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
         data_vectors = await self.embedding_engine.embed_text(query_texts)
         return await asyncio.gather(
             *[
-                self.search(collection_name, None, vector, limit, with_vectors)
+                self.search(
+                    collection_name,
+                    None,
+                    vector,
+                    limit,
+                    with_vectors,
+                    include_payload=include_payload,
+                )
                 for vector in data_vectors
             ]
         )

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

@@ -8,7 +8,7 @@ from typing import AsyncGenerator, List
 from contextlib import asynccontextmanager
 from sqlalchemy.orm import joinedload
 from sqlalchemy.exc import NoResultFound
-from sqlalchemy import NullPool, text, select, MetaData, Table, delete, inspect
+from sqlalchemy import NullPool, text, select, MetaData, Table, delete, inspect, URL
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
 
 from cognee.modules.data.models.Data import Data
@@ -549,6 +549,7 @@ class SQLAlchemyAdapter:
                             )
                             await connection.execute(drop_table_query)
                         metadata.clear()
+
         except Exception as e:
             logger.error(f"Error deleting database: {e}")
             raise e

cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py

@@ -355,6 +355,7 @@ class ChromaDBAdapter(VectorDBInterface):
         limit: Optional[int] = 15,
         with_vector: bool = False,
         normalized: bool = True,
+        include_payload: bool = False,  # TODO: Add support for this parameter when set to False
     ):
         """
         Search for items in a collection using either a text or a vector query.
@@ -441,6 +442,7 @@ class ChromaDBAdapter(VectorDBInterface):
         query_texts: List[str],
         limit: int = 5,
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         """
         Perform multiple searches in a single request for efficiency, returning results for each

cognee/infrastructure/databases/vector/config.py

@@ -29,6 +29,9 @@ class VectorConfig(BaseSettings):
     vector_db_key: str = ""
     vector_db_provider: str = "lancedb"
     vector_dataset_database_handler: str = "lancedb"
+    vector_db_username: str = ""
+    vector_db_password: str = ""
+    vector_db_host: str = ""
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -65,6 +68,9 @@ class VectorConfig(BaseSettings):
             "vector_db_key": self.vector_db_key,
             "vector_db_provider": self.vector_db_provider,
             "vector_dataset_database_handler": self.vector_dataset_database_handler,
+            "vector_db_username": self.vector_db_username,
+            "vector_db_password": self.vector_db_password,
+            "vector_db_host": self.vector_db_host,
         }