refactor(models&tools): switch to dify_config in models and tools. (#6394)

Co-authored-by: Poorandy <andymonicamua1@gmail.com>

api/core/tools/tool_file_manager.py

@@ -9,9 +9,9 @@ from mimetypes import guess_extension, guess_type
 from typing import Optional, Union
 from uuid import uuid4
 
-from flask import current_app
 from httpx import get
 
+from configs import dify_config
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models.model import MessageFile
@@ -26,25 +26,25 @@ class ToolFileManager:
         """
         sign file to get a temporary url
         """
-        base_url = current_app.config.get('FILES_URL')
+        base_url = dify_config.FILES_URL
         file_preview_url = f'{base_url}/files/tools/{tool_file_id}{extension}'
 
         timestamp = str(int(time.time()))
         nonce = os.urandom(16).hex()
-        data_to_sign = f"file-preview|{tool_file_id}|{timestamp}|{nonce}"
+        data_to_sign = f'file-preview|{tool_file_id}|{timestamp}|{nonce}'
-        secret_key = current_app.config['SECRET_KEY'].encode()
+        secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b''
         sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
         encoded_sign = base64.urlsafe_b64encode(sign).decode()
 
-        return f"{file_preview_url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
+        return f'{file_preview_url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}'
 
     @staticmethod
     def verify_file(file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
         """
         verify signature
         """
-        data_to_sign = f"file-preview|{file_id}|{timestamp}|{nonce}"
+        data_to_sign = f'file-preview|{file_id}|{timestamp}|{nonce}'
-        secret_key = current_app.config['SECRET_KEY'].encode()
+        secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b''
         recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
         recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()
 
@@ -53,23 +53,23 @@ class ToolFileManager:
             return False
 
         current_time = int(time.time())
-        return current_time - int(timestamp) <= current_app.config.get('FILES_ACCESS_TIMEOUT')
+        return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT
 
     @staticmethod
-    def create_file_by_raw(user_id: str, tenant_id: str,
+    def create_file_by_raw(
-                           conversation_id: Optional[str], file_binary: bytes,
+        user_id: str, tenant_id: str, conversation_id: Optional[str], file_binary: bytes, mimetype: str
-                           mimetype: str
     ) -> ToolFile:
         """
         create file
         """
         extension = guess_extension(mimetype) or '.bin'
         unique_name = uuid4().hex
-        filename = f"tools/{tenant_id}/{unique_name}{extension}"
+        filename = f'tools/{tenant_id}/{unique_name}{extension}'
         storage.save(filename, file_binary)
 
-        tool_file = ToolFile(user_id=user_id, tenant_id=tenant_id,
+        tool_file = ToolFile(
-                             conversation_id=conversation_id, file_key=filename, mimetype=mimetype)
+            user_id=user_id, tenant_id=tenant_id, conversation_id=conversation_id, file_key=filename, mimetype=mimetype
+        )
 
         db.session.add(tool_file)
         db.session.commit()
@@ -77,8 +77,11 @@ class ToolFileManager:
         return tool_file
 
     @staticmethod
-    def create_file_by_url(user_id: str, tenant_id: str,
+    def create_file_by_url(
-                           conversation_id: str, file_url: str,
+        user_id: str,
+        tenant_id: str,
+        conversation_id: str,
+        file_url: str,
     ) -> ToolFile:
         """
         create file
@@ -90,12 +93,17 @@ class ToolFileManager:
         mimetype = guess_type(file_url)[0] or 'octet/stream'
         extension = guess_extension(mimetype) or '.bin'
         unique_name = uuid4().hex
-        filename = f"tools/{tenant_id}/{unique_name}{extension}"
+        filename = f'tools/{tenant_id}/{unique_name}{extension}'
         storage.save(filename, blob)
 
-        tool_file = ToolFile(user_id=user_id, tenant_id=tenant_id,
+        tool_file = ToolFile(
-                             conversation_id=conversation_id, file_key=filename,
+            user_id=user_id,
-                             mimetype=mimetype, original_url=file_url)
+            tenant_id=tenant_id,
+            conversation_id=conversation_id,
+            file_key=filename,
+            mimetype=mimetype,
+            original_url=file_url,
+        )
 
         db.session.add(tool_file)
         db.session.commit()
@@ -103,15 +111,15 @@ class ToolFileManager:
         return tool_file
 
     @staticmethod
-    def create_file_by_key(user_id: str, tenant_id: str,
+    def create_file_by_key(
-                           conversation_id: str, file_key: str,
+        user_id: str, tenant_id: str, conversation_id: str, file_key: str, mimetype: str
-                           mimetype: str
     ) -> ToolFile:
         """
         create file
         """
-        tool_file = ToolFile(user_id=user_id, tenant_id=tenant_id,
+        tool_file = ToolFile(
-                             conversation_id=conversation_id, file_key=file_key, mimetype=mimetype)
+            user_id=user_id, tenant_id=tenant_id, conversation_id=conversation_id, file_key=file_key, mimetype=mimetype
+        )
         return tool_file
 
     @staticmethod
@@ -123,9 +131,13 @@ class ToolFileManager:
 
         :return: the binary of the file, mime type
         """
-        tool_file: ToolFile = db.session.query(ToolFile).filter(
+        tool_file: ToolFile = (
+            db.session.query(ToolFile)
+            .filter(
                 ToolFile.id == id,
-        ).first()
+            )
+            .first()
+        )
 
         if not tool_file:
             return None
@@ -143,18 +155,31 @@ class ToolFileManager:
 
         :return: the binary of the file, mime type
         """
-        message_file: MessageFile = db.session.query(MessageFile).filter(
+        message_file: MessageFile = (
+            db.session.query(MessageFile)
+            .filter(
                 MessageFile.id == id,
-        ).first()
+            )
+            .first()
+        )
 
+        # Check if message_file is not None
+        if message_file is not None:
             # get tool file id
             tool_file_id = message_file.url.split('/')[-1]
             # trim extension
             tool_file_id = tool_file_id.split('.')[0]
+        else:
+            tool_file_id = None
 
-        tool_file: ToolFile = db.session.query(ToolFile).filter(
+
+        tool_file: ToolFile = (
+            db.session.query(ToolFile)
+            .filter(
                 ToolFile.id == tool_file_id,
-        ).first()
+            )
+            .first()
+        )
 
         if not tool_file:
             return None
@@ -172,9 +197,13 @@ class ToolFileManager:
 
         :return: the binary of the file, mime type
         """
-        tool_file: ToolFile = db.session.query(ToolFile).filter(
+        tool_file: ToolFile = (
+            db.session.query(ToolFile)
+            .filter(
                 ToolFile.id == tool_file_id,
-        ).first()
+            )
+            .first()
+        )
 
         if not tool_file:
             return None

api/core/tools/tool_manager.py

@@ -6,8 +6,7 @@ from os import listdir, path
 from threading import Lock
 from typing import Any, Union
 
-from flask import current_app
+from configs import dify_config
-
 from core.agent.entities import AgentToolEntity
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.helper.module_import_helper import load_single_subclass_from_source
@@ -566,7 +565,7 @@ class ToolManager:
         provider_type = provider_type
         provider_id = provider_id
         if provider_type == 'builtin':
-            return (current_app.config.get("CONSOLE_API_URL")
+            return (dify_config.CONSOLE_API_URL
                     + "/console/api/workspaces/current/tool-provider/builtin/"
                     + provider_id
                     + "/icon")

api/core/workflow/workflow_engine_manager.py

@@ -2,8 +2,7 @@ import logging
 import time
 from typing import Optional, cast
 
-from flask import current_app
+from configs import dify_config
-
 from core.app.app_config.entities import FileExtraConfig
 from core.app.apps.base_app_queue_manager import GenerateTaskStoppedException
 from core.app.entities.app_invoke_entities import InvokeFrom
@@ -126,7 +125,7 @@ class WorkflowEngineManager:
                 user_inputs=user_inputs
             )
 
-        workflow_call_max_depth = current_app.config.get("WORKFLOW_CALL_MAX_DEPTH")
+        workflow_call_max_depth = dify_config.WORKFLOW_CALL_MAX_DEPTH
         if call_depth > workflow_call_max_depth:
             raise ValueError('Max workflow call depth {} reached.'.format(workflow_call_max_depth))
 
@@ -177,8 +176,8 @@ class WorkflowEngineManager:
             predecessor_node: BaseNode = None
             current_iteration_node: BaseIterationNode = None
             has_entry_node = False
-            max_execution_steps = current_app.config.get("WORKFLOW_MAX_EXECUTION_STEPS")
+            max_execution_steps = dify_config.WORKFLOW_MAX_EXECUTION_STEPS
-            max_execution_time = current_app.config.get("WORKFLOW_MAX_EXECUTION_TIME")
+            max_execution_time = dify_config.WORKFLOW_MAX_EXECUTION_TIME
             while True:
                 # get next node, multiple target nodes in the future
                 next_node = self._get_next_overall_node(

api/models/dataset.py

@@ -9,10 +9,10 @@ import re
 import time
 from json import JSONDecodeError
 
-from flask import current_app
 from sqlalchemy import func
 from sqlalchemy.dialects.postgresql import JSONB
 
+from configs import dify_config
 from core.rag.retrieval.retrival_methods import RetrievalMethod
 from extensions.ext_database import db
 from extensions.ext_storage import storage
@@ -528,7 +528,7 @@ class DocumentSegment(db.Model):
             nonce = os.urandom(16).hex()
             timestamp = str(int(time.time()))
             data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
-            secret_key = current_app.config['SECRET_KEY'].encode()
+            secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b''
             sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
             encoded_sign = base64.urlsafe_b64encode(sign).decode()
 

api/models/model.py

@@ -4,10 +4,11 @@ import uuid
 from enum import Enum
 from typing import Optional
 
-from flask import current_app, request
+from flask import request
 from flask_login import UserMixin
 from sqlalchemy import Float, func, text
 
+from configs import dify_config
 from core.file.tool_file_parser import ToolFileParser
 from core.file.upload_file_parser import UploadFileParser
 from extensions.ext_database import db
@@ -111,7 +112,7 @@ class App(db.Model):
 
     @property
     def api_base_url(self):
-        return (current_app.config['SERVICE_API_URL'] if current_app.config['SERVICE_API_URL']
+        return (dify_config.SERVICE_API_URL if dify_config.SERVICE_API_URL
                 else request.host_url.rstrip('/')) + '/v1'
 
     @property
@@ -1113,7 +1114,7 @@ class Site(db.Model):
     @property
     def app_base_url(self):
         return (
-            current_app.config['APP_WEB_URL'] if current_app.config['APP_WEB_URL'] else request.host_url.rstrip('/'))
+            dify_config.APP_WEB_URL if  dify_config.APP_WEB_URL else request.host_url.rstrip('/'))
 
 
 class ApiToken(db.Model):