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

cognee/infrastructure/databases/exceptions/__init__.py

@@ -8,4 +8,7 @@ from .exceptions import (
     EntityNotFoundError,
     EntityAlreadyExistsError,
     DatabaseNotCreatedError,
+    EmbeddingException,
+    MissingQueryParameterError,
+    MutuallyExclusiveQueryParametersError,
 )

cognee/infrastructure/databases/exceptions/exceptions.py

@@ -1,13 +1,13 @@
 from fastapi import status
-from cognee.exceptions import CogneeApiError, CriticalError
+from cognee.exceptions import CogneeSystemError, CogneeValidationError, CogneeConfigurationError
 
 
-class DatabaseNotCreatedError(CriticalError):
+class DatabaseNotCreatedError(CogneeSystemError):
     """
     Represents an error indicating that the database has not been created. This error should
     be raised when an attempt is made to access the database before it has been initialized.
 
-    Inherits from CriticalError. Overrides the constructor to include a default message and
+    Inherits from CogneeSystemError. Overrides the constructor to include a default message and
     status code.
     """
 
@@ -20,10 +20,10 @@ class DatabaseNotCreatedError(CriticalError):
         super().__init__(message, name, status_code)
 
 
-class EntityNotFoundError(CogneeApiError):
+class EntityNotFoundError(CogneeValidationError):
     """
     Represents an error when a requested entity is not found in the database. This class
-    inherits from CogneeApiError.
+    inherits from CogneeValidationError.
 
     Public methods:
 
@@ -49,11 +49,11 @@ class EntityNotFoundError(CogneeApiError):
         # super().__init__(message, name, status_code) :TODO: This is not an error anymore with the dynamic exception handling therefore we shouldn't log error
 
 
-class EntityAlreadyExistsError(CogneeApiError):
+class EntityAlreadyExistsError(CogneeValidationError):
     """
     Represents an error when an entity creation is attempted but the entity already exists.
 
-    This class is derived from CogneeApiError and is used to signal a conflict in operations
+    This class is derived from CogneeValidationError and is used to signal a conflict in operations
     involving resource creation.
     """
 
@@ -66,11 +66,11 @@ class EntityAlreadyExistsError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class NodesetFilterNotSupportedError(CogneeApiError):
+class NodesetFilterNotSupportedError(CogneeConfigurationError):
     """
     Raise an exception when a nodeset filter is not supported by the current database.
 
-    This exception inherits from `CogneeApiError` and is designed to provide information
+    This exception inherits from `CogneeConfigurationError` and is designed to provide information
     about the specific issue of unsupported nodeset filters in the context of graph
     databases.
     """
@@ -84,3 +84,51 @@ class NodesetFilterNotSupportedError(CogneeApiError):
         self.message = message
         self.name = name
         self.status_code = status_code
+
+
+class EmbeddingException(CogneeConfigurationError):
+    """
+    Custom exception for handling embedding-related errors.
+
+    This exception class is designed to indicate issues specifically related to embeddings
+    within the application. It extends the base exception class CogneeConfigurationError allows
+    for customization of the error message, name, and status code.
+    """
+
+    def __init__(
+        self,
+        message: str = "Embedding Exception.",
+        name: str = "EmbeddingException",
+        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class MissingQueryParameterError(CogneeValidationError):
+    """
+    Raised when neither 'query_text' nor 'query_vector' is provided,
+    and at least one is required to perform the operation.
+    """
+
+    def __init__(
+        self,
+        name: str = "MissingQueryParameterError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        message = "One of query_text or query_vector must be provided!"
+        super().__init__(message, name, status_code)
+
+
+class MutuallyExclusiveQueryParametersError(CogneeValidationError):
+    """
+    Raised when both 'text' and 'embedding' are provided to the search function,
+    but only one type of input is allowed at a time.
+    """
+
+    def __init__(
+        self,
+        name: str = "MutuallyExclusiveQueryParametersError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        message = "The search function accepts either text or embedding as input, but not both."
+        super().__init__(message, name, status_code)

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

@@ -138,8 +138,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 +1632,64 @@ 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})

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

@@ -1322,3 +1322,52 @@ 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)},
+        )

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

@@ -3,11 +3,16 @@
 This module defines custom exceptions for Neptune Analytics operations.
 """
 
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import (
+    CogneeSystemError,
+    CogneeTransientError,
+    CogneeValidationError,
+    CogneeConfigurationError,
+)
 from fastapi import status
 
 
-class NeptuneAnalyticsError(CogneeApiError):
+class NeptuneAnalyticsError(CogneeSystemError):
     """Base exception for Neptune Analytics operations."""
 
     def __init__(
@@ -19,7 +24,7 @@ class NeptuneAnalyticsError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsConnectionError(NeptuneAnalyticsError):
+class NeptuneAnalyticsConnectionError(CogneeTransientError):
     """Exception raised when connection to Neptune Analytics fails."""
 
     def __init__(
@@ -31,7 +36,7 @@ class NeptuneAnalyticsConnectionError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsQueryError(NeptuneAnalyticsError):
+class NeptuneAnalyticsQueryError(CogneeValidationError):
     """Exception raised when a query execution fails."""
 
     def __init__(
@@ -43,7 +48,7 @@ class NeptuneAnalyticsQueryError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsAuthenticationError(NeptuneAnalyticsError):
+class NeptuneAnalyticsAuthenticationError(CogneeConfigurationError):
     """Exception raised when authentication with Neptune Analytics fails."""
 
     def __init__(
@@ -55,7 +60,7 @@ class NeptuneAnalyticsAuthenticationError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsConfigurationError(NeptuneAnalyticsError):
+class NeptuneAnalyticsConfigurationError(CogneeConfigurationError):
     """Exception raised when Neptune Analytics configuration is invalid."""
 
     def __init__(
@@ -67,7 +72,7 @@ class NeptuneAnalyticsConfigurationError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsTimeoutError(NeptuneAnalyticsError):
+class NeptuneAnalyticsTimeoutError(CogneeTransientError):
     """Exception raised when a Neptune Analytics operation times out."""
 
     def __init__(
@@ -79,7 +84,7 @@ class NeptuneAnalyticsTimeoutError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsThrottlingError(NeptuneAnalyticsError):
+class NeptuneAnalyticsThrottlingError(CogneeTransientError):
     """Exception raised when requests are throttled by Neptune Analytics."""
 
     def __init__(
@@ -91,7 +96,7 @@ class NeptuneAnalyticsThrottlingError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsResourceNotFoundError(NeptuneAnalyticsError):
+class NeptuneAnalyticsResourceNotFoundError(CogneeValidationError):
     """Exception raised when a Neptune Analytics resource is not found."""
 
     def __init__(
@@ -103,7 +108,7 @@ class NeptuneAnalyticsResourceNotFoundError(NeptuneAnalyticsError):
         super().__init__(message, name, status_code)
 
 
-class NeptuneAnalyticsInvalidParameterError(NeptuneAnalyticsError):
+class NeptuneAnalyticsInvalidParameterError(CogneeValidationError):
     """Exception raised when invalid parameters are provided to Neptune Analytics."""
 
     def __init__(

cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

@@ -9,7 +9,7 @@ from typing import List, Dict, Any, Optional, Tuple, Type, Union
 
 from falkordb import FalkorDB
 
-from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 from cognee.infrastructure.databases.graph.graph_db_interface import (
     GraphDBInterface,
     record_graph_changes,
@@ -721,7 +721,7 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
             Returns the search results as a result set from the graph database.
         """
         if query_text is None and query_vector is None:
-            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+            raise MissingQueryParameterError()
 
         if query_text and not query_vector:
             query_vector = (await self.embed_data([query_text]))[0]

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

@@ -5,7 +5,8 @@ import json
 from typing import List, Optional, Any, Dict, Type, Tuple
 from uuid import UUID
 
-from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
+from cognee.infrastructure.databases.exceptions import MutuallyExclusiveQueryParametersError
 from cognee.infrastructure.databases.graph.neptune_driver.adapter import NeptuneGraphDB
 from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
 from cognee.infrastructure.engine import DataPoint
@@ -274,11 +275,9 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
             limit = self._TOPK_UPPER_BOUND
 
         if query_vector and query_text:
-            raise InvalidValueError(
-                message="The search function accepts either text or embedding as input, but not both."
-            )
+            raise MutuallyExclusiveQueryParametersError()
         elif query_text is None and query_vector is None:
-            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+            raise MissingQueryParameterError()
         elif query_vector:
             embedding = query_vector
         else:

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

@@ -4,13 +4,13 @@ from uuid import UUID
 from typing import List, Optional
 from chromadb import AsyncHttpClient, Settings
 
-from cognee.exceptions import InvalidValueError
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.storage.utils import get_own_properties
 from cognee.infrastructure.engine import DataPoint
 from cognee.infrastructure.engine.utils import parse_id
 from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
 from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
+from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 
 from ..embeddings.EmbeddingEngine import EmbeddingEngine
 from ..vector_db_interface import VectorDBInterface
@@ -378,7 +378,7 @@ class ChromaDBAdapter(VectorDBInterface):
             Returns a list of ScoredResult instances representing the search results.
         """
         if query_text is None and query_vector is None:
-            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+            raise MissingQueryParameterError()
 
         if query_text and not query_vector:
             query_vector = (await self.embedding_engine.embed_text([query_text]))[0]

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

@@ -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

@@ -6,7 +6,7 @@ import math
 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.Gemini import (
     GeminiTokenizer,
 )
@@ -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/get_embedding_engine.py

@@ -27,7 +27,7 @@ def get_embedding_engine() -> EmbeddingEngine:
         config.embedding_provider,
         config.embedding_model,
         config.embedding_dimensions,
-        config.embedding_max_tokens,
+        config.embedding_max_completion_tokens,
         config.embedding_endpoint,
         config.embedding_api_key,
         config.embedding_api_version,
@@ -41,7 +41,7 @@ def create_embedding_engine(
     embedding_provider,
     embedding_model,
     embedding_dimensions,
-    embedding_max_tokens,
+    embedding_max_completion_tokens,
     embedding_endpoint,
     embedding_api_key,
     embedding_api_version,
@@ -58,7 +58,7 @@ def create_embedding_engine(
           'ollama', or another supported provider.
         - embedding_model: The model to be used for the embedding engine.
         - embedding_dimensions: The number of dimensions for the embeddings.
-        - embedding_max_tokens: The maximum number of tokens for the embeddings.
+        - embedding_max_completion_tokens: The maximum number of tokens for the embeddings.
         - embedding_endpoint: The endpoint for the embedding service, relevant for certain
           providers.
         - embedding_api_key: API key to authenticate with the embedding service, if
@@ -81,7 +81,7 @@ def create_embedding_engine(
         return FastembedEmbeddingEngine(
             model=embedding_model,
             dimensions=embedding_dimensions,
-            max_tokens=embedding_max_tokens,
+            max_completion_tokens=embedding_max_completion_tokens,
         )
 
     if embedding_provider == "ollama":
@@ -90,7 +90,7 @@ def create_embedding_engine(
         return OllamaEmbeddingEngine(
             model=embedding_model,
             dimensions=embedding_dimensions,
-            max_tokens=embedding_max_tokens,
+            max_completion_tokens=embedding_max_completion_tokens,
             endpoint=embedding_endpoint,
             huggingface_tokenizer=huggingface_tokenizer,
         )
@@ -104,5 +104,5 @@ def create_embedding_engine(
         api_version=embedding_api_version,
         model=embedding_model,
         dimensions=embedding_dimensions,
-        max_tokens=embedding_max_tokens,
+        max_completion_tokens=embedding_max_completion_tokens,
     )

cognee/infrastructure/databases/vector/exceptions/exceptions.py

@@ -1,12 +1,12 @@
 from fastapi import status
-from cognee.exceptions import CriticalError
+from cognee.exceptions import CogneeValidationError
 
 
-class CollectionNotFoundError(CriticalError):
+class CollectionNotFoundError(CogneeValidationError):
     """
     Represents an error that occurs when a requested collection cannot be found.
 
-    This class extends the CriticalError to handle specific cases where a requested
+    This class extends the CogneeValidationError to handle specific cases where a requested
     collection is unavailable. It can be initialized with a custom message and allows for
     logging options including log level and whether to log the error.
     """

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel
 from lancedb.pydantic import LanceModel, Vector
 from typing import Generic, List, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
 
-from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 from cognee.infrastructure.engine import DataPoint
 from cognee.infrastructure.engine.utils import parse_id
 from cognee.infrastructure.files.storage import get_file_storage
@@ -228,7 +228,7 @@ class LanceDBAdapter(VectorDBInterface):
         normalized: bool = True,
     ):
         if query_text is None and query_vector is None:
-            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+            raise MissingQueryParameterError()
 
         if query_text and not query_vector:
             query_vector = (await self.embedding_engine.embed_text([query_text]))[0]