use file id instead of image url

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 ###

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(

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

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"