add user gen image

pull/21891/head
ytqh 1 year ago
parent 7fc884acdc
commit 2c7c76aec3

@ -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/<uuid:image_id>", endpoint="image_detail")

@ -26,3 +26,17 @@ end_users_infinite_scroll_pagination_fields = {
"total": fields.Integer, "total": fields.Integer,
"data": fields.List(fields.Nested(detailed_end_user_fields)), "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,
}

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

@ -1843,3 +1843,30 @@ class TraceAppConfig(db.Model): # type: ignore[name-defined]
"created_at": str(self.created_at) if self.created_at else None, "created_at": str(self.created_at) if self.created_at else None,
"updated_at": str(self.updated_at) if self.updated_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()

@ -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
Loading…
Cancel
Save