refactor: Make the authorized_user_datasets function more understandable
This commit is contained in:
parent
d4b23aa565
commit
9a4e8ddc4d
2 changed files with 18 additions and 3 deletions
|
|
@ -1,8 +1,9 @@
|
|||
from uuid import UUID
|
||||
from typing import Union
|
||||
from typing import Union, Tuple, List
|
||||
|
||||
from cognee.modules.users.methods import get_default_user
|
||||
from cognee.modules.users.models import User
|
||||
from cognee.modules.data.models import Dataset
|
||||
from cognee.modules.data.exceptions import DatasetNotFoundError
|
||||
from cognee.modules.data.methods import (
|
||||
get_authorized_existing_datasets,
|
||||
|
|
@ -11,7 +12,21 @@ from cognee.modules.data.methods import (
|
|||
)
|
||||
|
||||
|
||||
async def authorized_user_datasets(user: User, datasets: Union[str, list[str], list[UUID]]):
|
||||
async def authorized_user_datasets(
|
||||
datasets: Union[str, list[str], list[UUID]], user: User = None
|
||||
) -> Tuple[User, List[Dataset]]:
|
||||
"""
|
||||
Function handles creation and dataset authorization if datasets 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:
|
||||
user: Cognee User request is being processed for, if None default user will be used.
|
||||
datasets: Dataset names or Dataset UUID (in case Datasets already exist)
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
# If no user is provided use default user
|
||||
if user is None:
|
||||
user = await get_default_user()
|
||||
|
|
|
|||
|
|
@ -67,7 +67,7 @@ async def cognee_pipeline(
|
|||
await test_embedding_connection()
|
||||
cognee_pipeline.first_run = False # Update flag after first run
|
||||
|
||||
user, authorized_datasets = await authorized_user_datasets(user, datasets)
|
||||
user, authorized_datasets = await authorized_user_datasets(datasets, user)
|
||||
|
||||
for dataset in authorized_datasets:
|
||||
async for run_info in run_pipeline(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue