add fixes

This commit is contained in:
Hande 2025-12-02 18:06:28 +01:00
parent 023bab470a
commit 04aae3852d
7 changed files with 51 additions and 12 deletions

View file

@ -3,6 +3,7 @@ import cognee
import pathlib
from cognee.modules.users.exceptions import PermissionDeniedError
from cognee.modules.users.tenants.methods import select_tenant
from cognee.shared.logging_utils import get_logger
from cognee.modules.search.types import SearchType
from cognee.modules.users.methods import create_user
@ -26,7 +27,7 @@ async def main():
# 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"
# os.environ["LLM_API_KEY"] = "your-api-key"
# Create a clean slate for cognee -- reset data and system state
print("Resetting cognee data...")
@ -142,6 +143,9 @@ async def main():
print("User 2 is creating CogneeLab tenant/organization")
tenant_id = await create_tenant("CogneeLab", user_2.id)
print("User 2 is selecting CogneeLab tenant/organization as active tenant")
await select_tenant(user_id=user_2.id, tenant_id=tenant_id)
print("\nUser 2 is creating Researcher role")
role_id = await create_role(role_name="Researcher", owner_id=user_2.id)
@ -157,23 +161,59 @@ async def main():
)
await add_user_to_role(user_id=user_3.id, role_id=role_id, owner_id=user_2.id)
print("\nOperation as user_3 to select CogneeLab tenant/organization as active tenant")
await select_tenant(user_id=user_3.id, tenant_id=tenant_id)
print(
"\nOperation started as user_2 to give read permission to Researcher role for the dataset owned by user_2"
"\nOperation started as user_2, with CogneeLab as its active tenant, to give read permission to Researcher role for the dataset QUANTUM owned by user_2"
)
# Even though the dataset owner is user_2, the dataset doesn't belong to the tenant/organization CogneeLab.
# So we can't assign permissions to it when we're acting in the CogneeLab tenant.
try:
await authorized_give_permission_on_datasets(
role_id,
[quantum_dataset_id],
"read",
user_2.id,
)
except PermissionDeniedError:
print(
"User 2 could not give permission to the role as the QUANTUM dataset is not part of the CogneeLab tenant"
)
print(
"We will now create a new QUANTUM dataset with the QUANTUM_COGNEE_LAB name in the CogneeLab tenant so that permissions can be assigned to the Researcher role inside the tenant/organization"
)
# We can re-create the QUANTUM dataset in the CogneeLab tenant. The old QUANTUM dataset is still owned by user_2 personally
# and can still be accessed by selecting the personal tenant for user 2.
from cognee.modules.users.methods import get_user
# Note: We need to update user_2 from the database to refresh its tenant context changes
user_2 = await get_user(user_2.id)
await cognee.add([text], dataset_name="QUANTUM_COGNEE_LAB", user=user_2)
quantum_cognee_lab_cognify_result = await cognee.cognify(["QUANTUM_COGNEE_LAB"], user=user_2)
# The recreated Quantum dataset will now have a different dataset_id as it's a new dataset in a different organization
quantum_cognee_lab_dataset_id = extract_dataset_id_from_cognify(
quantum_cognee_lab_cognify_result
)
print(
"\nOperation started as user_2, with CogneeLab as its active tenant, to give read permission to Researcher role for the dataset QUANTUM owned by the CogneeLab tenant"
)
await authorized_give_permission_on_datasets(
role_id,
[quantum_dataset_id],
[quantum_cognee_lab_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:")
print("\nSearch result as user_3 on the QUANTUM dataset owned by the CogneeLab organization:")
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],
user=user_3,
dataset_ids=[quantum_cognee_lab_dataset_id],
)
for result in search_results:
print(f"{result}\n")

View file

@ -49,7 +49,7 @@ class Company(DataPoint):
ROOT = Path(__file__).resolve().parent
DATA_DIR = ROOT.parent / "data"
DATA_DIR = ROOT / "data"
COGNEE_DIR = ROOT / ".cognee_system"
ARTIFACTS_DIR = ROOT / ".artifacts"
GRAPH_HTML = ARTIFACTS_DIR / "graph_visualization.html"

View file

@ -3,4 +3,4 @@ Core Features Getting Started Example
Reference: https://colab.research.google.com/drive/12Vi9zID-M3fpKpKiaqDBvkk98ElkRPWy?usp=sharing
"""
"""

View file

@ -4,4 +4,4 @@ Custom Prompt Example
Reference: https://docs.cognee.ai/guides/custom-prompts
"""
"""

View file

@ -1,4 +1,4 @@
"""
Direct LLM Call for Structured Output Example
Reference: https://docs.cognee.ai/guides/low-level-llm
"""
"""

View file

@ -2,4 +2,4 @@
Graph Visualization Example
Reference: https://docs.cognee.ai/guides/graph-visualization
"""
"""

View file

@ -3,4 +3,3 @@ Retrievers and Search Examples
Reference: https://docs.cognee.ai/guides/search-basics
"""