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

cognee/infrastructure/llm/tokenizer/Mistral/adapter.py

@@ -16,17 +16,17 @@ class MistralTokenizer(TokenizerInterface):
 
     Instance variables include:
     - model: str
-    - max_tokens: int
+    - max_completion_tokens: int
     - tokenizer: MistralTokenizer
     """
 
     def __init__(
         self,
         model: str,
-        max_tokens: int = 3072,
+        max_completion_tokens: int = 3072,
     ):
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
         # Import here to make it an optional dependency
         from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

cognee/infrastructure/llm/tokenizer/TikToken/adapter.py

@@ -13,10 +13,10 @@ class TikTokenTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: Optional[str] = None,
-        max_tokens: int = 8191,
+        max_completion_tokens: int = 8191,
     ):
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         # Initialize TikToken for GPT based on model
         if model:
             self.tokenizer = tiktoken.encoding_for_model(self.model)
@@ -93,9 +93,9 @@ class TikTokenTokenizer(TokenizerInterface):
         num_tokens = len(self.tokenizer.encode(text))
         return num_tokens
 
-    def trim_text_to_max_tokens(self, text: str) -> str:
+    def trim_text_to_max_completion_tokens(self, text: str) -> str:
         """
-        Trim the text so that the number of tokens does not exceed max_tokens.
+        Trim the text so that the number of tokens does not exceed max_completion_tokens.
 
         Parameters:
         -----------
@@ -111,13 +111,13 @@ class TikTokenTokenizer(TokenizerInterface):
         num_tokens = self.count_tokens(text)
 
         # If the number of tokens is within the limit, return the text as is
-        if num_tokens <= self.max_tokens:
+        if num_tokens <= self.max_completion_tokens:
             return text
 
         # If the number exceeds the limit, trim the text
         # This is a simple trim, it may cut words in half; consider using word boundaries for a cleaner cut
         encoded_text = self.tokenizer.encode(text)
-        trimmed_encoded_text = encoded_text[: self.max_tokens]
+        trimmed_encoded_text = encoded_text[: self.max_completion_tokens]
         # Decoding the trimmed text
         trimmed_text = self.tokenizer.decode(trimmed_encoded_text)
         return trimmed_text

cognee/infrastructure/llm/utils.py

@@ -32,13 +32,13 @@ def get_max_chunk_tokens():
 
     # We need to make sure chunk size won't take more than half of LLM max context token size
     # but it also can't be bigger than the embedding engine max token size
-    llm_cutoff_point = llm_client.max_tokens // 2  # Round down the division
-    max_chunk_tokens = min(embedding_engine.max_tokens, llm_cutoff_point)
+    llm_cutoff_point = llm_client.max_completion_tokens // 2  # Round down the division
+    max_chunk_tokens = min(embedding_engine.max_completion_tokens, llm_cutoff_point)
 
     return max_chunk_tokens
 
 
-def get_model_max_tokens(model_name: str):
+def get_model_max_completion_tokens(model_name: str):
     """
     Retrieve the maximum token limit for a specified model name if it exists.
 
@@ -56,15 +56,15 @@ def get_model_max_tokens(model_name: str):
 
         Number of max tokens of model, or None if model is unknown
     """
-    max_tokens = None
+    max_completion_tokens = None
 
     if model_name in litellm.model_cost:
-        max_tokens = litellm.model_cost[model_name]["max_tokens"]
-        logger.debug(f"Max input tokens for {model_name}: {max_tokens}")
+        max_completion_tokens = litellm.model_cost[model_name]["max_tokens"]
+        logger.debug(f"Max input tokens for {model_name}: {max_completion_tokens}")
     else:
         logger.info("Model not found in LiteLLM's model_cost.")
 
-    return max_tokens
+    return max_completion_tokens
 
 
 async def test_llm_connection():

cognee/modules/data/exceptions/exceptions.py

@@ -1,8 +1,11 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import (
+    CogneeValidationError,
+    CogneeConfigurationError,
+)
 from fastapi import status
 
 
-class UnstructuredLibraryImportError(CogneeApiError):
+class UnstructuredLibraryImportError(CogneeConfigurationError):
     def __init__(
         self,
         message: str = "Import error. Unstructured library is not installed.",
@@ -12,7 +15,7 @@ class UnstructuredLibraryImportError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class UnauthorizedDataAccessError(CogneeApiError):
+class UnauthorizedDataAccessError(CogneeValidationError):
     def __init__(
         self,
         message: str = "User does not have permission to access this data.",
@@ -22,7 +25,7 @@ class UnauthorizedDataAccessError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class DatasetNotFoundError(CogneeApiError):
+class DatasetNotFoundError(CogneeValidationError):
     def __init__(
         self,
         message: str = "Dataset not found.",
@@ -32,7 +35,7 @@ class DatasetNotFoundError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class DatasetTypeError(CogneeApiError):
+class DatasetTypeError(CogneeValidationError):
     def __init__(
         self,
         message: str = "Dataset type not supported.",
@@ -40,3 +43,13 @@ class DatasetTypeError(CogneeApiError):
         status_code=status.HTTP_400_BAD_REQUEST,
     ):
         super().__init__(message, name, status_code)
+
+
+class InvalidTableAttributeError(CogneeValidationError):
+    def __init__(
+        self,
+        message: str = "The provided data object is missing the required '__tablename__' attribute.",
+        name: str = "InvalidTableAttributeError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/data/methods/__init__.py

@@ -7,6 +7,7 @@ from .get_datasets import get_datasets
 from .get_datasets_by_name import get_datasets_by_name
 from .get_dataset_data import get_dataset_data
 from .get_authorized_dataset import get_authorized_dataset
+from .get_authorized_dataset_by_name import get_authorized_dataset_by_name
 from .get_data import get_data
 from .get_unique_dataset_id import get_unique_dataset_id
 from .get_authorized_existing_datasets import get_authorized_existing_datasets
@@ -18,6 +19,7 @@ from .delete_data import delete_data
 
 # Create
 from .load_or_create_datasets import load_or_create_datasets
+from .create_authorized_dataset import create_authorized_dataset
 
 # Check
 from .check_dataset_name import check_dataset_name

cognee/modules/data/methods/create_authorized_dataset.py

@@ -0,0 +1,19 @@
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.users.models import User
+from cognee.modules.data.models import Dataset
+from cognee.modules.users.permissions.methods import give_permission_on_dataset
+from .create_dataset import create_dataset
+
+
+async def create_authorized_dataset(dataset_name: str, user: User) -> Dataset:
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        new_dataset = await create_dataset(dataset_name, user, session)
+
+    await give_permission_on_dataset(user, new_dataset.id, "read")
+    await give_permission_on_dataset(user, new_dataset.id, "write")
+    await give_permission_on_dataset(user, new_dataset.id, "delete")
+    await give_permission_on_dataset(user, new_dataset.id, "share")
+
+    return new_dataset

cognee/modules/data/methods/delete_data.py

@@ -1,4 +1,4 @@
-from cognee.exceptions import InvalidAttributeError
+from cognee.modules.data.exceptions.exceptions import InvalidTableAttributeError
 from cognee.modules.data.models import Data
 from cognee.infrastructure.databases.relational import get_relational_engine
 
@@ -13,9 +13,7 @@ async def delete_data(data: Data):
         ValueError: If the data object is invalid.
     """
     if not hasattr(data, "__tablename__"):
-        raise InvalidAttributeError(
-            message="The provided data object is missing the required '__tablename__' attribute."
-        )
+        raise InvalidTableAttributeError()
 
     db_engine = get_relational_engine()
 

cognee/modules/data/methods/get_authorized_dataset.py

@@ -1,11 +1,15 @@
-from typing import Optional
 from uuid import UUID
-from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
+from typing import Optional
+
+from cognee.modules.users.models import User
+from cognee.modules.data.methods.get_authorized_existing_datasets import (
+    get_authorized_existing_datasets,
+)
 from ..models import Dataset
 
 
 async def get_authorized_dataset(
-    user_id: UUID, dataset_id: UUID, permission_type="read"
+    user: User, dataset_id: UUID, permission_type="read"
 ) -> Optional[Dataset]:
     """
     Get a specific dataset with permissions for a user.
@@ -18,6 +22,8 @@ async def get_authorized_dataset(
     Returns:
         Optional[Dataset]: dataset with permissions
     """
-    datasets = await get_specific_user_permission_datasets(user_id, permission_type, [dataset_id])
+    authorized_datasets = await get_authorized_existing_datasets(
+        [dataset_id], permission_type, user
+    )
 
-    return datasets[0] if datasets else None
+    return authorized_datasets[0] if authorized_datasets else None

cognee/modules/data/methods/get_authorized_dataset_by_name.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+from cognee.modules.users.models import User
+from cognee.modules.data.methods.get_authorized_existing_datasets import (
+    get_authorized_existing_datasets,
+)
+
+from ..models import Dataset
+
+
+async def get_authorized_dataset_by_name(
+    dataset_name: str, user: User, permission_type: str
+) -> Optional[Dataset]:
+    authorized_datasets = await get_authorized_existing_datasets([], permission_type, user)
+
+    return next((dataset for dataset in authorized_datasets if dataset.name == dataset_name), None)

cognee/modules/data/methods/load_or_create_datasets.py

@@ -1,12 +1,9 @@
 from typing import List, Union
 from uuid import UUID
 
-from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.data.models import Dataset
-from cognee.modules.data.methods import create_dataset
-from cognee.modules.data.methods import get_unique_dataset_id
+from cognee.modules.data.methods import create_authorized_dataset
 from cognee.modules.data.exceptions import DatasetNotFoundError
-from cognee.modules.users.permissions.methods import give_permission_on_dataset
 
 
 async def load_or_create_datasets(
@@ -34,22 +31,7 @@ async def load_or_create_datasets(
         if isinstance(identifier, UUID):
             raise DatasetNotFoundError(f"Dataset with given UUID does not exist: {identifier}")
 
-        # Otherwise, create a new Dataset instance
-        new_dataset = Dataset(
-            id=await get_unique_dataset_id(dataset_name=identifier, user=user),
-            name=identifier,
-            owner_id=user.id,
-        )
-
-        # Save dataset to database
-        db_engine = get_relational_engine()
-        async with db_engine.get_async_session() as session:
-            await create_dataset(identifier, user, session)
-
-        await give_permission_on_dataset(user, new_dataset.id, "read")
-        await give_permission_on_dataset(user, new_dataset.id, "write")
-        await give_permission_on_dataset(user, new_dataset.id, "delete")
-        await give_permission_on_dataset(user, new_dataset.id, "share")
+        new_dataset = await create_authorized_dataset(identifier, user)
 
         result.append(new_dataset)
 

cognee/modules/data/processing/document_types/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeSystemError
 from fastapi import status
 
 
-class PyPdfInternalError(CogneeApiError):
+class PyPdfInternalError(CogneeSystemError):
     """Internal pypdf error"""
 
     def __init__(

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -2,8 +2,11 @@ import time
 from cognee.shared.logging_utils import get_logger
 from typing import List, Dict, Union, Optional, Type
 
-from cognee.exceptions import InvalidValueError
-from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError
+from cognee.modules.graph.exceptions import (
+    EntityNotFoundError,
+    EntityAlreadyExistsError,
+    InvalidDimensionsError,
+)
 from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
 from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph
@@ -66,8 +69,7 @@ class CogneeGraph(CogneeAbstractGraph):
         node_name: Optional[List[str]] = None,
     ) -> None:
         if node_dimension < 1 or edge_dimension < 1:
-            raise InvalidValueError(message="Dimensions must be positive integers")
-
+            raise InvalidDimensionsError()
         try:
             import time
 

cognee/modules/graph/cognee_graph/CogneeGraphElements.py

@@ -1,7 +1,6 @@
 import numpy as np
 from typing import List, Dict, Optional, Any, Union
-
-from cognee.exceptions import InvalidValueError
+from cognee.modules.graph.exceptions import InvalidDimensionsError, DimensionOutOfRangeError
 
 
 class Node:
@@ -24,7 +23,7 @@ class Node:
         self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1
     ):
         if dimension <= 0:
-            raise InvalidValueError(message="Dimension must be a positive integer")
+            raise InvalidDimensionsError()
         self.id = node_id
         self.attributes = attributes if attributes is not None else {}
         self.attributes["vector_distance"] = float("inf")
@@ -58,9 +57,7 @@ class Node:
 
     def is_node_alive_in_dimension(self, dimension: int) -> bool:
         if dimension < 0 or dimension >= len(self.status):
-            raise InvalidValueError(
-                message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}."
-            )
+            raise DimensionOutOfRangeError(dimension=dimension, max_index=len(self.status) - 1)
         return self.status[dimension] == 1
 
     def add_attribute(self, key: str, value: Any) -> None:
@@ -110,7 +107,7 @@ class Edge:
         dimension: int = 1,
     ):
         if dimension <= 0:
-            raise InvalidValueError(message="Dimensions must be a positive integer.")
+            raise InvalidDimensionsError()
         self.node1 = node1
         self.node2 = node2
         self.attributes = attributes if attributes is not None else {}
@@ -120,9 +117,7 @@ class Edge:
 
     def is_edge_alive_in_dimension(self, dimension: int) -> bool:
         if dimension < 0 or dimension >= len(self.status):
-            raise InvalidValueError(
-                message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}."
-            )
+            raise DimensionOutOfRangeError(dimension=dimension, max_index=len(self.status) - 1)
         return self.status[dimension] == 1
 
     def add_attribute(self, key: str, value: Any) -> None:

cognee/modules/graph/exceptions/__init__.py

@@ -7,4 +7,6 @@ This module defines a set of exceptions for handling various graph errors
 from .exceptions import (
     EntityNotFoundError,
     EntityAlreadyExistsError,
+    InvalidDimensionsError,
+    DimensionOutOfRangeError,
 )

cognee/modules/graph/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeValidationError
 from fastapi import status
 
 
-class EntityNotFoundError(CogneeApiError):
+class EntityNotFoundError(CogneeValidationError):
     """Database returns nothing"""
 
     def __init__(
@@ -14,7 +14,7 @@ class EntityNotFoundError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class EntityAlreadyExistsError(CogneeApiError):
+class EntityAlreadyExistsError(CogneeValidationError):
     """Conflict detected, like trying to create a resource that already exists"""
 
     def __init__(
@@ -24,3 +24,25 @@ class EntityAlreadyExistsError(CogneeApiError):
         status_code=status.HTTP_409_CONFLICT,
     ):
         super().__init__(message, name, status_code)
+
+
+class InvalidDimensionsError(CogneeValidationError):
+    def __init__(
+        self,
+        name: str = "InvalidDimensionsError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        message = "Dimensions must be positive integers."
+        super().__init__(message, name, status_code)
+
+
+class DimensionOutOfRangeError(CogneeValidationError):
+    def __init__(
+        self,
+        dimension: int,
+        max_index: int,
+        name: str = "DimensionOutOfRangeError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        message = f"Dimension {dimension} is out of range. Valid range is 0 to {max_index}."
+        super().__init__(message, name, status_code)

cognee/modules/graph/methods/get_formatted_graph_data.py

@@ -3,10 +3,11 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.context_global_variables import set_database_global_context_variables
 from cognee.modules.data.exceptions.exceptions import DatasetNotFoundError
 from cognee.modules.data.methods import get_authorized_dataset
+from cognee.modules.users.models import User
 
 
-async def get_formatted_graph_data(dataset_id: UUID, user_id: UUID):
-    dataset = await get_authorized_dataset(user_id, dataset_id)
+async def get_formatted_graph_data(dataset_id: UUID, user: User):
+    dataset = await get_authorized_dataset(user, dataset_id)
     if not dataset:
         raise DatasetNotFoundError(message="Dataset not found.")
 

cognee/modules/ingestion/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeValidationError
 from fastapi import status
 
 
-class IngestionError(CogneeApiError):
+class IngestionError(CogneeValidationError):
     def __init__(
         self,
         message: str = "Type of data sent to classify not supported.",

cognee/modules/ontology/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeSystemError
 from fastapi import status
 
 
-class OntologyInitializationError(CogneeApiError):
+class OntologyInitializationError(CogneeSystemError):
     def __init__(
         self,
         message: str = "Ontology initialization failed",
@@ -12,7 +12,7 @@ class OntologyInitializationError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class FindClosestMatchError(CogneeApiError):
+class FindClosestMatchError(CogneeSystemError):
     def __init__(
         self,
         message: str = "Error in find_closest_match",
@@ -22,7 +22,7 @@ class FindClosestMatchError(CogneeApiError):
         super().__init__(message, name, status_code)
 
 
-class GetSubgraphError(CogneeApiError):
+class GetSubgraphError(CogneeSystemError):
     def __init__(
         self,
         message: str = "Failed to retrieve subgraph",

cognee/modules/pipelines/__init__.py

@@ -1,4 +1,4 @@
 from .tasks.task import Task
 from .operations.run_tasks import run_tasks
 from .operations.run_parallel import run_tasks_parallel
-from .operations.pipeline import cognee_pipeline
+from .operations.pipeline import run_pipeline

cognee/modules/pipelines/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeSystemError
 from fastapi import status
 
 
-class PipelineRunFailedError(CogneeApiError):
+class PipelineRunFailedError(CogneeSystemError):
     def __init__(
         self,
         message: str = "Pipeline run failed.",

cognee/modules/pipelines/exceptions/tasks.py

@@ -0,0 +1,18 @@
+from fastapi import status
+from cognee.exceptions import CogneeValidationError
+
+
+class WrongTaskTypeError(CogneeValidationError):
+    """
+    Raised when the tasks argument is not a list of Task class instances.
+    """
+
+    def __init__(
+        self,
+        message: str = "tasks argument must be a list, containing Task class instances.",
+        name: str = "WrongTaskTypeError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        self.message = message
+        self.name = name
+        self.status_code = status_code

cognee/modules/pipelines/layers/__init__.py

@@ -0,0 +1 @@
+from .validate_pipeline_tasks import validate_pipeline_tasks

cognee/modules/pipelines/layers/check_pipeline_run_qualification.py

@@ -0,0 +1,59 @@
+from typing import Union, Optional
+from cognee.modules.data.models import Dataset
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRunStatus
+from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
+from cognee.modules.pipelines.methods import get_pipeline_run_by_dataset
+from cognee.shared.logging_utils import get_logger
+
+from cognee.modules.pipelines.models.PipelineRunInfo import (
+    PipelineRunCompleted,
+    PipelineRunStarted,
+)
+
+logger = get_logger(__name__)
+
+
+async def check_pipeline_run_qualification(
+    dataset: Dataset, data: list[Data], pipeline_name: str
+) -> Optional[Union[PipelineRunStarted, PipelineRunCompleted]]:
+    """
+    Function used to determine if pipeline is currently being processed or was already processed.
+    In case pipeline was or is being processed return value is returned and current pipline execution should be stopped.
+    In case pipeline is not or was not processed there will be no return value and pipeline processing can start.
+
+    Args:
+        dataset: Dataset object
+        data: List of Data
+        pipeline_name: pipeline name
+
+    Returns: Pipeline state if it is being processed or was already processed
+
+    """
+
+    # async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
+    if isinstance(dataset, Dataset):
+        task_status = await get_pipeline_status([dataset.id], pipeline_name)
+    else:
+        task_status = {}
+
+    if str(dataset.id) in task_status:
+        if task_status[str(dataset.id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED:
+            logger.info("Dataset %s is already being processed.", dataset.id)
+            pipeline_run = await get_pipeline_run_by_dataset(dataset.id, pipeline_name)
+            return PipelineRunStarted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+                payload=data,
+            )
+        elif task_status[str(dataset.id)] == PipelineRunStatus.DATASET_PROCESSING_COMPLETED:
+            logger.info("Dataset %s is already processed.", dataset.id)
+            pipeline_run = await get_pipeline_run_by_dataset(dataset.id, pipeline_name)
+            return PipelineRunCompleted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+            )
+
+    return