From 2c7c76aec38e78fd349206a5717b2ebf5af4d3de Mon Sep 17 00:00:00 2001 From: ytqh Date: Sat, 29 Mar 2025 21:41:07 +0800 Subject: [PATCH] add user gen image --- .../app/image_generate.py | 283 ++++++++++++++++++ api/fields/end_user_fields.py | 14 + ...29_2140-e4e52e0dfb56_add_user_gen_image.py | 47 +++ api/models/model.py | 27 ++ api/services/image_generation_service.py | 169 +++++++++++ 5 files changed, 540 insertions(+) create mode 100644 api/controllers/service_api_with_auth/app/image_generate.py create mode 100644 api/migrations/versions/2025_03_29_2140-e4e52e0dfb56_add_user_gen_image.py create mode 100644 api/services/image_generation_service.py diff --git a/api/controllers/service_api_with_auth/app/image_generate.py b/api/controllers/service_api_with_auth/app/image_generate.py new file mode 100644 index 0000000000..4e36954c16 --- /dev/null +++ b/api/controllers/service_api_with_auth/app/image_generate.py @@ -0,0 +1,283 @@ +import json +import logging +import os +import uuid +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Tuple + +from controllers.service_api_with_auth import api +from controllers.service_api_with_auth.app.error import NotChatAppError +from controllers.service_api_with_auth.wraps import validate_user_token_and_extract_info +from core.app.entities.app_invoke_entities import InvokeFrom +from core.file import FileTransferMethod, FileType +from extensions.ext_database import db +from fields.end_user_fields import image_fields, image_list_fields +from flask_restful import Resource, fields, marshal_with, reqparse # type: ignore +from libs.helper import TimestampField, uuid_value +from models.enums import CreatedByRole +from models.model import App, AppMode, Conversation, EndUser, Message, UserGeneratedImage +from models.types import StringUUID +from services.image_generation_service import ImageGenerationService +from sqlalchemy.orm import Session +from werkzeug.exceptions import BadRequest, InternalServerError, NotFound + +# Constants +DEFAULT_DAILY_LIMIT = 5 +MIN_CONVERSATION_ROUNDS = 10 + + +class ImageGenerateApi(Resource): + @validate_user_token_and_extract_info + def post(self, app_model: App, end_user: EndUser): + """Generate a personalized image based on conversation content. + --- + tags: + - service/image + summary: Generate a personalized image + description: Generate an image with encouraging text based on conversation + security: + - ApiKeyAuth: [] + parameters: + - name: body + in: body + required: true + schema: + type: object + required: + - conversation_id + - content_type + properties: + conversation_id: + type: string + format: uuid + description: ID of the conversation to use for image generation + content_type: + type: string + enum: [self_message, summary_advice] + description: Type of text content to generate + responses: + 200: + description: Image generation started + schema: + type: object + properties: + result: + type: string + example: success + message: + type: string + example: Image generation started + 400: + description: Invalid request or conversation not suitable + 401: + description: Invalid or missing token + 403: + description: Daily limit reached + 404: + description: Conversation not found or not a chat app + """ + app_mode = AppMode.value_of(app_model.mode) + if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}: + raise NotChatAppError() + + parser = reqparse.RequestParser() + parser.add_argument("conversation_id", required=True, type=uuid_value, location="json") + parser.add_argument( + "content_type", required=True, type=str, choices=["self_message", "summary_advice"], location="json" + ) + args = parser.parse_args() + + conversation_id = str(args["conversation_id"]) + content_type = args["content_type"] + + # Check if conversation exists + conversation = ( + db.session.query(Conversation) + .filter( + Conversation.id == conversation_id, + Conversation.app_id == app_model.id, + Conversation.from_end_user_id == end_user.id, + ) + .first() + ) + + if not conversation: + raise NotFound("Conversation not found") + + # Check if conversation has enough rounds + messages_count = db.session.query(Message).filter(Message.conversation_id == conversation_id).count() + + if messages_count < MIN_CONVERSATION_ROUNDS: + return { + "result": "error", + "message": "I need to know more about you before creating an image. Please continue our conversation.", + }, 400 + + # Check if user has reached daily limit + today = datetime.utcnow().date() + tomorrow = today + timedelta(days=1) + today_start = datetime.combine(today, datetime.min.time()) + today_end = datetime.combine(tomorrow, datetime.min.time()) + + daily_count = ( + db.session.query(UserGeneratedImage) + .filter( + UserGeneratedImage.end_user_id == end_user.id, + UserGeneratedImage.created_at >= today_start, + UserGeneratedImage.created_at < today_end, + ) + .count() + ) + + if daily_count >= DEFAULT_DAILY_LIMIT: + return { + "result": "error", + "message": ( + f"You've reached your daily limit of {DEFAULT_DAILY_LIMIT} generated images. Please try again tomorrow." + ), + }, 403 + + try: + # Use the service to generate the image + # This would typically be done asynchronously in a background task + # For simplicity, we're doing it synchronously here + image_id = ImageGenerationService.process_image_generation_request( + app_id=str(app_model.id), + conversation_id=conversation_id, + end_user_id=str(end_user.id), + content_type=content_type, + ) + + if image_id: + return {"result": "success", "message": "Image generated successfully.", "image_id": image_id} + else: + return {"result": "error", "message": "Failed to generate image. Please try again later."}, 500 + + except Exception as e: + raise InternalServerError("Failed to generate image") + + +class ImageListApi(Resource): + @validate_user_token_and_extract_info + @marshal_with(image_list_fields) + def get(self, app_model: App, end_user: EndUser): + """Get user-generated images list. + --- + tags: + - service/image + summary: List user generated images + description: Get a list of images generated for the current user + security: + - ApiKeyAuth: [] + parameters: + - name: limit + in: query + type: integer + minimum: 1 + maximum: 100 + default: 20 + description: Number of images to return + - name: offset + in: query + type: integer + minimum: 0 + default: 0 + description: Offset for pagination + responses: + 200: + description: Images retrieved successfully + schema: + type: object + properties: + data: + type: array + items: + type: object + has_more: + type: boolean + 401: + description: Invalid or missing token + 404: + description: Not a chat app + """ + app_mode = AppMode.value_of(app_model.mode) + if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}: + raise NotChatAppError() + + parser = reqparse.RequestParser() + parser.add_argument("limit", type=int, required=False, default=20, location="args") + parser.add_argument("offset", type=int, required=False, default=0, location="args") + args = parser.parse_args() + + limit = min(args["limit"], 100) + offset = max(args["offset"], 0) + + # Get images for the user + query = ( + db.session.query(UserGeneratedImage) + .filter(UserGeneratedImage.app_id == app_model.id, UserGeneratedImage.end_user_id == end_user.id) + .order_by(UserGeneratedImage.created_at.desc()) + ) + + total_count = query.count() + images = query.limit(limit).offset(offset).all() + + return {"data": images, "has_more": (offset + limit) < total_count} + + +class ImageDetailApi(Resource): + @validate_user_token_and_extract_info + @marshal_with(image_fields) + def get(self, app_model: App, end_user: EndUser, image_id): + """Get a specific generated image. + --- + tags: + - service/image + summary: Get image details + description: Get details of a specific generated image + security: + - ApiKeyAuth: [] + parameters: + - name: image_id + in: path + required: true + type: string + format: uuid + description: ID of the image to retrieve + responses: + 200: + description: Image retrieved successfully + schema: + type: object + 401: + description: Invalid or missing token + 404: + description: Image not found or not a chat app + """ + app_mode = AppMode.value_of(app_model.mode) + if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}: + raise NotChatAppError() + + image_id = str(image_id) + + # Get the image + image = ( + db.session.query(UserGeneratedImage) + .filter( + UserGeneratedImage.id == image_id, + UserGeneratedImage.app_id == app_model.id, + UserGeneratedImage.end_user_id == end_user.id, + ) + .first() + ) + + if not image: + raise NotFound("Image not found") + + return image + + +# Register API resources +api.add_resource(ImageGenerateApi, "/images/generate") +api.add_resource(ImageListApi, "/images") +api.add_resource(ImageDetailApi, "/images/", endpoint="image_detail") diff --git a/api/fields/end_user_fields.py b/api/fields/end_user_fields.py index 0c0d41cd04..eb03946919 100644 --- a/api/fields/end_user_fields.py +++ b/api/fields/end_user_fields.py @@ -26,3 +26,17 @@ end_users_infinite_scroll_pagination_fields = { "total": fields.Integer, "data": fields.List(fields.Nested(detailed_end_user_fields)), } + +# Image generation fields definition +image_fields = { + "id": fields.String, + "image_url": fields.String, + "content_type": fields.String, + "text_content": fields.String, + "created_at": TimestampField, +} + +image_list_fields = { + "data": fields.List(fields.Nested(image_fields)), + "has_more": fields.Boolean, +} diff --git a/api/migrations/versions/2025_03_29_2140-e4e52e0dfb56_add_user_gen_image.py b/api/migrations/versions/2025_03_29_2140-e4e52e0dfb56_add_user_gen_image.py new file mode 100644 index 0000000000..9d57717622 --- /dev/null +++ b/api/migrations/versions/2025_03_29_2140-e4e52e0dfb56_add_user_gen_image.py @@ -0,0 +1,47 @@ +"""add user gen image + +Revision ID: e4e52e0dfb56 +Revises: 4b37d4034604 +Create Date: 2025-03-29 21:40:54.197571 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = 'e4e52e0dfb56' +down_revision = '4b37d4034604' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('user_generated_images', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('app_id', models.types.StringUUID(), nullable=False), + sa.Column('end_user_id', models.types.StringUUID(), nullable=False), + sa.Column('conversation_id', models.types.StringUUID(), nullable=False), + sa.Column('image_url', sa.Text(), nullable=False), + sa.Column('content_type', sa.String(length=255), nullable=False), + sa.Column('text_content', sa.Text(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='user_generated_image_pkey') + ) + with op.batch_alter_table('user_generated_images', schema=None) as batch_op: + batch_op.create_index('user_generated_image_app_idx', ['app_id'], unique=False) + batch_op.create_index('user_generated_image_user_idx', ['end_user_id'], unique=False) + + # ### 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.drop_index('user_generated_image_user_idx') + batch_op.drop_index('user_generated_image_app_idx') + + op.drop_table('user_generated_images') + # ### end Alembic commands ### diff --git a/api/models/model.py b/api/models/model.py index 070dc440ef..c4dfe632d5 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1843,3 +1843,30 @@ class TraceAppConfig(db.Model): # type: ignore[name-defined] "created_at": str(self.created_at) if self.created_at else None, "updated_at": str(self.updated_at) if self.updated_at else None, } + + +# User generated image model +class UserGeneratedImage(db.Model): # type: ignore[name-defined] + __tablename__ = "user_generated_images" + __table_args__ = ( + db.PrimaryKeyConstraint("id", name="user_generated_image_pkey"), + db.Index("user_generated_image_user_idx", "end_user_id"), + db.Index("user_generated_image_app_idx", "app_id"), + ) + + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) + app_id = db.Column(StringUUID, nullable=False) + end_user_id = db.Column(StringUUID, nullable=False) + conversation_id = db.Column(StringUUID, nullable=False) + image_url = db.Column(db.Text, nullable=False) + content_type = db.Column(db.String(255), nullable=False) # 'self_message' or 'summary_advice' + text_content = db.Column(db.Text, nullable=False) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.current_timestamp()) + + @property + def end_user(self): + return db.session.query(EndUser).filter(EndUser.id == self.end_user_id).first() + + @property + def conversation(self): + return db.session.query(Conversation).filter(Conversation.id == self.conversation_id).first() diff --git a/api/services/image_generation_service.py b/api/services/image_generation_service.py new file mode 100644 index 0000000000..1ee60e894a --- /dev/null +++ b/api/services/image_generation_service.py @@ -0,0 +1,169 @@ +import json +import logging +import os +import random +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple, Union + +# Import the UserGeneratedImage model from our controller module +# This is a bit of a circular import, but it's the simplest solution for now +from controllers.service_api_with_auth.app.image_generate import UserGeneratedImage +from extensions.ext_database import db +from models.enums import CreatedByRole +from models.model import App, Conversation, EndUser, Message, UploadFile +from sqlalchemy.orm import Session + +# Configure logging +logger = logging.getLogger(__name__) + + +class ImageGenerationService: + @staticmethod + def generate_motivational_text(conversation_id: str, content_type: str) -> str: + """ + Generate motivational text based on conversation history. + + Args: + conversation_id: The ID of the conversation + content_type: Type of content to generate ('self_message' or 'summary_advice') + + Returns: + str: Generated text content + """ + # In a real implementation, this would call a large language model + # Here we'll just use placeholders based on the content type + + with Session(db.engine) as session: + # Get the last few messages from the conversation to understand context + messages = ( + session.query(Message) + .filter(Message.conversation_id == conversation_id) + .order_by(Message.created_at.desc()) + .limit(20) + .all() + ) + + # Reverse to get chronological order + messages.reverse() + + # Extract conversation context + context = "\n".join([f"User: {msg.query}\nAI: {msg.answer}" for msg in messages]) + + # In production, you would pass this context to a language model + # For demonstration, we'll return placeholder text + + if content_type == "self_message": + sample_messages = [ + "You've got this! Take one small step today.", + "Remember your strength - you've overcome challenges before.", + "Be kind to yourself today, you deserve it.", + "Your feelings are valid, and you have the power to work through them.", + "Small progress is still progress. Celebrate your wins today.", + ] + return random.choice(sample_messages) + else: # summary_advice + sample_advice = [ + "Based on our conversation, I notice you tend to be hard on yourself. Try practicing self-compassion by speaking to yourself as you would to a friend.", + "I've observed that you often describe feeling overwhelmed. Breaking tasks into smaller steps might help manage these feelings better.", + "In our discussions, I noticed patterns of negative self-talk. Consider challenging these thoughts by asking 'Is this really true?' when they arise.", + "From our conversations, it seems you might benefit from more self-care routines. Even 5 minutes of mindfulness daily could make a difference.", + "You've mentioned feeling anxious in social situations. Progressive exposure to small social interactions might help build confidence over time.", + ] + return random.choice(sample_advice) + + @staticmethod + def generate_background_image() -> str: + """ + Generate or select a background image. + + In a real implementation, this might call an image generation API + or select from pre-generated images. + + Returns: + str: URL of the generated/selected image + """ + # In a real implementation, this would integrate with an image generation API + # or select from a pool of pre-generated images + + # For this example, we'll return a placeholder + placeholder_images = [ + "https://example.com/background1.jpg", + "https://example.com/background2.jpg", + "https://example.com/background3.jpg", + "https://example.com/background4.jpg", + "https://example.com/background5.jpg", + ] + + return random.choice(placeholder_images) + + @staticmethod + def overlay_text_on_image(image_url: str, text: str) -> str: + """ + Overlay text on the image. + + In a real implementation, this would use image processing libraries. + + Args: + image_url: URL of the background image + text: Text to overlay on the image + + Returns: + str: URL of the final image with text + """ + # In a real implementation, this would use image processing libraries + # like Pillow to overlay text on the image + + # For this example, we'll just return the same URL + # In production, you would process the image and save it to storage + return image_url + + @staticmethod + def process_image_generation_request( + app_id: str, conversation_id: str, end_user_id: str, content_type: str + ) -> Optional[str]: + """ + Process an image generation request. + + Args: + app_id: The ID of the app + conversation_id: The ID of the conversation + end_user_id: The ID of the end user + content_type: Type of content to generate ('self_message' or 'summary_advice') + + Returns: + Optional[str]: ID of the generated image if successful, None otherwise + """ + try: + # 1. Generate motivational text based on conversation history + text_content = ImageGenerationService.generate_motivational_text( + conversation_id=conversation_id, content_type=content_type + ) + + # 2. Generate or select a background image + image_url = ImageGenerationService.generate_background_image() + + # 3. Overlay text on the image + final_image_url = ImageGenerationService.overlay_text_on_image(image_url=image_url, text=text_content) + + # 4. Create and save the user generated image record + with Session(db.engine) as session: + new_image = UserGeneratedImage( + app_id=app_id, + end_user_id=end_user_id, + conversation_id=conversation_id, + image_url=final_image_url, + content_type=content_type, + text_content=text_content, + ) + + session.add(new_image) + session.commit() + + image_id = str(new_image.id) + logger.info(f"Generated image {image_id} for user {end_user_id}") + + return image_id + + except Exception as e: + logger.error(f"Error generating image: {str(e)}") + return None