diff --git a/api/migrations/versions/2025_04_08_2116-69f4e59ecafb_add_file_id_instead_of_image_url.py b/api/migrations/versions/2025_04_08_2116-69f4e59ecafb_add_file_id_instead_of_image_url.py new file mode 100644 index 0000000000..5cd0e270ca --- /dev/null +++ b/api/migrations/versions/2025_04_08_2116-69f4e59ecafb_add_file_id_instead_of_image_url.py @@ -0,0 +1,35 @@ +"""add file id instead of image url + +Revision ID: 69f4e59ecafb +Revises: 382f0cad6652 +Create Date: 2025-04-08 21:16:17.846817 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '69f4e59ecafb' +down_revision = '382f0cad6652' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('user_generated_images', schema=None) as batch_op: + batch_op.add_column(sa.Column('file_id', models.types.StringUUID(), nullable=True)) + batch_op.drop_column('image_url') + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('user_generated_images', schema=None) as batch_op: + batch_op.add_column(sa.Column('image_url', sa.TEXT(), autoincrement=False, nullable=True)) + batch_op.drop_column('file_id') + + # ### end Alembic commands ### diff --git a/api/models/model.py b/api/models/model.py index 24a0f9bc48..70c7afcdb5 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1868,7 +1868,7 @@ class UserGeneratedImage(db.Model): # type: ignore[name-defined] end_user_id = db.Column(StringUUID, nullable=False) workflow_run_id = db.Column(StringUUID, nullable=True) # related generation id (nullable for pending status) content_type = db.Column(db.String(255), nullable=False) # 'self_message' or 'summary_advice' - image_url = db.Column(db.Text, nullable=True) + file_id = db.Column(StringUUID, nullable=True) text_content = db.Column(db.Text, nullable=True) raw_content = db.Column(db.JSON, nullable=True) # save raw llm outputs status = db.Column( diff --git a/api/services/image_generation_service.py b/api/services/image_generation_service.py index 66edf5c172..047328c8dd 100644 --- a/api/services/image_generation_service.py +++ b/api/services/image_generation_service.py @@ -1,13 +1,11 @@ -import json from enum import Enum from configs import dify_config -from core.app.entities.app_invoke_entities import InvokeFrom +from core.tools.tool_file_manager import ToolFileManager from extensions.ext_database import db from libs.helper import RateLimiter from libs.infinite_scroll_pagination import MultiPagePagination -from models.model import App, EndUser, Message, UserGeneratedImage -from services.app_generate_service import AppGenerateService +from models.model import EndUser, UserGeneratedImage from tasks.image_generation_task import generate_image_task @@ -17,6 +15,9 @@ class ContentType(str, Enum): SUMMARY_ADVICE = "summary_advice" +DEFAULT_IMAGE_EXTENSION = ".png" + + class ImageGenerationService: generate_image_rate_limiter = RateLimiter( @@ -71,6 +72,11 @@ class ImageGenerationService: total_count = query.count() images = query.limit(limit).offset(offset).all() + # sign file with file_id to get a temporary url + for image in images: + if image.file_id: + image.image_url = ToolFileManager.sign_file(image.file_id, DEFAULT_IMAGE_EXTENSION) + return MultiPagePagination(data=images, total=total_count) @staticmethod @@ -80,4 +86,8 @@ class ImageGenerationService: if image is None: raise Exception("Image not found") + # sign file with file_id to get a temporary url + if image.file_id: + image.image_url = ToolFileManager.sign_file(image.file_id, DEFAULT_IMAGE_EXTENSION) + return image diff --git a/api/tasks/image_generation_task.py b/api/tasks/image_generation_task.py index 1e822b9005..1dad827eb2 100644 --- a/api/tasks/image_generation_task.py +++ b/api/tasks/image_generation_task.py @@ -117,11 +117,14 @@ def generate_image_task( db.session.commit() raise Exception("Failed to generate image") + # parse file id from url: f"{base_url}/files/tools/{tool_file_id}{extension}" + file_id = url.split("/files/tools/")[1].split(".")[0] + text_content = raw_content.get("text") # Update the existing UserGeneratedImage with the generated content user_generated_image.workflow_run_id = workflow_run_id - user_generated_image.image_url = url + user_generated_image.file_id = file_id user_generated_image.text_content = text_content user_generated_image.raw_content = raw_content user_generated_image.status = "completed"