diff --git a/backend/connector/models.py b/backend/connector/models.py index cc58ee5be..bc37af041 100644 --- a/backend/connector/models.py +++ b/backend/connector/models.py @@ -1,10 +1,13 @@ +import json import uuid +from typing import Any -from account.models import User +from account.models import EncryptionSecret, User from connector.fields import ConnectorAuthJSONField from connector_auth.models import ConnectorAuth from connector_processor.connector_processor import ConnectorProcessor from connector_processor.constants import ConnectorKeys +from cryptography.fernet import Fernet from django.db import models from project.models import Project from utils.models.base_model import BaseModel @@ -107,6 +110,15 @@ def __str__(self) -> str: f" workflow: {self.workflow})" ) + @property + def metadata(self) -> Any: + encryption_secret: EncryptionSecret = EncryptionSecret.objects.get() + cipher_suite: Fernet = Fernet(encryption_secret.key.encode("utf-8")) + decrypted_value = cipher_suite.decrypt( + bytes(self.connector_metadata_b).decode("utf-8") + ) + return json.loads(decrypted_value) + class Meta: constraints = [ models.UniqueConstraint( diff --git a/backend/file_management/file_management_helper.py b/backend/file_management/file_management_helper.py index 5c834d852..3b5032101 100644 --- a/backend/file_management/file_management_helper.py +++ b/backend/file_management/file_management_helper.py @@ -36,7 +36,7 @@ class FileManagerHelper: @staticmethod def get_file_system(connector: ConnectorInstance) -> UnstractFileSystem: """Creates the `UnstractFileSystem` for the corresponding connector.""" - metadata = connector.connector_metadata + metadata = connector.metadata if connector.connector_id in fs_connectors: connector = fs_connectors[connector.connector_id]["metadata"][ "connector" diff --git a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py index 33bc189ce..361f87066 100644 --- a/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py @@ -388,7 +388,6 @@ def dynamic_indexer( vector_db = str(profile_manager.vector_store.id) x2text_adapter = str(profile_manager.x2text.id) file_hash = ToolUtils.get_hash_from_file(file_path=file_name) - extract_file_path = None if not is_summary: directory, filename = os.path.split(file_name) extract_file_path: str = os.path.join( diff --git a/backend/tool_instance/tool_instance_helper.py b/backend/tool_instance/tool_instance_helper.py index febbdc33a..102460276 100644 --- a/backend/tool_instance/tool_instance_helper.py +++ b/backend/tool_instance/tool_instance_helper.py @@ -59,12 +59,9 @@ def update_instance_metadata( output_connector = ConnectorInstanceHelper.get_output_connector_instance_by_name_for_workflow( # noqa tool_instance.workflow_id, output_connector_name ) - if ( - output_connector - and "path" in output_connector.connector_metadata - ): + if output_connector and "path" in output_connector.metadata: metadata[JsonSchemaKey.OUTPUT_FOLDER] = os.path.join( - output_connector.connector_metadata["path"], + output_connector.metadata["path"], *(metadata[JsonSchemaKey.OUTPUT_FOLDER].split("/")), ) if ( @@ -76,9 +73,9 @@ def update_instance_metadata( tool_instance.workflow_id, input_connector_name ) - if input_connector and "path" in input_connector.connector_metadata: + if input_connector and "path" in input_connector.metadata: metadata[JsonSchemaKey.ROOT_FOLDER] = os.path.join( - input_connector.connector_metadata["path"], + input_connector.metadata["path"], *(metadata[JsonSchemaKey.ROOT_FOLDER].split("/")), ) ToolInstanceHelper.update_metadata_with_adapter_instances( @@ -181,6 +178,7 @@ def update_metadata_with_adapter_instances( adapter_property=adapter_property, adapter_type=AdapterTypes.OCR, ) + @staticmethod def get_altered_metadata( tool_instance: ToolInstance, @@ -210,13 +208,10 @@ def get_altered_metadata( output_connector = ConnectorInstanceHelper.get_output_connector_instance_by_name_for_workflow( # noqa tool_instance.workflow_id, output_connector_name ) - if ( - output_connector - and "path" in output_connector.connector_metadata - ): + if output_connector and "path" in output_connector.metadata: relative_path = ToolInstanceHelper.get_relative_path( metadata[JsonSchemaKey.OUTPUT_FOLDER], - output_connector.connector_metadata["path"], + output_connector.metadata["path"], ) metadata[JsonSchemaKey.OUTPUT_FOLDER] = relative_path if ( @@ -227,10 +222,10 @@ def get_altered_metadata( input_connector = ConnectorInstanceHelper.get_input_connector_instance_by_name_for_workflow( # noqa tool_instance.workflow_id, input_connector_name ) - if input_connector and "path" in input_connector.connector_metadata: + if input_connector and "path" in input_connector.metadata: relative_path = ToolInstanceHelper.get_relative_path( metadata[JsonSchemaKey.ROOT_FOLDER], - input_connector.connector_metadata["path"], + input_connector.metadata["path"], ) metadata[JsonSchemaKey.ROOT_FOLDER] = relative_path return metadata diff --git a/backend/workflow_manager/endpoint/destination.py b/backend/workflow_manager/endpoint/destination.py index 2fb16eb03..cdaae988b 100644 --- a/backend/workflow_manager/endpoint/destination.py +++ b/backend/workflow_manager/endpoint/destination.py @@ -6,9 +6,7 @@ import fsspec import magic -from account.models import EncryptionSecret from connector.models import ConnectorInstance -from cryptography.fernet import Fernet from django.db import connection from fsspec.implementations.local import LocalFileSystem from unstract.sdk.constants import ToolExecKey @@ -78,15 +76,8 @@ def _get_endpoint_for_workflow( endpoint_type=WorkflowEndpoint.EndpointType.DESTINATION, ) if endpoint.connector_instance: - encryption_secret: EncryptionSecret = EncryptionSecret.objects.get() - f: Fernet = Fernet(encryption_secret.key.encode("utf-8")) - endpoint.connector_instance.connector_metadata = json.loads( - f.decrypt( - bytes(endpoint.connector_instance.connector_metadata_b - ).decode( - "utf-8" - ) - ) + endpoint.connector_instance.connector_metadata = ( + endpoint.connector_instance.metadata ) return endpoint @@ -191,9 +182,7 @@ def copy_output_to_output_directory(self) -> None: def insert_into_db(self, file_history: Optional[FileHistory]) -> None: """Insert data into the database.""" connector_instance: ConnectorInstance = self.endpoint.connector_instance - connector_settings: dict[ - str, Any - ] = connector_instance.connector_metadata + connector_settings: dict[str, Any] = connector_instance.metadata destination_configurations: dict[str, Any] = self.endpoint.configuration table_name: str = str( destination_configurations.get(DestinationKey.TABLE) diff --git a/backend/workflow_manager/endpoint/source.py b/backend/workflow_manager/endpoint/source.py index 7e9e3fd03..6dfe41e61 100644 --- a/backend/workflow_manager/endpoint/source.py +++ b/backend/workflow_manager/endpoint/source.py @@ -1,5 +1,4 @@ import fnmatch -import json import logging import os import shutil @@ -8,10 +7,8 @@ from typing import Any, Optional import fsspec -from account.models import EncryptionSecret from connector.models import ConnectorInstance from connector_processor.constants import ConnectorKeys -from cryptography.fernet import Fernet from django.core.files.uploadedfile import UploadedFile from django.db import connection from unstract.workflow_execution.enums import LogState @@ -93,17 +90,9 @@ def _get_endpoint_for_workflow( endpoint_type=WorkflowEndpoint.EndpointType.SOURCE, ) if endpoint.connector_instance: - encryption_secret: EncryptionSecret = EncryptionSecret.objects.get() - f: Fernet = Fernet(encryption_secret.key.encode("utf-8")) - endpoint.connector_instance.connector_metadata = json.loads( - f.decrypt( - bytes(endpoint.connector_instance.connector_metadata_b - ).decode( - "utf-8" - ) - ) + endpoint.connector_instance.connector_metadata = ( + endpoint.connector_instance.metadata ) - return endpoint def validate(self) -> None: