From 52c978faebe0db0110455115108cb4a5fd263d3d Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Mon, 29 Sep 2025 20:15:50 +0200 Subject: [PATCH] docs: Multi user authorization example (#1466) ## Description Add return value of creating role and tenant, add detailed permissions example to Cognee ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Code refactoring - [ ] Performance improvement - [ ] Other (please specify): ## Pre-submission Checklist - [x] **I have tested my changes thoroughly before submitting this PR** - [x] **This PR contains minimal changes necessary to address the issue/feature** - [x] My code follows the project's coding standards and style guidelines - [x] I have added tests that prove my fix is effective or that my feature works - [x] I have added necessary documentation (if applicable) - [x] All new and existing tests pass - [x] I have searched existing PRs to ensure this change hasn't been submitted already - [x] I have linked any relevant issues in the description - [x] My commits have clear and descriptive messages ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin. --------- Co-authored-by: Boris Co-authored-by: Hande <159312713+hande-k@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- .github/workflows/examples_tests.yml | 28 +++ .../routers/get_permissions_router.py | 12 +- .../users/roles/methods/create_role.py | 3 +- .../users/tenants/methods/create_tenant.py | 3 +- examples/python/permissions_example.py | 188 ++++++++++++++++++ 5 files changed, 228 insertions(+), 6 deletions(-) create mode 100644 examples/python/permissions_example.py diff --git a/.github/workflows/examples_tests.yml b/.github/workflows/examples_tests.yml index f4167a57a..4eb9e184f 100644 --- a/.github/workflows/examples_tests.yml +++ b/.github/workflows/examples_tests.yml @@ -1,5 +1,8 @@ name: Reusable Examples Tests +permissions: + contents: read + on: workflow_call: @@ -131,3 +134,28 @@ jobs: EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }} EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }} run: uv run python ./examples/python/memify_coding_agent_example.py + + test-permissions-example: + name: Run Permissions Example + runs-on: ubuntu-22.04 + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Cognee Setup + uses: ./.github/actions/cognee_setup + with: + python-version: '3.11.x' + + - name: Run Memify Tests + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + LLM_MODEL: ${{ secrets.LLM_MODEL }} + LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }} + LLM_API_KEY: ${{ secrets.LLM_API_KEY }} + LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }} + EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }} + EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }} + EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }} + EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }} + run: uv run python ./examples/python/permissions_example.py diff --git a/cognee/api/v1/permissions/routers/get_permissions_router.py b/cognee/api/v1/permissions/routers/get_permissions_router.py index 89603ac46..637293268 100644 --- a/cognee/api/v1/permissions/routers/get_permissions_router.py +++ b/cognee/api/v1/permissions/routers/get_permissions_router.py @@ -94,9 +94,11 @@ def get_permissions_router() -> APIRouter: from cognee.modules.users.roles.methods import create_role as create_role_method - await create_role_method(role_name=role_name, owner_id=user.id) + role_id = await create_role_method(role_name=role_name, owner_id=user.id) - return JSONResponse(status_code=200, content={"message": "Role created for tenant"}) + return JSONResponse( + status_code=200, content={"message": "Role created for tenant", "role_id": str(role_id)} + ) @permissions_router.post("/users/{user_id}/roles") async def add_user_to_role( @@ -212,8 +214,10 @@ def get_permissions_router() -> APIRouter: from cognee.modules.users.tenants.methods import create_tenant as create_tenant_method - await create_tenant_method(tenant_name=tenant_name, user_id=user.id) + tenant_id = await create_tenant_method(tenant_name=tenant_name, user_id=user.id) - return JSONResponse(status_code=200, content={"message": "Tenant created."}) + return JSONResponse( + status_code=200, content={"message": "Tenant created.", "tenant_id": str(tenant_id)} + ) return permissions_router diff --git a/cognee/modules/users/roles/methods/create_role.py b/cognee/modules/users/roles/methods/create_role.py index bdba4ad31..a57f000fe 100644 --- a/cognee/modules/users/roles/methods/create_role.py +++ b/cognee/modules/users/roles/methods/create_role.py @@ -15,7 +15,7 @@ from cognee.modules.users.models import ( async def create_role( role_name: str, owner_id: UUID, -): +) -> UUID: """ Create a new role with the given name, if the request owner with the given id has the necessary permission. @@ -45,3 +45,4 @@ async def create_role( await session.commit() await session.refresh(role) + return role.id diff --git a/cognee/modules/users/tenants/methods/create_tenant.py b/cognee/modules/users/tenants/methods/create_tenant.py index bd8abadd1..09dfc8855 100644 --- a/cognee/modules/users/tenants/methods/create_tenant.py +++ b/cognee/modules/users/tenants/methods/create_tenant.py @@ -7,7 +7,7 @@ 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): +async def create_tenant(tenant_name: str, user_id: UUID) -> UUID: """ Create a new tenant with the given name, for the user with the given id. This user is the owner of the tenant. @@ -34,5 +34,6 @@ async def create_tenant(tenant_name: str, user_id: UUID): user.tenant_id = tenant.id await session.merge(user) await session.commit() + return tenant.id except IntegrityError: raise EntityAlreadyExistsError(message="Tenant already exists.") diff --git a/examples/python/permissions_example.py b/examples/python/permissions_example.py new file mode 100644 index 000000000..4f51b660f --- /dev/null +++ b/examples/python/permissions_example.py @@ -0,0 +1,188 @@ +import os +import cognee +import pathlib + +from cognee.modules.users.exceptions import PermissionDeniedError +from cognee.shared.logging_utils import get_logger +from cognee.modules.search.types import SearchType +from cognee.modules.users.methods import create_user +from cognee.modules.users.permissions.methods import authorized_give_permission_on_datasets +from cognee.modules.users.roles.methods import add_user_to_role +from cognee.modules.users.roles.methods import create_role +from cognee.modules.users.tenants.methods import create_tenant +from cognee.modules.users.tenants.methods import add_user_to_tenant +from cognee.modules.engine.operations.setup import setup +from cognee.shared.logging_utils import setup_logging, CRITICAL + +logger = get_logger() + + +async def main(): + # ENABLE PERMISSIONS FEATURE + # Note: When ENABLE_BACKEND_ACCESS_CONTROL is enabled vector provider is automatically set to use LanceDB + # and graph provider is set to use Kuzu. + os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "True" + + # Set the rest of your environment variables as needed. By default OpenAI is used as the LLM provider + # Reference the .env.tempalte file for available option and how to change LLM provider: https://github.com/topoteretes/cognee/blob/main/.env.template + # For example to set your OpenAI LLM API key use: + # os.environ["LLM_API_KEY""] = "your-api-key" + + # Create a clean slate for cognee -- reset data and system state + print("Resetting cognee data...") + await cognee.prune.prune_data() + await cognee.prune.prune_system(metadata=True) + print("Data reset complete.\n") + + # Set up the necessary databases and tables for user management. + await setup() + + # NOTE: When a document is added in Cognee with permissions enabled only the owner of the document has permissions + # to work with the document initially. + # Add document for user_1, add it under dataset name AI + explanation_file_path = os.path.join( + pathlib.Path(__file__).parent, "../data/artificial_intelligence.pdf" + ) + + print("Creating user_1: user_1@example.com") + user_1 = await create_user("user_1@example.com", "example") + await cognee.add([explanation_file_path], dataset_name="AI", user=user_1) + + # Add document for user_2, add it under dataset name QUANTUM + text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena. + At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages + this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the + preparation and manipulation of quantum states. + """ + print("\nCreating user_2: user_2@example.com") + user_2 = await create_user("user_2@example.com", "example") + await cognee.add([text], dataset_name="QUANTUM", user=user_2) + + # Run cognify for both datasets as the appropriate user/owner + print("\nCreating different datasets for user_1 (AI dataset) and user_2 (QUANTUM dataset)") + ai_cognify_result = await cognee.cognify(["AI"], user=user_1) + quantum_cognify_result = await cognee.cognify(["QUANTUM"], user=user_2) + + # Extract dataset_ids from cognify results + def extract_dataset_id_from_cognify(cognify_result): + """Extract dataset_id from cognify output dictionary""" + for dataset_id, pipeline_result in cognify_result.items(): + return dataset_id # Return the first dataset_id + return None + + # Get dataset IDs from cognify results + # Note: When we want to work with datasets from other users (search, add, cognify and etc.) we must supply dataset + # information through dataset_id using dataset name only looks for datasets owned by current user + ai_dataset_id = extract_dataset_id_from_cognify(ai_cognify_result) + quantum_dataset_id = extract_dataset_id_from_cognify(quantum_cognify_result) + + # We can see here that user_1 can read his own dataset (AI dataset) + search_results = await cognee.search( + query_type=SearchType.GRAPH_COMPLETION, + query_text="What is in the document?", + user=user_1, + datasets=[ai_dataset_id], + ) + print("\nSearch results as user_1 on dataset owned by user_1:") + for result in search_results: + print(f"{result}\n") + + # But user_1 cant read the dataset owned by user_2 (QUANTUM dataset) + print("\nSearch result as user_1 on the dataset owned by user_2:") + try: + search_results = await cognee.search( + query_type=SearchType.GRAPH_COMPLETION, + query_text="What is in the document?", + user=user_1, + datasets=[quantum_dataset_id], + ) + except PermissionDeniedError: + print(f"User: {user_1} does not have permission to read from dataset: QUANTUM") + + # user_1 currently also can't add a document to user_2's dataset (QUANTUM dataset) + print("\nAttempting to add new data as user_1 to dataset owned by user_2:") + try: + await cognee.add( + [explanation_file_path], + dataset_id=quantum_dataset_id, + user=user_1, + ) + except PermissionDeniedError: + print(f"User: {user_1} does not have permission to write to dataset: QUANTUM") + + # We've shown that user_1 can't interact with the dataset from user_2 + # Now have user_2 give proper permission to user_1 to read QUANTUM dataset + # Note: supported permission types are "read", "write", "delete" and "share" + print( + "\nOperation started as user_2 to give read permission to user_1 for the dataset owned by user_2" + ) + await authorized_give_permission_on_datasets( + user_1.id, + [quantum_dataset_id], + "read", + user_2.id, + ) + + # Now user_1 can read from quantum dataset after proper permissions have been assigned by the QUANTUM dataset owner. + print("\nSearch result as user_1 on the dataset owned by user_2:") + search_results = await cognee.search( + query_type=SearchType.GRAPH_COMPLETION, + query_text="What is in the document?", + user=user_1, + dataset_ids=[quantum_dataset_id], + ) + for result in search_results: + print(f"{result}\n") + + # If we'd like for user_1 to add new documents to the QUANTUM dataset owned by user_2, user_1 would have to get + # "write" access permission, which user_1 currently does not have + + # Users can also be added to Roles and Tenants and then permission can be assigned on a Role/Tenant level as well + # To create a Role a user first must be an owner of a Tenant + print("User 2 is creating CogneeLab tenant/organization") + tenant_id = await create_tenant("CogneeLab", user_2.id) + + print("\nUser 2 is creating Researcher role") + role_id = await create_role(role_name="Researcher", owner_id=user_2.id) + + print("\nCreating user_3: user_3@example.com") + user_3 = await create_user("user_3@example.com", "example") + + # To add a user to a role he must be part of the same tenant/organization + print("\nOperation started as user_2 to add user_3 to CogneeLab tenant/organization") + await add_user_to_tenant(user_id=user_3.id, tenant_id=tenant_id, owner_id=user_2.id) + + print( + "\nOperation started by user_2, as tenant owner, to add user_3 to Researcher role inside the tenant/organization" + ) + await add_user_to_role(user_id=user_3.id, role_id=role_id, owner_id=user_2.id) + + print( + "\nOperation started as user_2 to give read permission to Researcher role for the dataset owned by user_2" + ) + await authorized_give_permission_on_datasets( + role_id, + [quantum_dataset_id], + "read", + user_2.id, + ) + + # Now user_3 can read from QUANTUM dataset as part of the Researcher role after proper permissions have been assigned by the QUANTUM dataset owner, user_2. + print("\nSearch result as user_3 on the dataset owned by user_2:") + search_results = await cognee.search( + query_type=SearchType.GRAPH_COMPLETION, + query_text="What is in the document?", + user=user_1, + dataset_ids=[quantum_dataset_id], + ) + for result in search_results: + print(f"{result}\n") + + # Note: All of these function calls and permission system is available through our backend endpoints as well + + +if __name__ == "__main__": + import asyncio + + logger = setup_logging(log_level=CRITICAL) + asyncio.run(main())