add user gen image
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")
|
||||
@ -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 ###
|
||||
@ -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…
Reference in New Issue