From 293a0e0053759686cd519e061a18e92e21ce32b7 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Wed, 17 Sep 2025 10:45:36 +0200 Subject: [PATCH] Fix formatiing --- .../embeddings/OllamaEmbeddingEngine.py | 6 +----- .../methods/get_authorized_dataset_by_name.py | 16 +++++++-------- .../layers/resolve_authorized_user_dataset.py | 20 +++++++++---------- 3 files changed, 19 insertions(+), 23 deletions(-) diff --git a/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py index 29c57ed2e..3ecc7dbe8 100644 --- a/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py @@ -94,11 +94,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine): """ Internal method to call the Ollama embeddings endpoint for a single prompt. """ - payload = { - "model": self.model, - "prompt": prompt, - "input": prompt - } + payload = {"model": self.model, "prompt": prompt, "input": prompt} headers = {} api_key = os.getenv("LLM_API_KEY") if api_key: diff --git a/cognee/modules/data/methods/get_authorized_dataset_by_name.py b/cognee/modules/data/methods/get_authorized_dataset_by_name.py index 5dc1d86a0..ad50e25e9 100644 --- a/cognee/modules/data/methods/get_authorized_dataset_by_name.py +++ b/cognee/modules/data/methods/get_authorized_dataset_by_name.py @@ -12,16 +12,16 @@ 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. + 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 + 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 - """ + 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) diff --git a/cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py b/cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py index e135b8351..7e3d1c124 100644 --- a/cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py +++ b/cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py @@ -12,18 +12,18 @@ 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. + 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. + 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. - """ + Returns: + Tuple[User, Dataset]: A tuple containing the user and the authorized dataset. + """ if not user: user = await get_default_user()