|
|
|
|
@ -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
|
|
|
|
|
) -> ToolFile:
|
|
|
|
|
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,9 +77,12 @@ class ToolFileManager:
|
|
|
|
|
return tool_file
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def create_file_by_url(user_id: str, tenant_id: str,
|
|
|
|
|
conversation_id: str, file_url: str,
|
|
|
|
|
) -> ToolFile:
|
|
|
|
|
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
|
|
|
|
|
) -> ToolFile:
|
|
|
|
|
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(
|
|
|
|
|
ToolFile.id == id,
|
|
|
|
|
).first()
|
|
|
|
|
tool_file: ToolFile = (
|
|
|
|
|
db.session.query(ToolFile)
|
|
|
|
|
.filter(
|
|
|
|
|
ToolFile.id == id,
|
|
|
|
|
)
|
|
|
|
|
.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(
|
|
|
|
|
MessageFile.id == id,
|
|
|
|
|
).first()
|
|
|
|
|
|
|
|
|
|
# get tool file id
|
|
|
|
|
tool_file_id = message_file.url.split('/')[-1]
|
|
|
|
|
# trim extension
|
|
|
|
|
tool_file_id = tool_file_id.split('.')[0]
|
|
|
|
|
|
|
|
|
|
tool_file: ToolFile = db.session.query(ToolFile).filter(
|
|
|
|
|
ToolFile.id == tool_file_id,
|
|
|
|
|
).first()
|
|
|
|
|
message_file: MessageFile = (
|
|
|
|
|
db.session.query(MessageFile)
|
|
|
|
|
.filter(
|
|
|
|
|
MessageFile.id == id,
|
|
|
|
|
)
|
|
|
|
|
.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(
|
|
|
|
|
ToolFile.id == tool_file_id,
|
|
|
|
|
)
|
|
|
|
|
.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(
|
|
|
|
|
ToolFile.id == tool_file_id,
|
|
|
|
|
).first()
|
|
|
|
|
tool_file: ToolFile = (
|
|
|
|
|
db.session.query(ToolFile)
|
|
|
|
|
.filter(
|
|
|
|
|
ToolFile.id == tool_file_id,
|
|
|
|
|
)
|
|
|
|
|
.first()
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
if not tool_file:
|
|
|
|
|
return None
|
|
|
|
|
|