cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/cli/reference.py

@@ -0,0 +1,26 @@
+from abc import abstractmethod
+from typing import Protocol, Optional
+import argparse
+
+
+class SupportsCliCommand(Protocol):
+    """Protocol for defining one cognee cli command"""
+
+    command_string: str
+    """name of the command"""
+    help_string: str
+    """the help string for argparse"""
+    description: Optional[str]
+    """the more detailed description for argparse, may include markdown for the docs"""
+    docs_url: Optional[str]
+    """the default docs url to be printed in case of an exception"""
+
+    @abstractmethod
+    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
+        """Configures the parser for the given argument"""
+        ...
+
+    @abstractmethod
+    def execute(self, args: argparse.Namespace) -> None:
+        """Executes the command with the given arguments"""
+        ...

cognee/cli/suppress_logging.py

@@ -0,0 +1,12 @@
+"""
+Module to suppress verbose logging before any cognee imports.
+This must be imported before any other cognee modules.
+"""
+
+import os
+
+# Set CLI mode to suppress verbose logging
+os.environ["COGNEE_CLI_MODE"] = "true"
+
+# Also set log level to ERROR for extra safety
+os.environ["LOG_LEVEL"] = "ERROR"

cognee/eval_framework/corpus_builder/corpus_builder_executor.py

@@ -5,7 +5,7 @@ from typing import Optional, Tuple, List, Dict, Union, Any, Callable, Awaitable
 from cognee.eval_framework.benchmark_adapters.benchmark_adapters import BenchmarkAdapter
 from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.modules.pipelines.tasks.task import Task
-from cognee.modules.pipelines import cognee_pipeline
+from cognee.modules.pipelines import run_pipeline
 
 logger = get_logger(level=ERROR)
 
@@ -61,7 +61,7 @@ class CorpusBuilderExecutor:
         await cognee.add(self.raw_corpus)
 
         tasks = await self.task_getter(chunk_size=chunk_size, chunker=chunker)
-        pipeline_run = cognee_pipeline(tasks=tasks)
+        pipeline_run = run_pipeline(tasks=tasks)
 
         async for run_info in pipeline_run:
             print(run_info)

cognee/eval_framework/eval_config.py

@@ -25,7 +25,7 @@ class EvalConfig(BaseSettings):
         "EM",
         "f1",
     ]  # Use only 'correctness' for DirectLLM
-    deepeval_model: str = "gpt-4o-mini"
+    deepeval_model: str = "gpt-5-mini"
 
     # Metrics params
     calculate_metrics: bool = True

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/get_graph_engine.py

@@ -21,10 +21,6 @@ async def get_graph_engine() -> GraphDBInterface:
     if hasattr(graph_client, "initialize"):
         await graph_client.initialize()
 
-    # Handle loading of graph for NetworkX
-    if config["graph_database_provider"].lower() == "networkx" and graph_client.graph is None:
-        await graph_client.load_graph_from_file()
-
     return graph_client
 
 
@@ -181,8 +177,7 @@ def create_graph_engine(
             graph_id=graph_identifier,
         )
 
-    from .networkx.adapter import NetworkXAdapter
-
-    graph_client = NetworkXAdapter(filename=graph_file_path)
-
-    return graph_client
+    raise EnvironmentError(
+        f"Unsupported graph database provider: {graph_database_provider}. "
+        f"Supported providers are: {', '.join(list(supported_databases.keys()) + ['neo4j', 'falkordb', 'kuzu', 'kuzu-remote', 'memgraph', 'neptune', 'neptune_analytics'])}"
+    )

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(
@@ -138,8 +152,9 @@ class KuzuAdapter(GraphDBInterface):
 
             s3_file_storage = S3FileStorage("")
 
-            async with self.KUZU_ASYNC_LOCK:
-                self.connection.execute("CHECKPOINT;")
+            if self.connection:
+                async with self.KUZU_ASYNC_LOCK:
+                    self.connection.execute("CHECKPOINT;")
 
             s3_file_storage.s3.put(self.temp_graph_file, self.db_path, recursive=True)
 
@@ -1631,3 +1646,185 @@ class KuzuAdapter(GraphDBInterface):
         """
         result = await self.query(query)
         return [record[0] for record in result] if result else []
+
+    async def get_last_user_interaction_ids(self, limit: int) -> List[str]:
+        """
+        Retrieve the IDs of the most recent CogneeUserInteraction nodes.
+        Parameters:
+        -----------
+        - limit (int): The maximum number of interaction IDs to return.
+        Returns:
+        --------
+        - List[str]: A list of interaction IDs, sorted by created_at descending.
+        """
+
+        query = """
+        MATCH (n)
+        WHERE n.type = 'CogneeUserInteraction'
+        RETURN n.id as id
+        ORDER BY n.created_at DESC
+        LIMIT $limit
+        """
+        rows = await self.query(query, {"limit": limit})
+
+        id_list = [row[0] for row in rows]
+        return id_list
+
+    async def apply_feedback_weight(
+        self,
+        node_ids: List[str],
+        weight: float,
+    ) -> None:
+        """
+        Increment `feedback_weight` inside r.properties JSON for edges where
+        relationship_name = 'used_graph_element_to_answer'.
+
+        """
+        # Step 1: fetch matching edges
+        query = """
+            MATCH (n:Node)-[r:EDGE]->()
+            WHERE n.id IN $node_ids AND r.relationship_name = 'used_graph_element_to_answer'
+            RETURN r.properties, n.id
+            """
+        results = await self.query(query, {"node_ids": node_ids})
+
+        # Step 2: update JSON client-side
+        updates = []
+        for props_json, source_id in results:
+            try:
+                props = json.loads(props_json) if props_json else {}
+            except json.JSONDecodeError:
+                props = {}
+
+            props["feedback_weight"] = props.get("feedback_weight", 0) + weight
+            updates.append((source_id, json.dumps(props)))
+
+        # Step 3: write back
+        for node_id, new_props in updates:
+            update_query = """
+                MATCH (n:Node)-[r:EDGE]->()
+                WHERE n.id = $node_id AND r.relationship_name = 'used_graph_element_to_answer'
+                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,
@@ -1322,3 +1324,139 @@ class Neo4jAdapter(GraphDBInterface):
         """
         result = await self.query(query)
         return [record["n"] for record in result] if result else []
+
+    async def get_last_user_interaction_ids(self, limit: int) -> List[str]:
+        """
+        Retrieve the IDs of the most recent CogneeUserInteraction nodes.
+        Parameters:
+        -----------
+        - limit (int): The maximum number of interaction IDs to return.
+        Returns:
+        --------
+        - List[str]: A list of interaction IDs, sorted by created_at descending.
+        """
+
+        query = """
+        MATCH (n)
+        WHERE n.type = 'CogneeUserInteraction'
+        RETURN n.id as id
+        ORDER BY n.created_at DESC
+        LIMIT $limit
+        """
+        rows = await self.query(query, {"limit": limit})
+
+        id_list = [row["id"] for row in rows if "id" in row]
+        return id_list
+
+    async def apply_feedback_weight(
+        self,
+        node_ids: List[str],
+        weight: float,
+    ) -> None:
+        """
+        Increment `feedback_weight` on relationships `:used_graph_element_to_answer`
+        outgoing from nodes whose `id` is in `node_ids`.
+
+        Args:
+            node_ids: List of node IDs to match.
+            weight: Amount to add to `r.feedback_weight` (can be negative).
+
+        Side effects:
+            Updates relationship property `feedback_weight`, defaulting missing values to 0.
+        """
+        query = """
+        MATCH (n)-[r]->()
+        WHERE n.id IN $node_ids AND r.relationship_name = 'used_graph_element_to_answer'
+        SET r.feedback_weight = coalesce(r.feedback_weight, 0) + $weight
+        """
+        await self.query(
+            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,
 )
@@ -41,11 +41,11 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
         self,
         model: Optional[str] = "openai/text-embedding-3-large",
         dimensions: Optional[int] = 3072,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
     ):
         self.model = model
         self.dimensions = dimensions
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.tokenizer = self.get_tokenizer()
         # self.retry_count = 0
         self.embedding_model = TextEmbedding(model_name=model)
@@ -112,7 +112,9 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
         """
         logger.debug("Loading tokenizer for FastembedEmbeddingEngine...")
 
-        tokenizer = TikTokenTokenizer(model="gpt-4o", max_tokens=self.max_tokens)
+        tokenizer = TikTokenTokenizer(
+            model="gpt-4o", max_completion_tokens=self.max_completion_tokens
+        )
 
         logger.debug("Tokenizer loaded for for FastembedEmbeddingEngine")
         return tokenizer