mirror of
https://github.com/langgenius/dify.git
synced 2026-02-22 14:01:50 -05:00
493 lines
18 KiB
Python
493 lines
18 KiB
Python
import logging
|
|
import threading
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from core.app.entities.app_asset_entities import (
|
|
AppAssetFileTree,
|
|
AppAssetNode,
|
|
AssetNodeType,
|
|
BatchUploadNode,
|
|
TreeNodeNotFoundError,
|
|
TreeParentNotFoundError,
|
|
TreePathConflictError,
|
|
)
|
|
from core.app_assets.builder import AssetBuildPipeline, BuildContext
|
|
from core.app_assets.builder.file_builder import FileBuilder
|
|
from core.app_assets.builder.skill_builder import SkillBuilder
|
|
from core.app_assets.converters import tree_to_asset_items
|
|
from core.app_assets.entities.assets import AssetItem
|
|
from core.app_assets.packager import AssetZipPackager
|
|
from core.app_assets.storage import AssetPath, app_asset_storage
|
|
from extensions.ext_database import db
|
|
from extensions.ext_redis import redis_client
|
|
from models.app_asset import AppAssets
|
|
from models.model import App
|
|
|
|
from .errors.app_asset import (
|
|
AppAssetNodeNotFoundError,
|
|
AppAssetNodeTooLargeError,
|
|
AppAssetParentNotFoundError,
|
|
AppAssetPathConflictError,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AppAssetService:
|
|
MAX_PREVIEW_CONTENT_SIZE = 5 * 1024 * 1024 # 5MB
|
|
_LOCK_TIMEOUT_SECONDS = 60
|
|
_DRAFT_CACHE_KEY_PREFIX = "app_asset:draft_download"
|
|
|
|
@staticmethod
|
|
def _lock(app_id: str):
|
|
return redis_client.lock(f"app_asset:lock:{app_id}", timeout=AppAssetService._LOCK_TIMEOUT_SECONDS)
|
|
|
|
@staticmethod
|
|
def get_or_create_assets(session: Session, app_model: App, account_id: str) -> AppAssets:
|
|
assets = (
|
|
session.query(AppAssets)
|
|
.filter(
|
|
AppAssets.tenant_id == app_model.tenant_id,
|
|
AppAssets.app_id == app_model.id,
|
|
AppAssets.version == AppAssets.VERSION_DRAFT,
|
|
)
|
|
.first()
|
|
)
|
|
if not assets:
|
|
assets = AppAssets(
|
|
id=str(uuid4()),
|
|
tenant_id=app_model.tenant_id,
|
|
app_id=app_model.id,
|
|
version=AppAssets.VERSION_DRAFT,
|
|
created_by=account_id,
|
|
)
|
|
session.add(assets)
|
|
session.commit()
|
|
return assets
|
|
|
|
@staticmethod
|
|
def get_tenant_app_assets(tenant_id: str, assets_id: str) -> AppAssets:
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
app_assets = (
|
|
session.query(AppAssets)
|
|
.filter(
|
|
AppAssets.tenant_id == tenant_id,
|
|
AppAssets.id == assets_id,
|
|
)
|
|
.first()
|
|
)
|
|
if not app_assets:
|
|
raise ValueError(f"App assets not found for tenant_id={tenant_id}, assets_id={assets_id}")
|
|
|
|
return app_assets
|
|
|
|
@staticmethod
|
|
def get_assets(tenant_id: str, app_id: str, user_id: str, *, is_draft: bool) -> AppAssets | None:
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
if is_draft:
|
|
stmt = session.query(AppAssets).filter(
|
|
AppAssets.tenant_id == tenant_id,
|
|
AppAssets.app_id == app_id,
|
|
AppAssets.version == AppAssets.VERSION_DRAFT,
|
|
)
|
|
if not stmt.first():
|
|
assets = AppAssets(
|
|
id=str(uuid4()),
|
|
tenant_id=tenant_id,
|
|
app_id=app_id,
|
|
version=AppAssets.VERSION_DRAFT,
|
|
created_by=user_id,
|
|
)
|
|
session.add(assets)
|
|
session.commit()
|
|
else:
|
|
stmt = (
|
|
session.query(AppAssets)
|
|
.filter(
|
|
AppAssets.tenant_id == tenant_id,
|
|
AppAssets.app_id == app_id,
|
|
AppAssets.version != AppAssets.VERSION_DRAFT,
|
|
)
|
|
.order_by(AppAssets.created_at.desc())
|
|
)
|
|
return stmt.first()
|
|
|
|
@staticmethod
|
|
def get_asset_tree(app_model: App, account_id: str) -> AppAssetFileTree:
|
|
with Session(db.engine) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
return assets.asset_tree
|
|
|
|
@staticmethod
|
|
def create_folder(
|
|
app_model: App,
|
|
account_id: str,
|
|
name: str,
|
|
parent_id: str | None = None,
|
|
) -> AppAssetNode:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
node = AppAssetNode.create_folder(str(uuid4()), name, parent_id)
|
|
|
|
try:
|
|
tree.add(node)
|
|
except TreeParentNotFoundError as e:
|
|
raise AppAssetParentNotFoundError(str(e)) from e
|
|
except TreePathConflictError as e:
|
|
raise AppAssetPathConflictError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def get_file_content(app_model: App, account_id: str, node_id: str) -> bytes:
|
|
with Session(db.engine) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
node = tree.get(node_id)
|
|
if not node or node.node_type != AssetNodeType.FILE:
|
|
raise AppAssetNodeNotFoundError(f"File node {node_id} not found")
|
|
|
|
if node.size > AppAssetService.MAX_PREVIEW_CONTENT_SIZE:
|
|
max_size_mb = AppAssetService.MAX_PREVIEW_CONTENT_SIZE / 1024 / 1024
|
|
raise AppAssetNodeTooLargeError(f"File node {node_id} size exceeded the limit: {max_size_mb} MB")
|
|
|
|
asset_storage = app_asset_storage
|
|
asset_path = AssetPath.draft(app_model.tenant_id, app_model.id, node_id)
|
|
return asset_storage.load(asset_path)
|
|
|
|
@staticmethod
|
|
def update_file_content(
|
|
app_model: App,
|
|
account_id: str,
|
|
node_id: str,
|
|
content: bytes,
|
|
) -> AppAssetNode:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
try:
|
|
node = tree.update(node_id, len(content))
|
|
except TreeNodeNotFoundError as e:
|
|
raise AppAssetNodeNotFoundError(str(e)) from e
|
|
|
|
asset_storage = app_asset_storage
|
|
asset_path = AssetPath.draft(app_model.tenant_id, app_model.id, node_id)
|
|
asset_storage.save(asset_path, content)
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def rename_node(
|
|
app_model: App,
|
|
account_id: str,
|
|
node_id: str,
|
|
new_name: str,
|
|
) -> AppAssetNode:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
try:
|
|
node = tree.rename(node_id, new_name)
|
|
except TreeNodeNotFoundError as e:
|
|
raise AppAssetNodeNotFoundError(str(e)) from e
|
|
except TreePathConflictError as e:
|
|
raise AppAssetPathConflictError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
return node
|
|
|
|
@staticmethod
|
|
def move_node(
|
|
app_model: App,
|
|
account_id: str,
|
|
node_id: str,
|
|
new_parent_id: str | None,
|
|
) -> AppAssetNode:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
try:
|
|
node = tree.move(node_id, new_parent_id)
|
|
except TreeNodeNotFoundError as e:
|
|
raise AppAssetNodeNotFoundError(str(e)) from e
|
|
except TreeParentNotFoundError as e:
|
|
raise AppAssetParentNotFoundError(str(e)) from e
|
|
except TreePathConflictError as e:
|
|
raise AppAssetPathConflictError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def reorder_node(
|
|
app_model: App,
|
|
account_id: str,
|
|
node_id: str,
|
|
after_node_id: str | None,
|
|
) -> AppAssetNode:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id=account_id)
|
|
tree = assets.asset_tree
|
|
|
|
try:
|
|
node = tree.reorder(node_id, after_node_id)
|
|
except TreeNodeNotFoundError as e:
|
|
raise AppAssetNodeNotFoundError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def delete_node(app_model: App, account_id: str, node_id: str) -> None:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
try:
|
|
removed_ids = tree.remove(node_id)
|
|
except TreeNodeNotFoundError as e:
|
|
raise AppAssetNodeNotFoundError(str(e)) from e
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
# FIXME(Mairuis): sync deletion queue, failed is fine
|
|
def _delete_file_from_storage(tenant_id: str, app_id: str, node_ids: list[str]) -> None:
|
|
asset_storage = app_asset_storage
|
|
for nid in node_ids:
|
|
asset_path = AssetPath.draft(tenant_id, app_id, nid)
|
|
try:
|
|
asset_storage.delete(asset_path)
|
|
except Exception:
|
|
logger.warning(
|
|
"Failed to delete storage file %s",
|
|
asset_storage.get_storage_key(asset_path),
|
|
exc_info=True,
|
|
)
|
|
|
|
threading.Thread(
|
|
target=lambda: _delete_file_from_storage(app_model.tenant_id, app_model.id, removed_ids)
|
|
).start()
|
|
|
|
@staticmethod
|
|
def publish(session: Session, app_model: App, account_id: str, workflow_id: str) -> AppAssets:
|
|
tenant_id = app_model.tenant_id
|
|
app_id = app_model.id
|
|
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
publish_id = str(uuid4())
|
|
|
|
published = AppAssets(
|
|
id=publish_id,
|
|
tenant_id=tenant_id,
|
|
app_id=app_id,
|
|
version=workflow_id,
|
|
created_by=account_id,
|
|
)
|
|
published.asset_tree = tree
|
|
session.add(published)
|
|
session.flush()
|
|
|
|
asset_storage = app_asset_storage
|
|
ctx = BuildContext(tenant_id=tenant_id, app_id=app_id, build_id=publish_id)
|
|
built_assets = AssetBuildPipeline(
|
|
[SkillBuilder(storage=asset_storage), FileBuilder(storage=asset_storage)]
|
|
).build_all(tree, ctx)
|
|
|
|
packager = AssetZipPackager(asset_storage.storage)
|
|
|
|
runtime_zip_bytes = packager.package(built_assets)
|
|
runtime_zip_path = AssetPath.build_zip(tenant_id, app_id, publish_id)
|
|
asset_storage.save(runtime_zip_path, runtime_zip_bytes)
|
|
|
|
source_items = tree_to_asset_items(tree, tenant_id, app_id, asset_storage)
|
|
source_zip_bytes = packager.package(source_items)
|
|
source_zip_path = AssetPath.source_zip(tenant_id, app_id, workflow_id)
|
|
asset_storage.save(source_zip_path, source_zip_bytes)
|
|
|
|
return published
|
|
|
|
@staticmethod
|
|
def build_assets(tenant_id: str, app_id: str, assets: AppAssets) -> None:
|
|
# Build resolved draft assets without packaging into a zip.
|
|
tree = assets.asset_tree
|
|
|
|
asset_storage = app_asset_storage
|
|
ctx = BuildContext(tenant_id=tenant_id, app_id=app_id, build_id=assets.id)
|
|
built_assets: list[AssetItem] = AssetBuildPipeline(
|
|
[SkillBuilder(storage=asset_storage), FileBuilder(storage=asset_storage)]
|
|
).build_all(tree, ctx)
|
|
|
|
packager = AssetZipPackager(storage=asset_storage.storage)
|
|
zip_bytes = packager.package(built_assets)
|
|
zip_path = AssetPath.build_zip(tenant_id, app_id, assets.id)
|
|
asset_storage.save(zip_path, zip_bytes)
|
|
|
|
@staticmethod
|
|
def get_file_download_url(
|
|
app_model: App,
|
|
account_id: str,
|
|
node_id: str,
|
|
expires_in: int = 3600,
|
|
) -> str:
|
|
with Session(db.engine) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
node = tree.get(node_id)
|
|
if not node or node.node_type != AssetNodeType.FILE:
|
|
raise AppAssetNodeNotFoundError(f"File node {node_id} not found")
|
|
|
|
asset_storage = app_asset_storage
|
|
asset_path = AssetPath.draft(app_model.tenant_id, app_model.id, node_id)
|
|
return asset_storage.get_download_url(asset_path, expires_in)
|
|
|
|
@staticmethod
|
|
def get_source_zip_bytes(tenant_id: str, app_id: str, workflow_id: str) -> bytes | None:
|
|
asset_storage = app_asset_storage
|
|
asset_path = AssetPath.source_zip(tenant_id, app_id, workflow_id)
|
|
source_zip = asset_storage.load_or_none(asset_path)
|
|
if source_zip is None:
|
|
logger.warning("Source zip not found: %s", asset_storage.get_storage_key(asset_path))
|
|
return source_zip
|
|
|
|
@staticmethod
|
|
def set_draft_assets(
|
|
app_model: App,
|
|
account_id: str,
|
|
new_tree: AppAssetFileTree,
|
|
) -> AppAssets:
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
assets.asset_tree = new_tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
return assets
|
|
|
|
@staticmethod
|
|
def get_file_upload_url(
|
|
app_model: App,
|
|
account_id: str,
|
|
name: str,
|
|
size: int,
|
|
parent_id: str | None = None,
|
|
expires_in: int = 3600,
|
|
) -> tuple[AppAssetNode, str]:
|
|
"""
|
|
Create a file node with metadata and return a pre-signed upload URL.
|
|
|
|
The file metadata is saved immediately. If the user doesn't upload,
|
|
the download will fail when the file is accessed.
|
|
|
|
Returns:
|
|
tuple of (node, upload_url)
|
|
"""
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
node_id = str(uuid4())
|
|
node = AppAssetNode.create_file(node_id, name, parent_id, size)
|
|
|
|
try:
|
|
tree.add(node)
|
|
except TreeParentNotFoundError as e:
|
|
raise AppAssetParentNotFoundError(str(e)) from e
|
|
except TreePathConflictError as e:
|
|
raise AppAssetPathConflictError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
asset_path = AssetPath.draft(app_model.tenant_id, app_model.id, node_id)
|
|
asset_storage = app_asset_storage
|
|
|
|
# put empty content to create the file record
|
|
# which avoids file not found error when uploading via presigned URL is never touched
|
|
# resulting in inconsistent state
|
|
asset_storage.save(asset_path, b"")
|
|
|
|
upload_url = asset_storage.get_upload_url(asset_path, expires_in)
|
|
|
|
return node, upload_url
|
|
|
|
@staticmethod
|
|
def batch_create_from_tree(
|
|
app_model: App,
|
|
account_id: str,
|
|
input_children: list[BatchUploadNode],
|
|
expires_in: int = 3600,
|
|
) -> list[BatchUploadNode]:
|
|
if not input_children:
|
|
return []
|
|
|
|
new_nodes: list[AppAssetNode] = []
|
|
for child in input_children:
|
|
new_nodes.extend(child.to_app_asset_nodes(None))
|
|
|
|
with AppAssetService._lock(app_model.id):
|
|
with Session(db.engine, expire_on_commit=False) as session:
|
|
assets = AppAssetService.get_or_create_assets(session, app_model, account_id)
|
|
tree = assets.asset_tree
|
|
|
|
try:
|
|
for node in new_nodes:
|
|
tree.add(node)
|
|
except TreeParentNotFoundError as e:
|
|
raise AppAssetParentNotFoundError(str(e)) from e
|
|
except TreePathConflictError as e:
|
|
raise AppAssetPathConflictError(str(e)) from e
|
|
|
|
assets.asset_tree = tree
|
|
assets.updated_by = account_id
|
|
session.commit()
|
|
|
|
asset_storage = app_asset_storage
|
|
|
|
def fill_urls(node: BatchUploadNode) -> None:
|
|
if node.node_type == AssetNodeType.FILE and node.id:
|
|
asset_path = AssetPath.draft(app_model.tenant_id, app_model.id, node.id)
|
|
node.upload_url = asset_storage.get_upload_url(asset_path, expires_in)
|
|
for child in node.children:
|
|
fill_urls(child)
|
|
|
|
for child in input_children:
|
|
fill_urls(child)
|
|
|
|
return input_children
|