cognee 0.2.4__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/api/v1/users/routers/get_auth_router.py

@@ -1,7 +1,19 @@
+from fastapi import Depends
+
 from cognee.modules.users.get_fastapi_users import get_fastapi_users
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
 from cognee.modules.users.authentication.get_client_auth_backend import get_client_auth_backend
 
 
 def get_auth_router():
     auth_backend = get_client_auth_backend()
-    return get_fastapi_users().get_auth_router(auth_backend)
+    auth_router = get_fastapi_users().get_auth_router(auth_backend)
+
+    @auth_router.get("/me")
+    async def get_me(user: User = Depends(get_authenticated_user)):
+        return {
+            "email": user.email,
+        }
+
+    return auth_router

cognee/base_config.py

@@ -1,15 +1,24 @@
 import os
 from typing import Optional
 from functools import lru_cache
-from cognee.root_dir import get_absolute_path
+from cognee.root_dir import get_absolute_path, ensure_absolute_path
 from cognee.modules.observability.observers import Observer
 from pydantic_settings import BaseSettings, SettingsConfigDict
+import pydantic
 
 
 class BaseConfig(BaseSettings):
     data_root_directory: str = get_absolute_path(".data_storage")
     system_root_directory: str = get_absolute_path(".cognee_system")
     monitoring_tool: object = Observer.LANGFUSE
+
+    @pydantic.model_validator(mode="after")
+    def validate_paths(self):
+        # Require absolute paths for root directories
+        self.data_root_directory = ensure_absolute_path(self.data_root_directory)
+        self.system_root_directory = ensure_absolute_path(self.system_root_directory)
+        return self
+
     langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY")
     langfuse_secret_key: Optional[str] = os.getenv("LANGFUSE_SECRET_KEY")
     langfuse_host: Optional[str] = os.getenv("LANGFUSE_HOST")

cognee/infrastructure/databases/graph/config.py

@@ -6,6 +6,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 import pydantic
 from pydantic import Field
 from cognee.base_config import get_base_config
+from cognee.root_dir import ensure_absolute_path
 from cognee.shared.data_models import KnowledgeGraph
 
 
@@ -51,15 +52,20 @@ class GraphConfig(BaseSettings):
     @pydantic.model_validator(mode="after")
     def fill_derived(cls, values):
         provider = values.graph_database_provider.lower()
+        base_config = get_base_config()
 
         # Set default filename if no filename is provided
         if not values.graph_filename:
             values.graph_filename = f"cognee_graph_{provider}"
 
-        # Set file path based on graph database provider if no file path is provided
-        if not values.graph_file_path:
-            base_config = get_base_config()
-
+        # Handle graph file path
+        if values.graph_file_path:
+            # Check if absolute path is provided
+            values.graph_file_path = ensure_absolute_path(
+                os.path.join(values.graph_file_path, values.graph_filename)
+            )
+        else:
+            # Default path
             databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
             values.graph_file_path = os.path.join(databases_directory_path, values.graph_filename)
 

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

@@ -21,6 +21,8 @@ from cognee.infrastructure.databases.graph.graph_db_interface import (
 )
 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
 
 logger = get_logger()
 
@@ -106,6 +108,18 @@ class KuzuAdapter(GraphDBInterface):
 
             self.db.init_database()
             self.connection = Connection(self.db)
+
+            try:
+                self.connection.execute("INSTALL JSON;")
+            except Exception as e:
+                logger.info(f"JSON extension already installed or not needed: {e}")
+
+            try:
+                self.connection.execute("LOAD EXTENSION JSON;")
+                logger.info("Loaded JSON extension")
+            except Exception as e:
+                logger.info(f"JSON extension already loaded or unavailable: {e}")
+
             # Create node table with essential fields and timestamp
             self.connection.execute("""
                 CREATE NODE TABLE IF NOT EXISTS Node(
@@ -1693,3 +1707,124 @@ class KuzuAdapter(GraphDBInterface):
                 SET r.properties = $props
                 """
             await self.query(update_query, {"node_id": node_id, "props": new_props})
+
+    async def collect_events(self, ids: List[str]) -> Any:
+        """
+        Collect all Event-type nodes reachable within 1..2 hops
+        from the given node IDs.
+
+        Args:
+            graph_engine: Object exposing an async .query(str) -> Any
+            ids: List of node IDs (strings)
+
+        Returns:
+            List of events
+        """
+
+        event_collection_cypher = """UNWIND [{quoted}] AS uid
+            MATCH (start {{id: uid}})
+            MATCH (start)-[*1..2]-(event)
+            WHERE event.type = 'Event'
+            WITH DISTINCT event
+            RETURN collect(event) AS events;
+        """
+
+        query = event_collection_cypher.format(quoted=ids)
+        result = await self.query(query)
+        events = []
+        for node in result[0][0]:
+            props = json.loads(node["properties"])
+
+            event = {
+                "id": node["id"],
+                "name": node["name"],
+                "description": props.get("description"),
+            }
+
+            if props.get("location"):
+                event["location"] = props["location"]
+
+            events.append(event)
+
+        return [{"events": events}]
+
+    async def collect_time_ids(
+        self,
+        time_from: Optional[Timestamp] = None,
+        time_to: Optional[Timestamp] = None,
+    ) -> str:
+        """
+        Collect IDs of Timestamp nodes between time_from and time_to.
+
+        Args:
+            graph_engine: Object exposing an async .query(query, params) -> list[dict]
+            time_from: Lower bound int (inclusive), optional
+            time_to: Upper bound int (inclusive), optional
+
+        Returns:
+            A string of quoted IDs:  "'id1', 'id2', 'id3'"
+            (ready for use in a Cypher UNWIND clause).
+        """
+
+        ids: List[str] = []
+
+        if time_from and time_to:
+            time_from = date_to_int(time_from)
+            time_to = date_to_int(time_to)
+
+            cypher = f"""
+            MATCH (n:Node)
+            WHERE n.type = 'Timestamp'
+            // Extract time_at from the JSON string and cast to INT64
+            WITH n, json_extract(n.properties, '$.time_at') AS t_str
+            WITH n,
+                 CASE
+                   WHEN t_str IS NULL OR t_str = '' THEN NULL
+                   ELSE CAST(t_str AS INT64)
+                 END AS t
+            WHERE t >= {time_from}
+            AND t <= {time_to}
+            RETURN n.id as id
+            """
+
+        elif time_from:
+            time_from = date_to_int(time_from)
+
+            cypher = f"""
+            MATCH (n:Node)
+            WHERE n.type = 'Timestamp'
+            // Extract time_at from the JSON string and cast to INT64
+            WITH n, json_extract(n.properties, '$.time_at') AS t_str
+            WITH n,
+                 CASE
+                   WHEN t_str IS NULL OR t_str = '' THEN NULL
+                   ELSE CAST(t_str AS INT64)
+                 END AS t
+            WHERE t >= {time_from}
+            RETURN n.id as id
+            """
+
+        elif time_to:
+            time_to = date_to_int(time_to)
+
+            cypher = f"""
+            MATCH (n:Node)
+            WHERE n.type = 'Timestamp'
+            // Extract time_at from the JSON string and cast to INT64
+            WITH n, json_extract(n.properties, '$.time_at') AS t_str
+            WITH n,
+                 CASE
+                   WHEN t_str IS NULL OR t_str = '' THEN NULL
+                   ELSE CAST(t_str AS INT64)
+                 END AS t
+            WHERE t <= {time_to}
+            RETURN n.id as id
+            """
+
+        else:
+            return ids
+
+        time_nodes = await self.query(cypher)
+        time_ids_list = [item[0] for item in time_nodes]
+
+        return ", ".join(f"'{uid}'" for uid in time_ids_list)

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

@@ -11,6 +11,8 @@ from contextlib import asynccontextmanager
 from typing import Optional, Any, List, Dict, Type, Tuple
 
 from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.utils.generate_timestamp_datapoint import date_to_int
+from cognee.tasks.temporal_graph.models import Timestamp
 from cognee.shared.logging_utils import get_logger, ERROR
 from cognee.infrastructure.databases.graph.graph_db_interface import (
     GraphDBInterface,
@@ -1371,3 +1373,90 @@ class Neo4jAdapter(GraphDBInterface):
             query,
             params={"weight": float(weight), "node_ids": list(node_ids)},
         )
+
+    async def collect_events(self, ids: List[str]) -> Any:
+        """
+        Collect all Event-type nodes reachable within 1..2 hops
+        from the given node IDs.
+
+        Args:
+            graph_engine: Object exposing an async .query(str) -> Any
+            ids: List of node IDs (strings)
+
+        Returns:
+            List of events
+        """
+
+        event_collection_cypher = """UNWIND [{quoted}] AS uid
+            MATCH (start {{id: uid}})
+            MATCH (start)-[*1..2]-(event)
+            WHERE event.type = 'Event'
+            WITH DISTINCT event
+            RETURN collect(event) AS events;
+        """
+
+        query = event_collection_cypher.format(quoted=ids)
+        return await self.query(query)
+
+    async def collect_time_ids(
+        self,
+        time_from: Optional[Timestamp] = None,
+        time_to: Optional[Timestamp] = None,
+    ) -> str:
+        """
+        Collect IDs of Timestamp nodes between time_from and time_to.
+
+        Args:
+            graph_engine: Object exposing an async .query(query, params) -> list[dict]
+            time_from: Lower bound int (inclusive), optional
+            time_to: Upper bound int (inclusive), optional
+
+        Returns:
+            A string of quoted IDs:  "'id1', 'id2', 'id3'"
+            (ready for use in a Cypher UNWIND clause).
+        """
+
+        ids: List[str] = []
+
+        if time_from and time_to:
+            time_from = date_to_int(time_from)
+            time_to = date_to_int(time_to)
+
+            cypher = """
+            MATCH (n)
+            WHERE n.type = 'Timestamp'
+              AND n.time_at >= $time_from
+              AND n.time_at <= $time_to
+            RETURN n.id AS id
+            """
+            params = {"time_from": time_from, "time_to": time_to}
+
+        elif time_from:
+            time_from = date_to_int(time_from)
+
+            cypher = """
+            MATCH (n)
+            WHERE n.type = 'Timestamp'
+              AND n.time_at >= $time_from
+            RETURN n.id AS id
+            """
+            params = {"time_from": time_from}
+
+        elif time_to:
+            time_to = date_to_int(time_to)
+
+            cypher = """
+            MATCH (n)
+            WHERE n.type = 'Timestamp'
+              AND n.time_at <= $time_to
+            RETURN n.id AS id
+            """
+            params = {"time_to": time_to}
+
+        else:
+            return ids
+
+        time_nodes = await self.query(cypher, params)
+        time_ids_list = [item["id"] for item in time_nodes if "id" in item]
+
+        return ", ".join(f"'{uid}'" for uid in time_ids_list)

cognee/infrastructure/databases/relational/__init__.py

@@ -1,6 +1,8 @@
 from .ModelBase import Base
 from .config import get_relational_config
 from .config import get_migration_config
+from .get_async_session import get_async_session
+from .with_async_session import with_async_session
 from .create_db_and_tables import create_db_and_tables
 from .get_relational_engine import get_relational_engine
 from .get_migration_relational_engine import get_migration_relational_engine

cognee/infrastructure/databases/relational/get_async_session.py

@@ -0,0 +1,15 @@
+from typing import AsyncGenerator
+from contextlib import asynccontextmanager
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from .get_relational_engine import get_relational_engine
+
+
+@asynccontextmanager
+async def get_async_session(auto_commit=False) -> AsyncGenerator[AsyncSession, None]:
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        yield session
+
+        if auto_commit:
+            await session.commit()

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

@@ -57,7 +57,12 @@ class SQLAlchemyAdapter:
             )
         else:
             self.engine = create_async_engine(
-                connection_string, pool_size=12, max_overflow=12, poolclass=None
+                connection_string,
+                pool_size=5,
+                max_overflow=10,
+                pool_recycle=280,
+                pool_pre_ping=True,
+                pool_timeout=280,
             )
 
         self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)

cognee/infrastructure/databases/relational/with_async_session.py

@@ -0,0 +1,25 @@
+from typing import Any, Callable, Optional
+from sqlalchemy.ext.asyncio import AsyncSession
+from .get_async_session import get_async_session
+
+
+def get_session_from_args(args):
+    last_arg = args[-1]
+    if isinstance(last_arg, AsyncSession):
+        return last_arg
+    return None
+
+
+def with_async_session(func: Callable[..., Any]) -> Callable[..., Any]:
+    async def wrapper(*args, **kwargs):
+        session = kwargs.get("session") or get_session_from_args(args)  # type: Optional[AsyncSession]
+
+        if session is None:
+            async with get_async_session() as session:
+                result = await func(*args, **kwargs, session=session)
+                await session.commit()
+                return result
+        else:
+            return await func(*args, **kwargs)
+
+    return wrapper

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

@@ -538,7 +538,7 @@ class ChromaDBAdapter(VectorDBInterface):
             Returns True upon successful deletion of all collections.
         """
         client = await self.get_connection()
-        collections = await self.list_collections()
+        collections = await client.list_collections()
         for collection_name in collections:
             await client.delete_collection(collection_name)
         return True

cognee/infrastructure/databases/vector/config.py

@@ -1,9 +1,11 @@
 import os
 import pydantic
+from pathlib import Path
 from functools import lru_cache
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from cognee.base_config import get_base_config
+from cognee.root_dir import ensure_absolute_path
 
 
 class VectorConfig(BaseSettings):
@@ -11,11 +13,9 @@ class VectorConfig(BaseSettings):
     Manage the configuration settings for the vector database.
 
     Public methods:
-
     - to_dict: Convert the configuration to a dictionary.
 
     Instance variables:
-
     - vector_db_url: The URL of the vector database.
     - vector_db_port: The port for the vector database.
     - vector_db_key: The key for accessing the vector database.
@@ -30,10 +30,17 @@ class VectorConfig(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
     @pydantic.model_validator(mode="after")
-    def fill_derived(cls, values):
-        # Set file path based on graph database provider if no file path is provided
-        if not values.vector_db_url:
-            base_config = get_base_config()
+    def validate_paths(cls, values):
+        base_config = get_base_config()
+
+        # If vector_db_url is provided and is not a path skip checking if path is absolute (as it can also be a url)
+        if values.vector_db_url and Path(values.vector_db_url).exists():
+            # Relative path to absolute
+            values.vector_db_url = ensure_absolute_path(
+                values.vector_db_url,
+            )
+        else:
+            # Default path
             databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
             values.vector_db_url = os.path.join(databases_directory_path, "cognee.lancedb")
 

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

@@ -4,7 +4,7 @@ from fastembed import TextEmbedding
 import litellm
 import os
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
-from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
+from cognee.infrastructure.databases.exceptions import EmbeddingException
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )

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

@@ -250,9 +250,7 @@ def embedding_rate_limit_sync(func):
             logger.warning(error_msg)
 
             # Create a custom embedding rate limit exception
-            from cognee.infrastructure.databases.exceptions.EmbeddingException import (
-                EmbeddingException,
-            )
+            from cognee.infrastructure.databases.exceptions import EmbeddingException
 
             raise EmbeddingException(error_msg)
 
@@ -307,9 +305,7 @@ def embedding_rate_limit_async(func):
             logger.warning(error_msg)
 
             # Create a custom embedding rate limit exception
-            from cognee.infrastructure.databases.exceptions.EmbeddingException import (
-                EmbeddingException,
-            )
+            from cognee.infrastructure.databases.exceptions import EmbeddingException
 
             raise EmbeddingException(error_msg)
 

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

@@ -33,6 +33,7 @@ def get_embedding_engine() -> EmbeddingEngine:
         config.embedding_api_version,
         config.huggingface_tokenizer,
         llm_config.llm_api_key,
+        llm_config.llm_provider,
     )
 
 
@@ -47,6 +48,7 @@ def create_embedding_engine(
     embedding_api_version,
     huggingface_tokenizer,
     llm_api_key,
+    llm_provider,
 ):
     """
     Create and return an embedding engine based on the specified provider.
@@ -99,7 +101,8 @@ def create_embedding_engine(
 
     return LiteLLMEmbeddingEngine(
         provider=embedding_provider,
-        api_key=embedding_api_key or llm_api_key,
+        api_key=embedding_api_key
+        or (embedding_api_key if llm_provider == "custom" else llm_api_key),
         endpoint=embedding_endpoint,
         api_version=embedding_api_version,
         model=embedding_model,

cognee/infrastructure/files/storage/LocalFileStorage.py

@@ -189,6 +189,15 @@ class LocalFileStorage(Storage):
 
         return os.path.isfile(os.path.join(parsed_storage_path, file_path))
 
+    def get_size(self, file_path: str) -> int:
+        parsed_storage_path = get_parsed_path(self.storage_path)
+
+        return (
+            os.path.getsize(os.path.join(parsed_storage_path, file_path))
+            if self.file_exists(file_path)
+            else 0
+        )
+
     def ensure_directory_exists(self, directory_path: str = ""):
         """
         Ensure that the specified directory exists, creating it if necessary.

cognee/infrastructure/files/storage/S3FileStorage.py

@@ -146,6 +146,11 @@ class S3FileStorage(Storage):
             self.s3.isfile, os.path.join(self.storage_path.replace("s3://", ""), file_path)
         )
 
+    async def get_size(self, file_path: str) -> int:
+        return await run_async(
+            self.s3.size, os.path.join(self.storage_path.replace("s3://", ""), file_path)
+        )
+
     async def ensure_directory_exists(self, directory_path: str = ""):
         """
         Ensure that the specified directory exists, creating it if necessary.

cognee/infrastructure/files/storage/StorageManager.py

@@ -46,6 +46,12 @@ class StorageManager:
         else:
             return self.storage.is_file(file_path)
 
+    async def get_size(self, file_path: str) -> int:
+        if inspect.iscoroutinefunction(self.storage.get_size):
+            return await self.storage.get_size(file_path)
+        else:
+            return self.storage.get_size(file_path)
+
     async def store(self, file_path: str, data: BinaryIO, overwrite: bool = False) -> str:
         """
         Store data at the specified file path.
@@ -84,7 +90,7 @@ class StorageManager:
         """
         # Check the actual storage type by class name to determine if open() is async or sync
 
-        if self.storage.__class__.__name__ == "S3FileStorage" and file_path.startswith("s3://"):
+        if self.storage.__class__.__name__ == "S3FileStorage":
             # S3FileStorage.open() is async
             async with self.storage.open(file_path, *args, **kwargs) as file:
                 yield file

cognee/infrastructure/files/storage/storage.py

@@ -40,6 +40,22 @@ class Storage(Protocol):
         """
         pass
 
+    def get_size(self, file_path: str) -> int:
+        """
+        Get the size of a specified file in bytes.
+
+        Parameters:
+        -----------
+
+            - file_path (str): The path of the file to get the size of.
+
+        Returns:
+        --------
+
+            - int: The size of the file in bytes.
+        """
+        pass
+
     def store(self, file_path: str, data: Union[BinaryIO, str], overwrite: bool):
         """
         Store data at the specified file path.

cognee/infrastructure/llm/LLMGateway.py

@@ -144,3 +144,21 @@ class LLMGateway:
             )
 
             return extract_summary(content=content, response_model=response_model)
+
+    @staticmethod
+    def extract_event_graph(content: str, response_model: Type[BaseModel]) -> Coroutine:
+        # TODO: Add BAML version of category and extraction and update function (consulted with Igor)
+        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
+            extract_event_graph,
+        )
+
+        return extract_event_graph(content=content, response_model=response_model)
+
+    @staticmethod
+    def extract_event_entities(content: str, response_model: Type[BaseModel]) -> Coroutine:
+        # TODO: Add BAML version of category and extraction and update function (consulted with Igor)
+        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
+            extract_event_entities,
+        )
+
+        return extract_event_entities(content=content, response_model=response_model)

cognee/infrastructure/llm/config.py

@@ -35,7 +35,7 @@ class LLMConfig(BaseSettings):
 
     structured_output_framework: str = "instructor"
     llm_provider: str = "openai"
-    llm_model: str = "gpt-5-mini"
+    llm_model: str = "openai/gpt-4o-mini"
     llm_endpoint: str = ""
     llm_api_key: Optional[str] = None
     llm_api_version: Optional[str] = None
@@ -44,7 +44,7 @@ class LLMConfig(BaseSettings):
     llm_max_completion_tokens: int = 16384
 
     baml_llm_provider: str = "openai"
-    baml_llm_model: str = "gpt-5-mini"
+    baml_llm_model: str = "gpt-4o-mini"
     baml_llm_endpoint: str = ""
     baml_llm_api_key: Optional[str] = None
     baml_llm_temperature: float = 0.0
@@ -52,6 +52,8 @@ class LLMConfig(BaseSettings):
 
     transcription_model: str = "whisper-1"
     graph_prompt_path: str = "generate_graph_prompt.txt"
+    temporal_graph_prompt_path: str = "generate_event_graph_prompt.txt"
+    event_entity_prompt_path: str = "generate_event_entity_prompt.txt"
     llm_rate_limit_enabled: bool = False
     llm_rate_limit_requests: int = 60
     llm_rate_limit_interval: int = 60  # in seconds (default is 60 requests per minute)

cognee/infrastructure/llm/prompts/extract_query_time.txt

@@ -0,0 +1,15 @@
+For the purposes of identifying timestamps in a query, you are tasked with extracting relevant timestamps from the query.
+## Timestamp requirements
+- If the query contains interval extrack both starts_at and ends_at  properties
+- If the query contains an instantaneous timestamp, starts_at and ends_at should be the same
+- If the query its open-ended (before 2009 or after 2009), the corresponding non defined end of the time should be none
+    -For example: "before 2009" -- starts_at: None, ends_at: 2009 or  "after 2009" -- starts_at: 2009, ends_at: None
+- Put always the data that comes first in time as starts_at and the timestamps that comes second in time as ends_at
+- If starts_at or ends_at cannot be extracted both of them has to be None
+## Output Format
+Your reply should be a JSON: list of dictionaries with the following structure:
+```python
+class QueryInterval(BaseModel):
+    starts_at: Optional[Timestamp] = None
+    ends_at: Optional[Timestamp] = None
+```