feat: Implement image thumbnail generation, storage, and API endpoints for assets, including a regeneration utility.

api/endpoints/assets_router.py

@@ -21,17 +21,31 @@ router = APIRouter(prefix="/api/assets", tags=["Assets"])
 
 
 @router.get("/{asset_id}")
-async def get_asset(asset_id: str, request: Request,dao: DAO = Depends(get_dao),) -> Response:
+async def get_asset(
-    logger.debug(f"get_asset called for ID: {asset_id}")
+    asset_id: str, 
+    request: Request,
+    thumbnail: bool = False,
+    dao: DAO = Depends(get_dao)
+) -> Response:
+    logger.debug(f"get_asset called for ID: {asset_id}, thumbnail={thumbnail}")
     asset = await dao.assets.get_asset(asset_id)
     # 2. Проверка на существование
     if not asset:
         raise HTTPException(status_code=404, detail="Asset not found")
+    
     headers = {
         # Кэшировать на 1 год (31536000 сек)
         "Cache-Control": "public, max-age=31536000, immutable"
     }
-    return Response(content=asset.data, media_type="image/png", headers=headers)
+    
+    content = asset.data
+    media_type = "image/png" # Default, or detect
+    
+    if thumbnail and asset.thumbnail:
+        content = asset.thumbnail
+        media_type = "image/jpeg"
+        
+    return Response(content=content, media_type=media_type, headers=headers)
 
 
 @router.get("")
@@ -41,7 +55,13 @@ async def get_assets(request: Request, dao: DAO = Depends(get_dao), limit: int =
     # assets = await dao.assets.get_assets() # This line seemed redundant/conflicting in original code
     total_count = await dao.assets.get_asset_count()
 
-    return AssetsResponse(assets=assets, total_count=total_count)
+    # Manually map to DTO to trigger computed fields validation if necessary, 
+    # but primarily to ensure valid Pydantic models for the response list.
+    # Asset.model_dump() generally includes computed fields (url) if configured.
+    # Let's ensure strict conversion.
+    asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets]
+    
+    return AssetsResponse(assets=asset_responses, total_count=total_count)
 
 
 
@@ -62,11 +82,16 @@ async def upload_asset(
     if not data:
         raise HTTPException(status_code=400, detail="Empty file")
 
+    # Generate thumbnail
+    from utils.image_utils import create_thumbnail
+    thumbnail_bytes = await asyncio.to_thread(create_thumbnail, data)
+
     asset = Asset(
         name=file.filename or "upload",
         type=AssetType.IMAGE,
         linked_char_id=linked_char_id,
         data=data,
+        thumbnail=thumbnail_bytes
     )
 
     asset_id = await dao.assets.create_asset(asset)
@@ -79,4 +104,39 @@ async def upload_asset(
         type=asset.type.value if hasattr(asset.type, "value") else asset.type,
         linked_char_id=asset.linked_char_id,
         created_at=asset.created_at,
-    )
+        url=asset.url
+    )
+
+
+@router.post("/regenerate_thumbnails")
+async def regenerate_thumbnails(dao: DAO = Depends(get_dao)):
+    """
+    Regenerates thumbnails for all existing image assets that don't have one.
+    """
+    logger.info("Starting thumbnail regeneration task")
+    from utils.image_utils import create_thumbnail
+    import asyncio
+
+    # Get all assets (pagination loop might be needed for huge datasets, but simple list for now)
+    # We'll rely on DAO providing a method or just fetch large chunk.
+    # Assuming get_assets might have limit, let's fetch in chunks or just all if possible within limit.
+    # Ideally should use a specific repo method for iteration.
+    # For now, let's fetch first 1000 or similar.
+    assets = await dao.assets.get_assets(limit=1000, offset=0, with_data=True) 
+    logger.info(f"Found {len(assets)} assets")
+    count = 0
+    updated = 0
+    
+    for asset in assets:
+        if asset.type == AssetType.IMAGE and asset.data :
+            try:
+                thumb = await asyncio.to_thread(create_thumbnail, asset.data)
+                if thumb:
+                    asset.thumbnail = thumb
+                    await dao.assets.update_asset(asset.id, asset)
+                    updated += 1
+            except Exception as e:
+                logger.error(f"Failed to regenerate thumbnail for asset {asset.id}: {e}")
+        count += 1
+        
+    return {"status": "completed", "processed": count, "updated": updated}

api/endpoints/character_router.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
 
-from api.models.AssetDTO import AssetsResponse
+from api.models.AssetDTO import AssetsResponse, AssetResponse
 from api.models.GenerationRequest import GenerationRequest, GenerationResponse
 from models.Asset import Asset
 from models.Character import Character
@@ -35,7 +35,9 @@ async def get_character_assets(character_id: str, dao: DAO = Depends(get_dao), l
         raise HTTPException(status_code=404, detail="Character not found")
     assets = await dao.assets.get_assets_by_char_id(character_id, limit, offset)
     total_count = await dao.assets.get_asset_count(character_id)
-    return AssetsResponse(assets=assets, total_count=total_count)
+    
+    asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets]
+    return AssetsResponse(assets=asset_responses, total_count=total_count)
 
 
 @router.get("/{character_id}", response_model=Character)

api/models/AssetDTO.py

@@ -6,14 +6,14 @@ from pydantic import BaseModel
 from models.Asset import Asset
 
 
-class AssetsResponse(BaseModel):
-    assets: List[Asset]
-    total_count: int
-
-
 class AssetResponse(BaseModel):
     id: str
     name: str
     type: str
     linked_char_id: Optional[str] = None
-    created_at: datetime
+    created_at: datetime
+    url: Optional[str] = None
+
+class AssetsResponse(BaseModel):
+    assets: List[AssetResponse]
+    total_count: int

api/service/generation_service.py

@@ -66,7 +66,7 @@ class GenerationService:
     async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str:
         future_prompt = """You are an prompt-assistant. You improving user-entered prompts for image generation. User may upload reference image too.
          I will provide sources prompt entered by user. Understand user needs and generate best variation of prompt.
-          ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT:"""
+          ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT: """
         future_prompt += prompt
         assets_data = []
         if assets is not None:
@@ -88,7 +88,7 @@ class GenerationService:
 
     async def get_generations(self, character_id: Optional[str] = None, limit: int = 10, offset: int = 0) -> List[
         Generation]:
-        return await self.dao.generations.get_generations(limit=limit, offset=offset)
+        return await self.dao.generations.get_generations(character_id = character_id,limit=limit, offset=offset)
 
     async def get_generation(self, generation_id: str) -> Optional[GenerationResponse]:
         gen = await self.dao.generations.get_generation(generation_id)
@@ -162,7 +162,7 @@ class GenerationService:
             for asset in reference_assets
             if asset.data is not None and asset.type == AssetType.IMAGE
         )
-        generation_prompt+=f"PROMPT: {generation.prompt}"
+        generation_prompt+=f" PROMPT: {generation.prompt}"
         logger.info(f"Final generation prompt assembled. Length: {len(generation_prompt)}. Media count: {len(media_group_bytes)}")
 
         # 3. Запускаем процесс генерации и симуляцию прогресса
@@ -209,11 +209,19 @@ class GenerationService:
         created_assets: List[Asset] = []
 
         for idx, img_bytes in enumerate(generated_bytes_list):
+            # Generate thumbnail
+            thumbnail_bytes = None
+            # Assuming AssetType.IMAGE since we are in generated_bytes_list which are images usually
+            # Or use explicit check if we have distinct types in list (not currently)
+            from utils.image_utils import create_thumbnail
+            thumbnail_bytes = await asyncio.to_thread(create_thumbnail, img_bytes)
+            
             new_asset = Asset(
                 name=f"Generated_{generation.linked_character_id}_{random.randint(1000, 9999)}",
                 type=AssetType.IMAGE,
                 linked_char_id=generation.linked_character_id,  # Если генерация привязана к персонажу
                 data=img_bytes,
+                thumbnail=thumbnail_bytes
                 # Остальные поля заполнятся дефолтными значениями (created_at)
             )
 
@@ -236,6 +244,8 @@ class GenerationService:
         end_time = datetime.now()
         generation.execution_time_seconds = (end_time - start_time).total_seconds()
         
+        logger.info(f"DEBUG: Saving generation {generation.id}. Metrics: api_exec={generation.api_execution_time_seconds}, tokens={generation.token_usage}, exec={generation.execution_time_seconds}")
+        
         await self.dao.generations.update_generation(generation)
         logger.info(f"Generation {generation.id} completed successfully. {len(created_assets)} assets created. Total Time: {generation.execution_time_seconds:.2f}s")
 

models/Asset.py

@@ -18,6 +18,7 @@ class Asset(BaseModel):
     data: Optional[bytes] = None
     tg_doc_file_id: Optional[str] = None
     tg_photo_file_id: Optional[str] = None
+    thumbnail: Optional[bytes] = None
     tags: List[str] = []
     created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
     updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

repos/assets_repo.py

@@ -15,8 +15,12 @@ class AssetsRepo:
 
         return str(res.inserted_id)
 
-    async def get_assets(self, limit: int = 10, offset: int = 0) -> List[Asset]:
+    async def get_assets(self, limit: int = 10, offset: int = 0, with_data: bool = False) -> List[Asset]:
-        res = await self.collection.find({}, {"data": 0}).sort("created_at", -1).skip(offset).limit(limit).to_list(None)
+        args = {}
+        if not with_data:
+            args["data"] = 0
+            args["thumbnail"] = 0
+        res = await self.collection.find({}, args).sort("created_at", -1).skip(offset).limit(limit).to_list(None)
         assets = []
         for doc in res:
             # Конвертируем ObjectId в строку и кладем в поле id
@@ -32,6 +36,7 @@ class AssetsRepo:
         projection = {"_id": 1, "name": 1, "type": 1, "tg_doc_file_id": 1}
         if with_data:
             projection["data"] = 1
+            projection["thumbnail"] = 1
 
         res = await self.collection.find_one({"_id": ObjectId(asset_id)},
                                              projection)
@@ -63,7 +68,7 @@ class AssetsRepo:
 
     async def get_assets_by_ids(self, asset_ids: List[str]) -> List[Asset]:
         object_ids = [ObjectId(asset_id) for asset_id in asset_ids]
-        res = self.collection.find({"_id": {"$in": object_ids}})
+        res = self.collection.find({"_id": {"$in": object_ids}}, {"thumbnail": 0})
         assets = []
         async for doc in res:
             doc["id"] = str(doc.pop("_id"))

repos/generation_repo.py

@@ -28,7 +28,9 @@ class GenerationRepo:
                               limit: int = 10, offset: int = 10) -> List[Generation]:
         args = {}
         if character_id is not None:
-            args["character_id"] = character_id
+            args["linked_character_id"] = character_id
+        else:
+            args["linked_character_id"] = None
         if status is not None:
             args["status"] = status
         res = await self.collection.find(args).sort("created_at", -1).skip(

utils/image_utils.py

@@ -0,0 +1,27 @@
+from io import BytesIO
+from typing import Tuple, Optional
+from PIL import Image
+import logging
+
+logger = logging.getLogger(__name__)
+
+def create_thumbnail(image_data: bytes, size: Tuple[int, int] = (800, 800)) -> Optional[bytes]:
+    """
+    Creates a thumbnail from image bytes.
+    Returns the thumbnail as bytes (JPEG format) or None if failed.
+    """
+    try:
+        with Image.open(BytesIO(image_data)) as img:
+            # Convert to RGB if necessary (e.g. for RGBA/P images saving as JPEG)
+            if img.mode in ('RGBA', 'P'):
+                img = img.convert('RGB')
+            
+            img.thumbnail(size)
+            
+            thumb_io = BytesIO()
+            img.save(thumb_io, format='JPEG', quality=85)
+            thumb_io.seek(0)
+            return thumb_io.read()
+    except Exception as e:
+        logger.error(f"Failed to create thumbnail: {e}")
+        return None