cognee 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/modules/users/methods/get_authenticated_user.py

@@ -5,14 +5,15 @@ from ..models import User
 from ..get_fastapi_users import get_fastapi_users
 from .get_default_user import get_default_user
 from cognee.shared.logging_utils import get_logger
+from cognee.context_global_variables import backend_access_control_enabled
 
 
 logger = get_logger("get_authenticated_user")
 
 # Check environment variable to determine authentication requirement
 REQUIRE_AUTHENTICATION = (
-    os.getenv("REQUIRE_AUTHENTICATION", "false").lower() == "true"
-    or os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true"
+    os.getenv("REQUIRE_AUTHENTICATION", "true").lower() == "true"
+    or os.environ.get("ENABLE_BACKEND_ACCESS_CONTROL", "true").lower() == "true"
 )
 
 fastapi_users = get_fastapi_users()

cognee/modules/users/methods/get_default_user.py

@@ -10,7 +10,7 @@ from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.users.methods.create_default_user import create_default_user
 
 
-async def get_default_user() -> SimpleNamespace:
+async def get_default_user() -> User:
     db_engine = get_relational_engine()
     base_config = get_base_config()
     default_email = base_config.default_user_email or "default_user@example.com"
@@ -18,7 +18,9 @@ async def get_default_user() -> SimpleNamespace:
     try:
         async with db_engine.get_async_session() as session:
             query = (
-                select(User).options(selectinload(User.roles)).where(User.email == default_email)
+                select(User)
+                .options(selectinload(User.roles), selectinload(User.tenants))
+                .where(User.email == default_email)
             )
 
             result = await session.execute(query)

cognee/modules/users/methods/get_user.py

@@ -14,7 +14,7 @@ async def get_user(user_id: UUID):
         user = (
             await session.execute(
                 select(User)
-                .options(selectinload(User.roles), selectinload(User.tenant))
+                .options(selectinload(User.roles), selectinload(User.tenants))
                 .where(User.id == user_id)
             )
         ).scalar()

cognee/modules/users/methods/get_user_by_email.py

@@ -13,7 +13,7 @@ async def get_user_by_email(user_email: str):
         user = (
             await session.execute(
                 select(User)
-                .options(joinedload(User.roles), joinedload(User.tenant))
+                .options(joinedload(User.roles), joinedload(User.tenants))
                 .where(User.email == user_email)
             )
         ).scalar()

cognee/modules/users/models/DatasetDatabase.py

@@ -1,6 +1,6 @@
 from datetime import datetime, timezone
 
-from sqlalchemy import Column, DateTime, String, UUID, ForeignKey
+from sqlalchemy import Column, DateTime, String, UUID, ForeignKey, JSON, text
 from cognee.infrastructure.databases.relational import Base
 
 
@@ -12,8 +12,29 @@ class DatasetDatabase(Base):
         UUID, ForeignKey("datasets.id", ondelete="CASCADE"), primary_key=True, index=True
     )
 
-    vector_database_name = Column(String, unique=True, nullable=False)
-    graph_database_name = Column(String, unique=True, nullable=False)
+    vector_database_name = Column(String, unique=False, nullable=False)
+    graph_database_name = Column(String, unique=False, nullable=False)
+
+    vector_database_provider = Column(String, unique=False, nullable=False)
+    graph_database_provider = Column(String, unique=False, nullable=False)
+
+    graph_dataset_database_handler = Column(String, unique=False, nullable=False)
+    vector_dataset_database_handler = Column(String, unique=False, nullable=False)
+
+    vector_database_url = Column(String, unique=False, nullable=True)
+    graph_database_url = Column(String, unique=False, nullable=True)
+
+    vector_database_key = Column(String, unique=False, nullable=True)
+    graph_database_key = Column(String, unique=False, nullable=True)
+
+    # configuration details for different database types. This would make it more flexible to add new database types
+    # without changing the database schema.
+    graph_database_connection_info = Column(
+        JSON, unique=False, nullable=False, server_default=text("'{}'")
+    )
+    vector_database_connection_info = Column(
+        JSON, unique=False, nullable=False, server_default=text("'{}'")
+    )
 
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))

cognee/modules/users/models/Tenant.py

@@ -1,7 +1,7 @@
-from sqlalchemy.orm import relationship
+from sqlalchemy.orm import relationship, Mapped
 from sqlalchemy import Column, String, ForeignKey, UUID
 from .Principal import Principal
-from .User import User
+from .UserTenant import UserTenant
 from .Role import Role
 
 
@@ -13,14 +13,13 @@ class Tenant(Principal):
 
     owner_id = Column(UUID, index=True)
 
-    # One-to-Many relationship with User; specify the join via User.tenant_id
-    users = relationship(
+    users: Mapped[list["User"]] = relationship(  # noqa: F821
         "User",
-        back_populates="tenant",
-        foreign_keys=lambda: [User.tenant_id],
+        secondary=UserTenant.__tablename__,
+        back_populates="tenants",
     )
 
-    # One-to-Many relationship with Role (if needed; similar fix)
+    # One-to-Many relationship with Role
     roles = relationship(
         "Role",
         back_populates="tenant",

cognee/modules/users/models/User.py

@@ -6,8 +6,10 @@ from sqlalchemy import ForeignKey, Column, UUID
 from sqlalchemy.orm import relationship, Mapped
 
 from .Principal import Principal
+from .UserTenant import UserTenant
 from .UserRole import UserRole
 from .Role import Role
+from .Tenant import Tenant
 
 
 class User(SQLAlchemyBaseUserTableUUID, Principal):
@@ -15,7 +17,7 @@ class User(SQLAlchemyBaseUserTableUUID, Principal):
 
     id = Column(UUID, ForeignKey("principals.id", ondelete="CASCADE"), primary_key=True)
 
-    # Foreign key to Tenant (Many-to-One relationship)
+    # Foreign key to current Tenant (Many-to-One relationship)
     tenant_id = Column(UUID, ForeignKey("tenants.id"))
 
     # Many-to-Many Relationship with Roles
@@ -25,11 +27,11 @@ class User(SQLAlchemyBaseUserTableUUID, Principal):
         back_populates="users",
     )
 
-    # Relationship to Tenant
-    tenant = relationship(
+    # Many-to-Many Relationship with Tenants user is a part of
+    tenants: Mapped[list["Tenant"]] = relationship(
         "Tenant",
+        secondary=UserTenant.__tablename__,
         back_populates="users",
-        foreign_keys=[tenant_id],
     )
 
     # ACL Relationship (One-to-Many)
@@ -46,7 +48,6 @@ class UserRead(schemas.BaseUser[uuid_UUID]):
 
 
 class UserCreate(schemas.BaseUserCreate):
-    tenant_id: Optional[uuid_UUID] = None
     is_verified: bool = True
 
 

cognee/modules/users/models/UserTenant.py

@@ -0,0 +1,12 @@
+from datetime import datetime, timezone
+from sqlalchemy import Column, ForeignKey, DateTime, UUID
+from cognee.infrastructure.databases.relational import Base
+
+
+class UserTenant(Base):
+    __tablename__ = "user_tenants"
+
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+
+    user_id = Column(UUID, ForeignKey("users.id"), primary_key=True)
+    tenant_id = Column(UUID, ForeignKey("tenants.id"), primary_key=True)

cognee/modules/users/models/__init__.py

@@ -1,6 +1,7 @@
 from .User import User
 from .Role import Role
 from .UserRole import UserRole
+from .UserTenant import UserTenant
 from .DatasetDatabase import DatasetDatabase
 from .RoleDefaultPermissions import RoleDefaultPermissions
 from .UserDefaultPermissions import UserDefaultPermissions

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

@@ -1,11 +1,8 @@
-from types import SimpleNamespace
-
 from cognee.shared.logging_utils import get_logger
 
 from ...models.User import User
 from cognee.modules.data.models.Dataset import Dataset
 from cognee.modules.users.permissions.methods import get_principal_datasets
-from cognee.modules.users.permissions.methods import get_role, get_tenant
 
 logger = get_logger()
 
@@ -25,17 +22,14 @@ async def get_all_user_permission_datasets(user: User, permission_type: str) ->
     # Get all datasets User has explicit access to
     datasets.extend(await get_principal_datasets(user, permission_type))
 
-    if user.tenant_id:
-        # Get all datasets all tenants have access to
-        tenant = await get_tenant(user.tenant_id)
+    # Get all tenants user is a part of
+    tenants = await user.awaitable_attrs.tenants
+    for tenant in tenants:
+        # Get all datasets all tenant members have access to
         datasets.extend(await get_principal_datasets(tenant, permission_type))
 
-        # Get all datasets Users roles have access to
-        if isinstance(user, SimpleNamespace):
-            # If simple namespace use roles defined in user
-            roles = user.roles
-        else:
-            roles = await user.awaitable_attrs.roles
+        # Get all datasets accessible by roles user is a part of
+        roles = await user.awaitable_attrs.roles
         for role in roles:
             datasets.extend(await get_principal_datasets(role, permission_type))
 
@@ -45,4 +39,10 @@ async def get_all_user_permission_datasets(user: User, permission_type: str) ->
         # If the dataset id key already exists, leave the dictionary unchanged.
         unique.setdefault(dataset.id, dataset)
 
-    return list(unique.values())
+    # Filter out dataset that aren't part of the selected user's tenant
+    filtered_datasets = []
+    for dataset in list(unique.values()):
+        if dataset.tenant_id == user.tenant_id:
+            filtered_datasets.append(dataset)
+
+    return filtered_datasets

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

@@ -42,11 +42,13 @@ async def add_user_to_role(user_id: UUID, role_id: UUID, owner_id: UUID):
             .first()
         )
 
+        user_tenants = await user.awaitable_attrs.tenants
+
         if not user:
             raise UserNotFoundError
         elif not role:
             raise RoleNotFoundError
-        elif user.tenant_id != role.tenant_id:
+        elif role.tenant_id not in [tenant.id for tenant in user_tenants]:
             raise TenantNotFoundError(
                 message="User tenant does not match role tenant. User cannot be added to role."
             )

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

@@ -1,2 +1,3 @@
 from .create_tenant import create_tenant
 from .add_user_to_tenant import add_user_to_tenant
+from .select_tenant import select_tenant

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

@@ -1,8 +1,11 @@
+from typing import Optional
 from uuid import UUID
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy import insert
 
 from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.users.models.UserTenant import UserTenant
 from cognee.modules.users.methods import get_user
 from cognee.modules.users.permissions.methods import get_tenant
 from cognee.modules.users.exceptions import (
@@ -12,14 +15,19 @@ from cognee.modules.users.exceptions import (
 )
 
 
-async def add_user_to_tenant(user_id: UUID, tenant_id: UUID, owner_id: UUID):
+async def add_user_to_tenant(
+    user_id: UUID, tenant_id: UUID, owner_id: UUID, set_as_active_tenant: Optional[bool] = False
+):
     """
         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.
+
+        If set_as_active_tenant is true it will automatically set the users active tenant to provided tenant.
     Args:
         user_id: Id of the user.
         tenant_id: Id of the tenant.
         owner_id: Id of the request owner.
+        set_as_active_tenant: If set_as_active_tenant is true it will automatically set the users active tenant to provided tenant.
 
     Returns:
         None
@@ -40,17 +48,18 @@ async def add_user_to_tenant(user_id: UUID, tenant_id: UUID, owner_id: UUID):
                 message="Only tenant owner can add other users to organization."
             )
 
-        try:
-            if user.tenant_id is None:
-                user.tenant_id = tenant_id
-            elif user.tenant_id == tenant_id:
-                return
-            else:
-                raise IntegrityError
-
+        if set_as_active_tenant:
+            user.tenant_id = tenant_id
             await session.merge(user)
             await session.commit()
-        except IntegrityError:
-            raise EntityAlreadyExistsError(
-                message="User is already part of a tenant. Only one tenant can be assigned to user."
+
+        try:
+            # Add association directly to the association table
+            create_user_tenant_statement = insert(UserTenant).values(
+                user_id=user_id, tenant_id=tenant_id
             )
+            await session.execute(create_user_tenant_statement)
+            await session.commit()
+
+        except IntegrityError:
+            raise EntityAlreadyExistsError(message="User is already part of group.")

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

@@ -1,19 +1,25 @@
 from uuid import UUID
+from sqlalchemy import insert
 from sqlalchemy.exc import IntegrityError
+from typing import Optional
 
+from cognee.modules.users.models.UserTenant import UserTenant
 from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.users.models import Tenant
 from cognee.modules.users.methods import get_user
 
 
-async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
+async def create_tenant(
+    tenant_name: str, user_id: UUID, set_as_active_tenant: Optional[bool] = True
+) -> 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.
+        set_as_active_tenant: If true, set the newly created tenant as the active tenant for the user.
 
     Returns:
         None
@@ -22,18 +28,26 @@ async def create_tenant(tenant_name: str, user_id: UUID) -> UUID:
     async with db_engine.get_async_session() as session:
         try:
             user = await get_user(user_id)
-            if user.tenant_id:
-                raise EntityAlreadyExistsError(
-                    message="User already has a tenant. New tenant cannot be created."
-                )
 
             tenant = Tenant(name=tenant_name, owner_id=user_id)
             session.add(tenant)
             await session.flush()
 
-            user.tenant_id = tenant.id
-            await session.merge(user)
-            await session.commit()
+            if set_as_active_tenant:
+                user.tenant_id = tenant.id
+                await session.merge(user)
+                await session.commit()
+
+            try:
+                # Add association directly to the association table
+                create_user_tenant_statement = insert(UserTenant).values(
+                    user_id=user_id, tenant_id=tenant.id
+                )
+                await session.execute(create_user_tenant_statement)
+                await session.commit()
+            except IntegrityError:
+                raise EntityAlreadyExistsError(message="User is already part of tenant.")
+
             return tenant.id
         except IntegrityError as e:
             raise EntityAlreadyExistsError(message="Tenant already exists.") from e

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

@@ -0,0 +1,62 @@
+from uuid import UUID
+from typing import Union
+
+import sqlalchemy.exc
+from sqlalchemy import select
+
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.users.methods.get_user import get_user
+from cognee.modules.users.models.UserTenant import UserTenant
+from cognee.modules.users.models.User import User
+from cognee.modules.users.permissions.methods import get_tenant
+from cognee.modules.users.exceptions import UserNotFoundError, TenantNotFoundError
+
+
+async def select_tenant(user_id: UUID, tenant_id: Union[UUID, None]) -> User:
+    """
+        Set the users active tenant to provided tenant.
+
+        If None tenant_id is provided set current Tenant to the default single user-tenant
+    Args:
+        user_id: UUID of the user.
+        tenant_id: Id of the tenant.
+
+    Returns:
+        None
+
+    """
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        user = await get_user(user_id)
+        if tenant_id is None:
+            # If no tenant_id is provided set current Tenant to the single user-tenant
+            user.tenant_id = None
+            await session.merge(user)
+            await session.commit()
+            return user
+
+        tenant = await get_tenant(tenant_id)
+
+        if not user:
+            raise UserNotFoundError
+        elif not tenant:
+            raise TenantNotFoundError
+
+        # Check if User is part of Tenant
+        result = await session.execute(
+            select(UserTenant)
+            .where(UserTenant.user_id == user.id)
+            .where(UserTenant.tenant_id == tenant_id)
+        )
+
+        try:
+            result = result.scalar_one()
+        except sqlalchemy.exc.NoResultFound as e:
+            raise TenantNotFoundError("User is not part of the tenant.") from e
+
+        if result:
+            # If user is part of tenant update current tenant of user
+            user.tenant_id = tenant_id
+            await session.merge(user)
+            await session.commit()
+            return user

cognee/shared/logging_utils.py

@@ -450,6 +450,8 @@ def setup_logging(log_level=None, name=None):
             try:
                 msg = self.format(record)
                 stream = self.stream
+                if hasattr(stream, "closed") and stream.closed:
+                    return
                 stream.write("\n" + msg + self.terminator)
                 self.flush()
             except Exception:
@@ -532,6 +534,10 @@ def setup_logging(log_level=None, name=None):
     # Get a configured logger and log system information
     logger = structlog.get_logger(name if name else __name__)
 
+    logger.warning(
+        "From version 0.5.0 onwards, Cognee will run with multi-user access control mode set to on by default. Data isolation between different users and datasets will be enforced and data created before multi-user access control mode was turned on won't be accessible by default. To disable multi-user access control mode and regain access to old data set the environment variable ENABLE_BACKEND_ACCESS_CONTROL to false before starting Cognee. For more information, please refer to the Cognee documentation."
+    )
+
     if logs_dir is not None:
         logger.info(f"Log file created at: {log_file_path}", log_file=log_file_path)
 

cognee/shared/rate_limiting.py

@@ -0,0 +1,30 @@
+from aiolimiter import AsyncLimiter
+from contextlib import nullcontext
+from cognee.infrastructure.llm.config import get_llm_config
+
+llm_config = get_llm_config()
+
+llm_rate_limiter = AsyncLimiter(
+    llm_config.llm_rate_limit_requests, llm_config.embedding_rate_limit_interval
+)
+embedding_rate_limiter = AsyncLimiter(
+    llm_config.embedding_rate_limit_requests, llm_config.embedding_rate_limit_interval
+)
+
+
+def llm_rate_limiter_context_manager():
+    global llm_rate_limiter
+    if llm_config.llm_rate_limit_enabled:
+        return llm_rate_limiter
+    else:
+        #  Return a no-op context manager if rate limiting is disabled
+        return nullcontext()
+
+
+def embedding_rate_limiter_context_manager():
+    global embedding_rate_limiter
+    if llm_config.embedding_rate_limit_enabled:
+        return embedding_rate_limiter
+    else:
+        #  Return a no-op context manager if rate limiting is disabled
+        return nullcontext()

cognee/tasks/chunks/__init__.py

@@ -1,4 +1,5 @@
 from .chunk_by_word import chunk_by_word
 from .chunk_by_sentence import chunk_by_sentence
 from .chunk_by_paragraph import chunk_by_paragraph
+from .chunk_by_row import chunk_by_row
 from .remove_disconnected_chunks import remove_disconnected_chunks

cognee/tasks/chunks/chunk_by_row.py

@@ -0,0 +1,94 @@
+from typing import Any, Dict, Iterator
+from uuid import NAMESPACE_OID, uuid5
+
+from cognee.infrastructure.databases.vector.embeddings import get_embedding_engine
+
+
+def _get_pair_size(pair_text: str) -> int:
+    """
+    Calculate the size of a given text in terms of tokens.
+
+    If an embedding engine's tokenizer is available, count the tokens for the provided word.
+    If the tokenizer is not available, assume the word counts as one token.
+
+    Parameters:
+    -----------
+
+        - pair_text (str): The key:value pair text for which the token size is to be calculated.
+
+    Returns:
+    --------
+
+        - int: The number of tokens representing the text, typically an integer, depending
+          on the tokenizer's output.
+    """
+    embedding_engine = get_embedding_engine()
+    if embedding_engine.tokenizer:
+        return embedding_engine.tokenizer.count_tokens(pair_text)
+    else:
+        return 3
+
+
+def chunk_by_row(
+    data: str,
+    max_chunk_size,
+) -> Iterator[Dict[str, Any]]:
+    """
+    Chunk the input text by row while enabling exact text reconstruction.
+
+    This function divides the given text data into smaller chunks on a line-by-line basis,
+    ensuring that the size of each chunk is less than or equal to the specified maximum
+    chunk size. It guarantees that when the generated chunks are concatenated, they
+    reproduce the original text accurately. The tokenization process is handled by
+    adapters compatible with the vector engine's embedding model.
+
+    Parameters:
+    -----------
+
+        - data (str): The input text to be chunked.
+        - max_chunk_size: The maximum allowed size for each chunk, in terms of tokens or
+          words.
+    """
+    current_chunk_list = []
+    chunk_index = 0
+    current_chunk_size = 0
+
+    lines = data.split("\n\n")
+    for line in lines:
+        pairs_text = line.split(", ")
+
+        for pair_text in pairs_text:
+            pair_size = _get_pair_size(pair_text)
+            if current_chunk_size > 0 and (current_chunk_size + pair_size > max_chunk_size):
+                # Yield current cut chunk
+                current_chunk = ", ".join(current_chunk_list)
+                chunk_dict = {
+                    "text": current_chunk,
+                    "chunk_size": current_chunk_size,
+                    "chunk_id": uuid5(NAMESPACE_OID, current_chunk),
+                    "chunk_index": chunk_index,
+                    "cut_type": "row_cut",
+                }
+
+                yield chunk_dict
+
+                # Start new chunk with current pair text
+                current_chunk_list = []
+                current_chunk_size = 0
+                chunk_index += 1
+
+            current_chunk_list.append(pair_text)
+            current_chunk_size += pair_size
+
+        # Yield row chunk
+        current_chunk = ", ".join(current_chunk_list)
+        if current_chunk:
+            chunk_dict = {
+                "text": current_chunk,
+                "chunk_size": current_chunk_size,
+                "chunk_id": uuid5(NAMESPACE_OID, current_chunk),
+                "chunk_index": chunk_index,
+                "cut_type": "row_end",
+            }
+
+            yield chunk_dict

cognee/tasks/documents/__init__.py

@@ -1,3 +1,2 @@
 from .classify_documents import classify_documents
 from .extract_chunks_from_documents import extract_chunks_from_documents
-from .check_permissions_on_dataset import check_permissions_on_dataset

cognee/tasks/documents/classify_documents.py

@@ -7,6 +7,7 @@ from cognee.modules.data.processing.document_types import (
     ImageDocument,
     TextDocument,
     UnstructuredDocument,
+    CsvDocument,
 )
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.engine.utils.generate_node_id import generate_node_id
@@ -15,6 +16,7 @@ from cognee.tasks.documents.exceptions import WrongDataDocumentInputError
 EXTENSION_TO_DOCUMENT_CLASS = {
     "pdf": PdfDocument,  # Text documents
     "txt": TextDocument,
+    "csv": CsvDocument,
     "docx": UnstructuredDocument,
     "doc": UnstructuredDocument,
     "odt": UnstructuredDocument,

cognee/tasks/feedback/generate_improved_answers.py

@@ -61,7 +61,7 @@ async def _generate_improved_answer_for_single_interaction(
         )
 
         retrieved_context = await retriever.get_context(query_text)
-        completion = await retriever.get_structured_completion(
+        completion = await retriever.get_completion(
             query=query_text,
             context=retrieved_context,
             response_model=ImprovedAnswerResponse,
@@ -70,9 +70,9 @@ async def _generate_improved_answer_for_single_interaction(
         new_context_text = await retriever.resolve_edges_to_text(retrieved_context)
 
         if completion:
-            enrichment.improved_answer = completion.answer
+            enrichment.improved_answer = completion[0].answer
             enrichment.new_context = new_context_text
-            enrichment.explanation = completion.explanation
+            enrichment.explanation = completion[0].explanation
             return enrichment
         else:
             logger.warning(