Merge branch 'main' into dev

README.md

@@ -176,16 +176,6 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-### Installation for UI
-
-To use the cognee UI with full functionality, you need to install cognee with API dependencies:
-
-```bash
-pip install 'cognee[api]'
-```
-
-The UI requires backend server functionality (uvicorn and other API dependencies) which are not included in the default cognee installation to keep it lightweight.
-
 ### Running the UI
 
 Try cognee UI by running ``` cognee-cli -ui ``` command on your terminal.

cognee/modules/data/methods/create_authorized_dataset.py

@@ -6,6 +6,15 @@ from .create_dataset import create_dataset
 
 
 async def create_authorized_dataset(dataset_name: str, user: User) -> Dataset:
+    """
+        Create a new dataset and give all permissions on this dataset to the given user.
+    Args:
+        dataset_name: Name of the dataset.
+        user: The user object.
+
+    Returns:
+        Dataset: The new authorized dataset.
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/data/methods/get_authorized_dataset.py

@@ -15,7 +15,7 @@ async def get_authorized_dataset(
     Get a specific dataset with permissions for a user.
 
     Args:
-        user_id (UUID): user id
+        user: User object
         dataset_id (UUID): dataset id
         permission_type (str): permission type(read, write, delete, share), default is read
 

cognee/modules/data/methods/get_authorized_dataset_by_name.py

@@ -11,6 +11,17 @@ from ..models import Dataset
 async def get_authorized_dataset_by_name(
     dataset_name: str, user: User, permission_type: str
 ) -> Optional[Dataset]:
+    """
+    Get a specific dataset with the given name, with permissions for a given user.
+
+    Args:
+        dataset_name: Name of the dataset.
+        user: User object.
+        permission_type (str): permission type(read, write, delete, share), default is read
+
+    Returns:
+        Optional[Dataset]: dataset with permissions
+    """
     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/graph/utils/retrieve_existing_edges.py

@@ -23,8 +23,6 @@ async def retrieve_existing_edges(
         chunk_graphs (list[KnowledgeGraph]): List of knowledge graphs corresponding to each
             data chunk. Each graph contains nodes (entities) and edges (relationships) that
             were extracted from the chunk content.
-        graph_engine (GraphDBInterface): Interface to the graph database that will be queried
-            to check for existing edges. Must implement the has_edges() method.
 
     Returns:
         dict[str, bool]: A mapping of edge keys to boolean values indicating existence.

cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py

@@ -11,6 +11,19 @@ from cognee.modules.data.methods import (
 
 
 async def resolve_authorized_user_dataset(dataset_id: UUID, dataset_name: str, user: User):
+    """
+    Function handles creation and dataset authorization if dataset already exist for Cognee.
+    Verifies that provided user has necessary permission for provided Dataset.
+    If Dataset does not exist creates the Dataset and gives permission for the user creating the dataset.
+
+    Args:
+        dataset_id: Id of the dataset.
+        dataset_name: Name of the dataset.
+        user: Cognee User request is being processed for, if None default user will be used.
+
+    Returns:
+        Tuple[User, Dataset]: A tuple containing the user and the authorized dataset.
+    """
     if not user:
         user = await get_default_user()
 

cognee/modules/pipelines/layers/resolve_authorized_user_datasets.py

@@ -25,7 +25,7 @@ async def resolve_authorized_user_datasets(
         datasets: Dataset names or Dataset UUID (in case Datasets already exist)
 
     Returns:
-
+        Tuple[User, List[Dataset]]: A tuple containing the user and the list of authorized datasets.
     """
     # If no user is provided use default user
     if user is None:

cognee/modules/users/permissions/methods/authorized_give_permission_on_datasets.py

@@ -9,6 +9,18 @@ from uuid import UUID
 async def authorized_give_permission_on_datasets(
     principal_id: UUID, dataset_ids: Union[List[UUID], UUID], permission_name: str, owner_id: UUID
 ):
+    """
+        Give permission to certain datasets to a user.
+        The request owner must have the necessary permission to share the datasets.
+    Args:
+        principal_id: Id of user to whom datasets are shared
+        dataset_ids: Ids of datasets to share
+        permission_name: Name of permission to give
+        owner_id: Id of the request owner
+
+    Returns:
+        None
+    """
     # If only a single dataset UUID is provided transform it to a list
     if not isinstance(dataset_ids, list):
         dataset_ids = [dataset_ids]

cognee/modules/users/permissions/methods/check_permission_on_dataset.py

@@ -10,6 +10,17 @@ logger = get_logger()
 
 
 async def check_permission_on_dataset(user: User, permission_type: str, dataset_id: UUID):
+    """
+        Check if a user has a specific permission on a dataset.
+    Args:
+        user: User whose permission is checked
+        permission_type: Type of permission to check
+        dataset_id: Id of the dataset
+
+    Returns:
+        None
+
+    """
     if user is None:
         user = await get_default_user()
 

cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py

@@ -11,6 +11,16 @@ logger = get_logger()
 
 
 async def get_all_user_permission_datasets(user: User, permission_type: str) -> list[Dataset]:
+    """
+        Return a list of datasets the user has permission for.
+        If the user is part of a tenant, return datasets his roles have permission for.
+    Args:
+        user
+        permission_type
+
+    Returns:
+        list[Dataset]: List of datasets user has permission for
+    """
     datasets = list()
     # Get all datasets User has explicit access to
     datasets.extend(await get_principal_datasets(user, permission_type))

cognee/modules/users/permissions/methods/get_document_ids_for_user.py

@@ -8,6 +8,16 @@ from ...models import ACL, Permission
 
 
 async def get_document_ids_for_user(user_id: UUID, datasets: list[str] = None) -> list[str]:
+    """
+        Return a list of documents ids for which the user has read permission.
+        If datasets are specified, return only documents from those datasets.
+    Args:
+        user_id: Id of the user
+        datasets: List of datasets
+
+    Returns:
+        list[str]: List of documents for which the user has read permission
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_principal.py

@@ -6,6 +6,15 @@ from ...models.Principal import Principal
 
 
 async def get_principal(principal_id: UUID):
+    """
+        Return information about a user based on their id
+    Args:
+        principal_id: Id of the user
+
+    Returns:
+        principal: Information about the user (principal)
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_principal_datasets.py

@@ -9,6 +9,17 @@ from ...models.ACL import ACL
 
 
 async def get_principal_datasets(principal: Principal, permission_type: str) -> list[Dataset]:
+    """
+        Return a list of datasets for which the user (principal) has a certain permission.
+    Args:
+        principal: Information about the user
+        permission_type: Type of permission
+
+    Returns:
+        list[Dataset]: List of datasets for which the user (principal)
+        has the permission (permission_type).
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_role.py

@@ -9,6 +9,16 @@ from ...models.Role import Role
 
 
 async def get_role(tenant_id: UUID, role_name: str):
+    """
+        Return the role with the name role_name of the given tenant.
+    Args:
+        tenant_id: Id of the given tenant
+        role_name: Name of the role
+
+    Returns
+        The role for the given tenant.
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/get_specific_user_permission_datasets.py

@@ -15,9 +15,9 @@ async def get_specific_user_permission_datasets(
         Return a list of datasets user has given permission for. If a list of datasets is provided,
         verify for which datasets user has appropriate permission for and return list of datasets he has permission for.
     Args:
-        user_id:
+        user_id: Id of the user.
-        permission_type:
+        permission_type: Type of the permission.
-        dataset_ids:
+        dataset_ids: Ids of the provided datasets
 
     Returns:
         list[Dataset]: List of datasets user has permission for

cognee/modules/users/permissions/methods/get_tenant.py

@@ -8,6 +8,15 @@ from ...models.Tenant import Tenant
 
 
 async def get_tenant(tenant_id: UUID):
+    """
+        Return information about the tenant based on the given id.
+    Args:
+        tenant_id: Id of the given tenant
+
+    Returns
+        Information about the given tenant.
+
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/give_default_permission_to_role.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_role(role_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the role with the given id as a default permission.
+    Args:
+        role_id: Id of the role
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/permissions/methods/give_default_permission_to_tenant.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_tenant(tenant_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the tenant with the given id as a default permission.
+    Args:
+        tenant_id: Id of the tenant
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         tenant = (

cognee/modules/users/permissions/methods/give_default_permission_to_user.py

@@ -16,6 +16,15 @@ from cognee.modules.users.models import (
 
 
 async def give_default_permission_to_user(user_id: UUID, permission_name: str):
+    """
+        Give the permission with given name to the user with the given id as a default permission.
+    Args:
+        user_id: Id of the tenant
+        permission_name: Name of the permission
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = (await session.execute(select(User).where(User.id == user_id))).scalars().first()

cognee/modules/users/permissions/methods/give_permission_on_dataset.py

@@ -24,6 +24,16 @@ async def give_permission_on_dataset(
     dataset_id: UUID,
     permission_name: str,
 ):
+    """
+        Give a specific permission on a dataset to a user.
+    Args:
+        principal: User who is being given the permission on the dataset
+        dataset_id: Id of the dataset
+        permission_name: Name of permission to give
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/users/roles/methods/add_user_to_role.py

@@ -21,6 +21,17 @@ from cognee.modules.users.models import (
 
 
 async def add_user_to_role(user_id: UUID, role_id: UUID, owner_id: UUID):
+    """
+        Add a user with the given id to the role with the given id.
+    Args:
+        user_id: Id of the user.
+        role_id: Id of the role.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = (await session.execute(select(User).where(User.id == user_id))).scalars().first()

cognee/modules/users/roles/methods/create_role.py

@@ -16,6 +16,16 @@ async def create_role(
     role_name: str,
     owner_id: UUID,
 ):
+    """
+        Create a new role with the given name, if the request owner with the given id
+        has the necessary permission.
+    Args:
+        role_name: Name of the new role.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(owner_id)

cognee/modules/users/tenants/methods/add_user_to_tenant.py

@@ -13,6 +13,18 @@ from cognee.modules.users.exceptions import (
 
 
 async def add_user_to_tenant(user_id: UUID, tenant_id: UUID, owner_id: UUID):
+    """
+        Add a user with the given id to the tenant with the given id.
+        This can only be successful if the request owner with the given id is the tenant owner.
+    Args:
+        user_id: Id of the user.
+        tenant_id: Id of the tenant.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(user_id)

cognee/modules/users/tenants/methods/create_tenant.py

@@ -8,6 +8,16 @@ from cognee.modules.users.methods import get_user
 
 
 async def create_tenant(tenant_name: str, user_id: UUID):
+    """
+        Create a new tenant with the given name, for the user with the given id.
+        This user is the owner of the tenant.
+    Args:
+        tenant_name: Name of the new tenant.
+        user_id: Id of the user.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         try:

pyproject.toml

@@ -63,14 +63,13 @@ dependencies = [
     "pylance>=0.22.0,<1.0.0",
     "kuzu (==0.11.0)",
     "python-magic-bin<0.5 ; platform_system == 'Windows'", # Only needed for Windows
-]
-
-[project.optional-dependencies]
-api = [
     "uvicorn>=0.34.0,<1.0.0",
     "gunicorn>=20.1.0,<24",
     "websockets>=15.0.1,<16.0.0"
 ]
+
+[project.optional-dependencies]
+api=[]
 distributed = [
     "modal>=1.0.5,<2.0.0",
 ]

uv.lock

@@ -823,6 +823,7 @@ dependencies = [
     { name = "fastapi" },
     { name = "fastapi-users", extra = ["sqlalchemy"] },
     { name = "filetype" },
+    { name = "gunicorn" },
     { name = "instructor" },
     { name = "jinja2" },
     { name = "kuzu" },
@@ -857,17 +858,14 @@ dependencies = [
     { name = "structlog" },
     { name = "tiktoken" },
     { name = "typing-extensions" },
+    { name = "uvicorn" },
+    { name = "websockets" },
 ]
 
 [package.optional-dependencies]
 anthropic = [
     { name = "anthropic" },
 ]
-api = [
-    { name = "gunicorn" },
-    { name = "uvicorn" },
-    { name = "websockets" },
-]
 aws = [
     { name = "s3fs", extra = ["boto3"] },
 ]
@@ -994,7 +992,7 @@ requires-dist = [
     { name = "google-generativeai", marker = "extra == 'gemini'", specifier = ">=0.8.4,<0.9" },
     { name = "graphiti-core", marker = "extra == 'graphiti'", specifier = ">=0.7.0,<0.8" },
     { name = "groq", marker = "extra == 'groq'", specifier = ">=0.8.0,<1.0.0" },
-    { name = "gunicorn", marker = "extra == 'api'", specifier = ">=20.1.0,<24" },
+    { name = "gunicorn", specifier = ">=20.1.0,<24" },
     { name = "instructor", specifier = ">=1.9.1,<2.0.0" },
     { name = "jinja2", specifier = ">=3.1.3,<4" },
     { name = "kuzu", specifier = "==0.11.0" },
@@ -1062,8 +1060,8 @@ requires-dist = [
     { name = "tweepy", marker = "extra == 'dev'", specifier = ">=4.14.0,<5.0.0" },
     { name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" },
     { name = "unstructured", extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], marker = "extra == 'docs'", specifier = ">=0.18.1,<19" },
-    { name = "uvicorn", marker = "extra == 'api'", specifier = ">=0.34.0,<1.0.0" },
+    { name = "uvicorn", specifier = ">=0.34.0,<1.0.0" },
-    { name = "websockets", marker = "extra == 'api'", specifier = ">=15.0.1,<16.0.0" },
+    { name = "websockets", specifier = ">=15.0.1,<16.0.0" },
 ]
 provides-extras = ["api", "distributed", "neo4j", "neptune", "postgres", "postgres-binary", "notebook", "langchain", "llama-index", "gemini", "huggingface", "ollama", "mistral", "anthropic", "deepeval", "posthog", "falkordb", "groq", "chromadb", "docs", "codegraph", "evals", "gui", "graphiti", "aws", "dev", "debug"]