AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-14 19:26:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Poorandy <andymonicamua1@gmail.com>
This commit is contained in:
Poorandy 2024-07-17 22:26:18 +08:00 committed by GitHub
parent 35f4a264d6
commit 8a80af39c9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 103 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

95
api/core/tools/tool_file_manager.py
View File

@ -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}"
secret_key = current_app.config['SECRET_KEY'].encode()
data_to_sign = f'file-preview|{tool_file_id}|{timestamp}|{nonce}'
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}"
secret_key = current_app.config['SECRET_KEY'].encode()
data_to_sign = f'file-preview|{file_id}|{timestamp}|{nonce}'
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,
conversation_id: Optional[str], file_binary: bytes,
mimetype: str
def create_file_by_raw(
user_id: str, tenant_id: str, conversation_id: Optional[str], file_binary: bytes, 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,
conversation_id=conversation_id, file_key=filename, mimetype=mimetype)
tool_file = ToolFile(
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,
conversation_id: str, file_url: str,
def create_file_by_url(
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,
conversation_id=conversation_id, file_key=filename,
mimetype=mimetype, original_url=file_url)
tool_file = ToolFile(
user_id=user_id,
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,
conversation_id: str, file_key: str,
mimetype: str
def create_file_by_key(
user_id: str, tenant_id: str, conversation_id: str, file_key: str, mimetype: str
) -> ToolFile:
"""
create file
"""
tool_file = ToolFile(user_id=user_id, tenant_id=tenant_id,
conversation_id=conversation_id, file_key=file_key, mimetype=mimetype)
tool_file = ToolFile(
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

5
api/core/tools/tool_manager.py
View File

@ -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")

9
api/core/workflow/workflow_engine_manager.py
View File

@ -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_time = current_app.config.get("WORKFLOW_MAX_EXECUTION_TIME")
max_execution_steps = dify_config.WORKFLOW_MAX_EXECUTION_STEPS
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(

4
api/models/dataset.py
View File

@ -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()

7
api/models/model.py
View File

@ -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):