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

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

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

@@ -57,7 +57,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         api_key: str = None,
         endpoint: str = None,
         api_version: str = None,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
     ):
         self.api_key = api_key
         self.endpoint = endpoint
@@ -65,7 +65,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         self.provider = provider
         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
 
@@ -179,20 +179,29 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         model = self.model.split("/")[-1]
 
         if "openai" in self.provider.lower():
-            tokenizer = TikTokenTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = TikTokenTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         elif "gemini" in self.provider.lower():
-            tokenizer = GeminiTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = GeminiTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         elif "mistral" in self.provider.lower():
-            tokenizer = MistralTokenizer(model=model, max_tokens=self.max_tokens)
+            tokenizer = MistralTokenizer(
+                model=model, max_completion_tokens=self.max_completion_tokens
+            )
         else:
             try:
                 tokenizer = HuggingFaceTokenizer(
-                    model=self.model.replace("hosted_vllm/", ""), max_tokens=self.max_tokens
+                    model=self.model.replace("hosted_vllm/", ""),
+                    max_completion_tokens=self.max_completion_tokens,
                 )
             except Exception as e:
                 logger.warning(f"Could not get tokenizer from HuggingFace due to: {e}")
                 logger.info("Switching to TikToken default tokenizer.")
-                tokenizer = TikTokenTokenizer(model=None, max_tokens=self.max_tokens)
+                tokenizer = TikTokenTokenizer(
+                    model=None, max_completion_tokens=self.max_completion_tokens
+                )
 
         logger.debug(f"Tokenizer loaded for model: {self.model}")
         return tokenizer

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

@@ -30,7 +30,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
     Instance variables:
     - model
     - dimensions
-    - max_tokens
+    - max_completion_tokens
     - endpoint
     - mock
     - huggingface_tokenizer_name
@@ -39,7 +39,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
 
     model: str
     dimensions: int
-    max_tokens: int
+    max_completion_tokens: int
     endpoint: str
     mock: bool
     huggingface_tokenizer_name: str
@@ -50,13 +50,13 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         self,
         model: Optional[str] = "avr/sfr-embedding-mistral:latest",
         dimensions: Optional[int] = 1024,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
         endpoint: Optional[str] = "http://localhost:11434/api/embeddings",
         huggingface_tokenizer: str = "Salesforce/SFR-Embedding-Mistral",
     ):
         self.model = model
         self.dimensions = dimensions
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.endpoint = endpoint
         self.huggingface_tokenizer_name = huggingface_tokenizer
         self.tokenizer = self.get_tokenizer()
@@ -132,7 +132,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         logger.debug("Loading HuggingfaceTokenizer for OllamaEmbeddingEngine...")
         tokenizer = HuggingFaceTokenizer(
-            model=self.huggingface_tokenizer_name, max_tokens=self.max_tokens
+            model=self.huggingface_tokenizer_name, max_completion_tokens=self.max_completion_tokens
         )
         logger.debug("Tokenizer loaded for OllamaEmbeddingEngine")
         return tokenizer

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

@@ -18,7 +18,7 @@ class EmbeddingConfig(BaseSettings):
     embedding_endpoint: Optional[str] = None
     embedding_api_key: Optional[str] = None
     embedding_api_version: Optional[str] = None
-    embedding_max_tokens: Optional[int] = 8191
+    embedding_max_completion_tokens: Optional[int] = 8191
     huggingface_tokenizer: Optional[str] = None
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -38,7 +38,7 @@ class EmbeddingConfig(BaseSettings):
             "embedding_endpoint": self.embedding_endpoint,
             "embedding_api_key": self.embedding_api_key,
             "embedding_api_version": self.embedding_api_version,
-            "embedding_max_tokens": self.embedding_max_tokens,
+            "embedding_max_completion_tokens": self.embedding_max_completion_tokens,
             "huggingface_tokenizer": self.huggingface_tokenizer,
         }
 

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)