diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..7408157 Binary files /dev/null and b/.DS_Store differ diff --git a/.env b/.env index 33f109c..27f717c 100644 --- a/.env +++ b/.env @@ -1,4 +1,5 @@ -BOT_TOKEN=8495170789:AAHyjjhHwwVtd9_ROnjHqPHRdnmyVr1aeaY +#BOT_TOKEN=8495170789:AAHyjjhHwwVtd9_ROnjHqPHRdnmyVr1aeaY +BOT_TOKEN=8011562605:AAF3kyzrZJgii0Jx-H8Sum5Njbo0BdbsiAo GEMINI_API_KEY=AIzaSyAHzDYhgjOqZZnvOnOFRGaSkKu4OAN3kZE MONGO_HOST=mongodb://admin:super_secure_password@31.59.58.220:27017/ ADMIN_ID=567047 \ No newline at end of file diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..13566b8 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,8 @@ +# Default ignored files +/shelf/ +/workspace.xml +# Editor-based HTTP Client requests +/httpRequests/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/.idea/ai-char-bot.iml b/.idea/ai-char-bot.iml new file mode 100644 index 0000000..552e188 --- /dev/null +++ b/.idea/ai-char-bot.iml @@ -0,0 +1,10 @@ + + + + + + + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml new file mode 100644 index 0000000..476b22b --- /dev/null +++ b/.idea/inspectionProfiles/Project_Default.xml @@ -0,0 +1,16 @@ + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100644 index 0000000..105ce2d --- /dev/null +++ b/.idea/inspectionProfiles/profiles_settings.xml @@ -0,0 +1,6 @@ + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..6d15bbe --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,7 @@ + + + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..b3ccf67 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..97b8ad3 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,5 @@ +{ + "configurations": [ + {"name":"Python Debugger: FastAPI","type":"debugpy","request":"launch","module":"uvicorn","args":["main:app","--reload", "--port", "8090"],"jinja":true} + ] +} \ No newline at end of file diff --git a/__pycache__/keyboards.cpython-313.pyc b/__pycache__/keyboards.cpython-313.pyc new file mode 100644 index 0000000..65f1ac8 Binary files /dev/null and b/__pycache__/keyboards.cpython-313.pyc differ diff --git a/__pycache__/main.cpython-313.pyc b/__pycache__/main.cpython-313.pyc new file mode 100644 index 0000000..12cde9d Binary files /dev/null and b/__pycache__/main.cpython-313.pyc differ diff --git a/adapters/Exception.py b/adapters/Exception.py index 173335b..b330748 100644 --- a/adapters/Exception.py +++ b/adapters/Exception.py @@ -1,3 +1,4 @@ class GoogleGenerationException(Exception): - message: str - pass \ No newline at end of file + def __init__(self, message: str): + self.message = message + super().__init__(message) \ No newline at end of file diff --git a/adapters/__pycache__/Exception.cpython-313.pyc b/adapters/__pycache__/Exception.cpython-313.pyc new file mode 100644 index 0000000..6528c4b Binary files /dev/null and b/adapters/__pycache__/Exception.cpython-313.pyc differ diff --git a/adapters/__pycache__/google_adapter.cpython-313.pyc b/adapters/__pycache__/google_adapter.cpython-313.pyc index bb7344a..c1a7312 100644 Binary files a/adapters/__pycache__/google_adapter.cpython-313.pyc and b/adapters/__pycache__/google_adapter.cpython-313.pyc differ diff --git a/adapters/__pycache__/kling_adapter.cpython-313.pyc b/adapters/__pycache__/kling_adapter.cpython-313.pyc new file mode 100644 index 0000000..236ecec Binary files /dev/null and b/adapters/__pycache__/kling_adapter.cpython-313.pyc differ diff --git a/adapters/google_adapter.py b/adapters/google_adapter.py index bbabfd8..7298066 100644 --- a/adapters/google_adapter.py +++ b/adapters/google_adapter.py @@ -1,7 +1,7 @@ import io import logging from datetime import datetime -from typing import List, Union +from typing import List, Union, Tuple, Dict, Any from PIL import Image from google import genai @@ -27,6 +27,7 @@ class GoogleAdapter: """Вспомогательный метод для подготовки контента (текст + картинки)""" contents = [prompt] if images_list: + logger.info(f"Preparing content with {len(images_list)} images") for img_bytes in images_list: try: # Gemini API требует PIL Image на входе @@ -34,6 +35,8 @@ class GoogleAdapter: contents.append(image) except Exception as e: logger.error(f"Error processing input image: {e}") + else: + logger.info("Preparing content with no images") return contents def generate_text(self, prompt: str, images_list: List[bytes] = None) -> str: @@ -59,20 +62,25 @@ class GoogleAdapter: for part in response.parts: if part.text: result_text += part.text - logger.error(f"Generated text: {result_text}") + logger.info(f"Generated text length: {len(result_text)}") return result_text except Exception as e: logger.error(f"Gemini Text API Error: {e}") - return f"Ошибка генерации текста: {e}" + raise GoogleGenerationException(f"Gemini Text API Error: {e}") - def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] = None, ) -> List[io.BytesIO]: + def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]: """ Генерация изображений (Text-to-Image или Image-to-Image). Возвращает список байтовых потоков (готовых к отправке). """ - contents = self._prepare_contents(prompt, images_list) + contents = self._prepare_contents(prompt, images_list) + logger.info(f"Generating image. Prompt length: {len(prompt)}, Ratio: {aspect_ratio}, Quality: {quality}") + + start_time = datetime.now() + token_usage = 0 + try: response = self.client.models.generate_content( model=self.IMAGE_MODEL, @@ -86,6 +94,13 @@ class GoogleAdapter: ), ) ) + + end_time = datetime.now() + api_duration = (end_time - start_time).total_seconds() + + if response.usage_metadata: + token_usage = response.usage_metadata.total_token_count + if response.parts is None and response.candidates[0].finish_reason is not None: raise GoogleGenerationException(f"Generation blocked in cause of {response.candidates[0].finish_reason.value}") @@ -111,9 +126,17 @@ class GoogleAdapter: except Exception as e: logger.error(f"Error processing output image: {e}") - return generated_images + if generated_images: + logger.info(f"Successfully generated {len(generated_images)} images in {api_duration:.2f}s. Tokens: {token_usage}") + else: + logger.warning("No images text generated from parts") + + metrics = { + "api_execution_time_seconds": api_duration, + "token_usage": token_usage + } + return generated_images, metrics except Exception as e: logger.error(f"Gemini Image API Error: {e}") - # В случае ошибки возвращаем пустой список (или можно рейзить исключение) - return [] \ No newline at end of file + raise GoogleGenerationException(f"Gemini Image API Error: {e}") \ No newline at end of file diff --git a/api/__pycache__/__init__.cpython-313.pyc b/api/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..b8ed943 Binary files /dev/null and b/api/__pycache__/__init__.cpython-313.pyc differ diff --git a/api/__pycache__/dependency.cpython-313.pyc b/api/__pycache__/dependency.cpython-313.pyc new file mode 100644 index 0000000..4b06fb7 Binary files /dev/null and b/api/__pycache__/dependency.cpython-313.pyc differ diff --git a/api/dependency.py b/api/dependency.py index 70f4479..4455940 100644 --- a/api/dependency.py +++ b/api/dependency.py @@ -25,6 +25,6 @@ def get_dao(mongo_client: AsyncIOMotorClient = Depends(get_mongo_client)) -> DAO # Провайдер сервиса (собирается из DAO и Gemini) def get_generation_service( dao: DAO = Depends(get_dao), - gemini: GoogleAdapter = Depends(get_gemini_client) + gemini: GoogleAdapter = Depends(get_gemini_client), ) -> GenerationService: return GenerationService(dao, gemini) \ No newline at end of file diff --git a/api/endpoints/__pycache__/__init__.cpython-313.pyc b/api/endpoints/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..287cc9f Binary files /dev/null and b/api/endpoints/__pycache__/__init__.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/assets_router.cpython-313.pyc b/api/endpoints/__pycache__/assets_router.cpython-313.pyc new file mode 100644 index 0000000..ff27036 Binary files /dev/null and b/api/endpoints/__pycache__/assets_router.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/character_router.cpython-313.pyc b/api/endpoints/__pycache__/character_router.cpython-313.pyc new file mode 100644 index 0000000..dbd9780 Binary files /dev/null and b/api/endpoints/__pycache__/character_router.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/generation_router.cpython-313.pyc b/api/endpoints/__pycache__/generation_router.cpython-313.pyc new file mode 100644 index 0000000..117aef6 Binary files /dev/null and b/api/endpoints/__pycache__/generation_router.cpython-313.pyc differ diff --git a/api/endpoints/assets_router.py b/api/endpoints/assets_router.py index 815fee8..89985db 100644 --- a/api/endpoints/assets_router.py +++ b/api/endpoints/assets_router.py @@ -13,12 +13,16 @@ from models.Asset import Asset, AssetType from repos.dao import DAO from api.dependency import get_dao +import logging + +logger = logging.getLogger(__name__) + 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: - + logger.debug(f"get_asset called for ID: {asset_id}") asset = await dao.assets.get_asset(asset_id) # 2. Проверка на существование if not asset: @@ -32,8 +36,9 @@ async def get_asset(asset_id: str, request: Request,dao: DAO = Depends(get_dao), @router.get("") async def get_assets(request: Request, dao: DAO = Depends(get_dao), limit: int = 10, offset: int = 0) -> AssetsResponse: + logger.info(f"get_assets called. Limit: {limit}, Offset: {offset}") assets = await dao.assets.get_assets(limit, offset) - assets = await dao.assets.get_assets() + # 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) @@ -46,6 +51,7 @@ async def upload_asset( linked_char_id: Optional[str] = Form(None), dao: DAO = Depends(get_dao), ): + logger.info(f"upload_asset called. Filename: {file.filename}, ContentType: {file.content_type}, LinkedCharId: {linked_char_id}") if not file.content_type: raise HTTPException(status_code=400, detail="Unknown file type") @@ -65,6 +71,7 @@ async def upload_asset( asset_id = await dao.assets.create_asset(asset) asset.id = str(asset_id) + logger.info(f"Asset created successfully. ID: {asset_id}") return AssetResponse( id=asset.id, diff --git a/api/endpoints/character_router.py b/api/endpoints/character_router.py index 8f20bba..ae72b9b 100644 --- a/api/endpoints/character_router.py +++ b/api/endpoints/character_router.py @@ -12,11 +12,16 @@ from models.Character import Character from repos.dao import DAO from api.dependency import get_dao +import logging + +logger = logging.getLogger(__name__) + router = APIRouter(prefix="/api/characters", tags=["Characters"]) @router.get("/", response_model=List[Character]) async def get_characters(request: Request, dao: DAO = Depends(get_dao), ) -> List[Character]: + logger.info("get_characters called") characters = await dao.chars.get_all_characters() return characters @@ -24,6 +29,7 @@ async def get_characters(request: Request, dao: DAO = Depends(get_dao), ) -> Lis @router.get("/{character_id}/assets", response_model=AssetsResponse) async def get_character_assets(character_id: str, dao: DAO = Depends(get_dao), limit: int = 10, offset: int = 0, ) -> AssetsResponse: + logger.info(f"get_character_assets called. CharacterID: {character_id}, Limit: {limit}, Offset: {offset}") character = await dao.chars.get_character(character_id) if character is None: raise HTTPException(status_code=404, detail="Character not found") @@ -34,11 +40,13 @@ async def get_character_assets(character_id: str, dao: DAO = Depends(get_dao), l @router.get("/{character_id}", response_model=Character) async def get_character_by_id(character_id: str, request: Request, dao: DAO = Depends(get_dao)) -> Character: + logger.debug(f"get_character_by_id called. ID: {character_id}") character = await dao.chars.get_character(character_id) return character -@router.post("/{character_id}/_run", response_model=Asset) +@router.post("/{character_id}/_run", response_model=GenerationResponse) async def post_character_generation(character_id: str, generation: GenerationRequest, request: Request) -> GenerationResponse: + logger.info(f"post_character_generation called. CharacterID: {character_id}") generation_service = request.app.state.generation_service diff --git a/api/endpoints/generation_router.py b/api/endpoints/generation_router.py index 72d103a..0b4a092 100644 --- a/api/endpoints/generation_router.py +++ b/api/endpoints/generation_router.py @@ -1,6 +1,6 @@ from typing import List, Optional -from fastapi import APIRouter +from fastapi import APIRouter, UploadFile, File, Form from fastapi.params import Depends from starlette.requests import Request @@ -11,6 +11,10 @@ from api.models.GenerationRequest import GenerationResponse, GenerationRequest, from api.service.generation_service import GenerationService from models.Generation import Generation +import logging + +logger = logging.getLogger(__name__) + router = APIRouter(prefix='/api/generations', tags=["Generation"]) @@ -18,13 +22,31 @@ router = APIRouter(prefix='/api/generations', tags=["Generation"]) async def ask_prompt_assistant(prompt_request: PromptRequest, request: Request, generation_service: GenerationService = Depends( get_generation_service)) -> PromptResponse: + logger.info(f"ask_prompt_assistant called with prompt length: {len(prompt_request.prompt)}. Linked assets: {len(prompt_request.linked_assets) if prompt_request.linked_assets else 0}") generated_prompt = await generation_service.ask_prompt_assistant(prompt_request.prompt, prompt_request.linked_assets) return PromptResponse(prompt=generated_prompt) +@router.post("/prompt-from-image", response_model=PromptResponse) +async def prompt_from_image( + prompt: Optional[str] = Form(None), + images: List[UploadFile] = File(...), + generation_service: GenerationService = Depends(get_generation_service) +) -> PromptResponse: + logger.info(f"prompt_from_image called. Images count: {len(images)}. Prompt provided: {bool(prompt)}") + images_bytes = [] + for image in images: + content = await image.read() + images_bytes.append(content) + + generated_prompt = await generation_service.generate_prompt_from_images(images_bytes, prompt) + return PromptResponse(prompt=generated_prompt) + + @router.get("", response_model=List[GenerationResponse]) -async def get_generations(character_id: Optional[str], limit: int = 10, offset: int = 0, +async def get_generations(character_id: Optional[str] = None, limit: int = 10, offset: int = 0, generation_service: GenerationService = Depends(get_generation_service)): + logger.info(f"get_generations called. CharacterId: {character_id}, Limit: {limit}, Offset: {offset}") return await generation_service.get_generations(character_id, limit=limit, offset=offset) @@ -32,12 +54,14 @@ async def get_generations(character_id: Optional[str], limit: int = 10, offset: async def post_generation(generation: GenerationRequest, request: Request, generation_service: GenerationService = Depends( get_generation_service)) -> GenerationResponse: + logger.info(f"post_generation (run) called. LinkedCharId: {generation.linked_character_id}, PromptLength: {len(generation.prompt)}") return await generation_service.create_generation_task(generation) @router.get("/{generation_id}", response_model=GenerationResponse) async def get_generation(generation_id: str, generation_service: GenerationService = Depends(get_generation_service)) -> GenerationResponse: + logger.debug(f"get_generation called for ID: {generation_id}") return await generation_service.get_generation(generation_id) diff --git a/api/models/GenerationRequest.py b/api/models/GenerationRequest.py index 5a4c512..0b71ce8 100644 --- a/api/models/GenerationRequest.py +++ b/api/models/GenerationRequest.py @@ -5,7 +5,7 @@ from pydantic import BaseModel from models.Asset import Asset from models.Generation import GenerationStatus -from models.enums import AspectRatios, Quality +from models.enums import AspectRatios, Quality, GenType class GenerationRequest(BaseModel): @@ -20,12 +20,18 @@ class GenerationResponse(BaseModel): id: str status: GenerationStatus failed_reason: Optional[str] = None + linked_character_id: Optional[str] = None aspect_ratio: AspectRatios quality: Quality prompt: str + tech_prompt: Optional[str] = None assets_list: List[str] result: Optional[str] = None + execution_time_seconds: Optional[float] = None + api_execution_time_seconds: Optional[float] = None + token_usage: Optional[int] = None + progress: int = 0 created_at: datetime = datetime.now(UTC) updated_at: datetime = datetime.now(UTC) diff --git a/api/models/__pycache__/AssetDTO.cpython-313.pyc b/api/models/__pycache__/AssetDTO.cpython-313.pyc new file mode 100644 index 0000000..8a355fb Binary files /dev/null and b/api/models/__pycache__/AssetDTO.cpython-313.pyc differ diff --git a/api/models/__pycache__/GenerationRequest.cpython-313.pyc b/api/models/__pycache__/GenerationRequest.cpython-313.pyc new file mode 100644 index 0000000..152561e Binary files /dev/null and b/api/models/__pycache__/GenerationRequest.cpython-313.pyc differ diff --git a/api/models/__pycache__/__init__.cpython-313.pyc b/api/models/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..dca3a3a Binary files /dev/null and b/api/models/__pycache__/__init__.cpython-313.pyc differ diff --git a/api/service/__pycache__/__init__.cpython-313.pyc b/api/service/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..8f1c002 Binary files /dev/null and b/api/service/__pycache__/__init__.cpython-313.pyc differ diff --git a/api/service/__pycache__/generation_service.cpython-313.pyc b/api/service/__pycache__/generation_service.cpython-313.pyc new file mode 100644 index 0000000..62c6ea2 Binary files /dev/null and b/api/service/__pycache__/generation_service.cpython-313.pyc differ diff --git a/api/service/generation_service.py b/api/service/generation_service.py index 0091d1f..3bcce54 100644 --- a/api/service/generation_service.py +++ b/api/service/generation_service.py @@ -2,7 +2,7 @@ import asyncio import logging import random from datetime import datetime, UTC -from typing import List, Optional +from typing import List, Optional, Tuple, Any, Dict from io import BytesIO from adapters.Exception import GoogleGenerationException @@ -11,7 +11,7 @@ from api.models.GenerationRequest import GenerationRequest, GenerationResponse # Импортируйте ваши модели DAO, Asset, Generation корректно from models.Asset import Asset, AssetType from models.Generation import Generation, GenerationStatus -from models.enums import AspectRatios, Quality +from models.enums import AspectRatios, Quality, GenType from repos.dao import DAO logger = logging.getLogger(__name__) @@ -24,20 +24,24 @@ async def generate_image_task( aspect_ratio: AspectRatios, quality: Quality, gemini: GoogleAdapter -) -> List[bytes]: +) -> Tuple[List[bytes], Dict[str, Any]]: """ Обертка для вызова синхронного метода Gemini в отдельном потоке. Возвращает список байтов сгенерированных изображений. """ try : + logger.info(f"Starting generate_image_task with prompt length: {len(prompt)}") # Запускаем блокирующую операцию в отдельном потоке, чтобы не тормозить Event Loop - generated_images_io: List[BytesIO] = await asyncio.to_thread( + result = await asyncio.to_thread( gemini.generate_image, prompt=prompt, images_list=media_group_bytes, aspect_ratio=aspect_ratio, quality=quality, ) + generated_images_io, metrics = result + + logger.info(f"generate_image_task completed, received {len(generated_images_io) if generated_images_io else 0} images") except GoogleGenerationException as e: raise e images_bytes = [] @@ -51,13 +55,13 @@ async def generate_image_task( # Закрываем поток img_io.close() - return images_bytes - + return images_bytes, metrics class GenerationService: def __init__(self, dao: DAO, gemini: GoogleAdapter): self.dao = dao self.gemini = gemini + 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. @@ -68,11 +72,20 @@ class GenerationService: if assets is not None: assets_db = await self.dao.assets.get_assets_by_ids(assets) assets_data.extend(asset.data for asset in assets_db) - generated_prompt = self.gemini.generate_text(future_prompt, assets_data) + generated_prompt = await asyncio.to_thread(self.gemini.generate_text, future_prompt, assets_data) logger.info(future_prompt) logger.info(generated_prompt) return generated_prompt + async def generate_prompt_from_images(self, images: List[bytes], user_prompt: Optional[str] = None) -> str: + technical_prompt = "You are a prompt engineer. Describe this image in detail to create a stable diffusion using this image as reference. " + if user_prompt: + technical_prompt += f"User also provided this context: {user_prompt}. " + + technical_prompt += "Provide ONLY the detailed prompt." + + return await asyncio.to_thread(self.gemini.generate_text, prompt=technical_prompt, images_list=images) + 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) @@ -97,8 +110,10 @@ class GenerationService: generation_model.id = gen_id async def runner(gen): + logger.info(f"Starting background generation task for ID: {gen.id}") try: await self.create_generation(gen) + logger.info(f"Background generation task finished for ID: {gen.id}") except Exception: # если генерация уже пошла и упала — пометим FAILED try: @@ -125,6 +140,8 @@ class GenerationService: raise async def create_generation(self, generation: Generation): + start_time = datetime.now() + logger.info(f"Processing generation {generation.id}. Character ID: {generation.linked_character_id}") # 2. Получаем ассеты-референсы (если они есть) reference_assets: List[Asset] = [] @@ -146,27 +163,47 @@ class GenerationService: if asset.data is not None and asset.type == AssetType.IMAGE ) generation_prompt+=f"PROMPT: {generation.prompt}" + logger.info(f"Final generation prompt assembled. Length: {len(generation_prompt)}. Media count: {len(media_group_bytes)}") - # 3. Запускаем процесс генерации + # 3. Запускаем процесс генерации и симуляцию прогресса + progress_task = asyncio.create_task(self._simulate_progress(generation)) + try: - generated_bytes_list = await generate_image_task( + + # Default to Image Generation (Gemini) + generated_bytes_list, metrics = await generate_image_task( prompt=generation_prompt, # или request.prompt media_group_bytes=media_group_bytes, aspect_ratio=generation.aspect_ratio, # предполагаем поля в request quality=generation.quality, gemini=self.gemini ) + + # Update metrics from API (Common for both) + generation.api_execution_time_seconds = metrics.get("api_execution_time_seconds") + generation.token_usage = metrics.get("token_usage") + except GoogleGenerationException as e: generation.status = GenerationStatus.FAILED - generation.failed_reason = str(e.message) + generation.failed_reason = str(e) generation.updated_at = datetime.now(UTC) await self.dao.generations.update_generation(generation) - raise + raise e except Exception as e: # Тут стоит добавить логирование ошибки logging.error(f"Generation failed: {e}") - # Можно обновить статус генерации на FAILED в БД + generation.status = GenerationStatus.FAILED + generation.failed_reason = str(e) + generation.updated_at = datetime.now(UTC) + await self.dao.generations.update_generation(generation) raise e + finally: + if not progress_task.done(): + progress_task.cancel() + try: + await progress_task + except asyncio.CancelledError: + pass # 4. Сохраняем полученные изображения как новые Ассеты created_assets: List[Asset] = [] @@ -192,6 +229,37 @@ class GenerationService: generation.assets_list = result_ids generation.status = GenerationStatus.DONE + generation.progress = 100 generation.updated_at = datetime.now(UTC) generation.tech_prompt = generation_prompt + + end_time = datetime.now() + generation.execution_time_seconds = (end_time - start_time).total_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") + + + async def _simulate_progress(self, generation: Generation): + """ + Increments progress from 0 to 90 over ~20 seconds. + """ + current_progress = 0 + try: + while current_progress < 90: + await asyncio.sleep(4) + # Random increment between 5 and 15 + increment = random.randint(5, 15) + current_progress = min(current_progress + increment, 90) + + # Fetch latest state (optional, but good practice to avoid overwriting unrelated fields) + # But for simplicity here we just use the object we have and save it. + # Ideally, we should fetch-update-save or use partial update if DAO supports it. + # Assuming simple update is fine for now. + generation.progress = current_progress + await self.dao.generations.update_generation(generation) + except asyncio.CancelledError: + # Task cancelled, generation finished (or failed) + pass + except Exception as e: + logger.error(f"Error in progress simulation: {e}") diff --git a/main.py b/main.py index de7960a..86e4723 100644 --- a/main.py +++ b/main.py @@ -43,6 +43,7 @@ load_dotenv() # --- КОНФИГУРАЦИЯ --- BOT_TOKEN = os.getenv("BOT_TOKEN") GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") + MONGO_HOST = os.getenv("MONGO_HOST") # Например: mongodb://localhost:27017 DB_NAME = os.getenv("DB_NAME", "my_bot_db") # Имя базы данных ADMIN_ID = int(os.getenv("ADMIN_ID", 0)) @@ -63,6 +64,7 @@ mongo_client = AsyncIOMotorClient(MONGO_HOST) users_repo = UsersRepo(mongo_client) char_repo = CharacterRepo(mongo_client) dao = DAO(mongo_client) # Главный DAO для бота +dao = DAO(mongo_client) # Главный DAO для бота gemini = GoogleAdapter(api_key=GEMINI_API_KEY) generation_service = GenerationService(dao, gemini) @@ -113,6 +115,7 @@ async def lifespan(app: FastAPI): # Инициализируем DAO для ассетов и кладем в state приложения # Теперь в эндпоинтах можно делать request.app.state.assets_dao + app.state.mongo_client = mongo_client app.state.mongo_client = mongo_client app.state.gemini_client = gemini diff --git a/middlewares/__pycache__/album.cpython-313.pyc b/middlewares/__pycache__/album.cpython-313.pyc new file mode 100644 index 0000000..f878b76 Binary files /dev/null and b/middlewares/__pycache__/album.cpython-313.pyc differ diff --git a/middlewares/__pycache__/dao.cpython-313.pyc b/middlewares/__pycache__/dao.cpython-313.pyc new file mode 100644 index 0000000..b586fc3 Binary files /dev/null and b/middlewares/__pycache__/dao.cpython-313.pyc differ diff --git a/models/Generation.py b/models/Generation.py index 9476060..2f1f510 100644 --- a/models/Generation.py +++ b/models/Generation.py @@ -5,7 +5,7 @@ from typing import List, Optional from pydantic import BaseModel, Field from models.Asset import Asset -from models.enums import AspectRatios, Quality +from models.enums import AspectRatios, Quality, GenType class GenerationStatus(str, Enum): @@ -24,5 +24,10 @@ class Generation(BaseModel): tech_prompt: Optional[str] = None assets_list: List[str] result: Optional[str] = None + progress: int = 0 + execution_time_seconds: Optional[float] = None + api_execution_time_seconds: Optional[float] = None + token_usage: Optional[int] = None + created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/__pycache__/Asset.cpython-313.pyc b/models/__pycache__/Asset.cpython-313.pyc new file mode 100644 index 0000000..97a91d8 Binary files /dev/null and b/models/__pycache__/Asset.cpython-313.pyc differ diff --git a/models/__pycache__/Character.cpython-313.pyc b/models/__pycache__/Character.cpython-313.pyc new file mode 100644 index 0000000..a89257a Binary files /dev/null and b/models/__pycache__/Character.cpython-313.pyc differ diff --git a/models/__pycache__/Generation.cpython-313.pyc b/models/__pycache__/Generation.cpython-313.pyc new file mode 100644 index 0000000..c3f8a9f Binary files /dev/null and b/models/__pycache__/Generation.cpython-313.pyc differ diff --git a/models/__pycache__/__init__.cpython-313.pyc b/models/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..4530e55 Binary files /dev/null and b/models/__pycache__/__init__.cpython-313.pyc differ diff --git a/models/__pycache__/enums.cpython-313.pyc b/models/__pycache__/enums.cpython-313.pyc new file mode 100644 index 0000000..254c746 Binary files /dev/null and b/models/__pycache__/enums.cpython-313.pyc differ diff --git a/models/enums.py b/models/enums.py index 099622f..a0fd856 100644 --- a/models/enums.py +++ b/models/enums.py @@ -39,5 +39,5 @@ class GenType(str, Enum): def value_type(self) -> str: return { GenType.TEXT: 'Text', - GenType.IMAGE: 'Image' + GenType.IMAGE: 'Image', }[self] diff --git a/repos/__pycache__/assets_repo.cpython-313.pyc b/repos/__pycache__/assets_repo.cpython-313.pyc new file mode 100644 index 0000000..3da7588 Binary files /dev/null and b/repos/__pycache__/assets_repo.cpython-313.pyc differ diff --git a/repos/__pycache__/char_repo.cpython-313.pyc b/repos/__pycache__/char_repo.cpython-313.pyc new file mode 100644 index 0000000..5fd4d21 Binary files /dev/null and b/repos/__pycache__/char_repo.cpython-313.pyc differ diff --git a/repos/__pycache__/dao.cpython-313.pyc b/repos/__pycache__/dao.cpython-313.pyc new file mode 100644 index 0000000..94cbfc8 Binary files /dev/null and b/repos/__pycache__/dao.cpython-313.pyc differ diff --git a/repos/__pycache__/generation_repo.cpython-313.pyc b/repos/__pycache__/generation_repo.cpython-313.pyc new file mode 100644 index 0000000..e21a10d Binary files /dev/null and b/repos/__pycache__/generation_repo.cpython-313.pyc differ diff --git a/repos/__pycache__/user_repo.cpython-313.pyc b/repos/__pycache__/user_repo.cpython-313.pyc index 6632145..52728bf 100644 Binary files a/repos/__pycache__/user_repo.cpython-313.pyc and b/repos/__pycache__/user_repo.cpython-313.pyc differ diff --git a/routers/__pycache__/__init__.cpython-313.pyc b/routers/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..282c6e2 Binary files /dev/null and b/routers/__pycache__/__init__.cpython-313.pyc differ diff --git a/routers/__pycache__/assets_router.cpython-313.pyc b/routers/__pycache__/assets_router.cpython-313.pyc new file mode 100644 index 0000000..f17008a Binary files /dev/null and b/routers/__pycache__/assets_router.cpython-313.pyc differ diff --git a/routers/__pycache__/auth_router.cpython-313.pyc b/routers/__pycache__/auth_router.cpython-313.pyc new file mode 100644 index 0000000..a621a37 Binary files /dev/null and b/routers/__pycache__/auth_router.cpython-313.pyc differ diff --git a/routers/__pycache__/char_router.cpython-313.pyc b/routers/__pycache__/char_router.cpython-313.pyc new file mode 100644 index 0000000..6356e12 Binary files /dev/null and b/routers/__pycache__/char_router.cpython-313.pyc differ diff --git a/routers/__pycache__/gen_router.cpython-313.pyc b/routers/__pycache__/gen_router.cpython-313.pyc new file mode 100644 index 0000000..ae783b4 Binary files /dev/null and b/routers/__pycache__/gen_router.cpython-313.pyc differ diff --git a/routers/gen_router.py b/routers/gen_router.py index 38d76e9..fdcfe6a 100644 --- a/routers/gen_router.py +++ b/routers/gen_router.py @@ -236,9 +236,6 @@ async def handle_album( for msg in album: if msg.photo: file_ids.append(msg.photo[-1].file_id) - elif msg.video: - # Если нужно, можно добавить обработку видео (пока пропускаем) - pass await message.answer(f"📥 Принято {len(album)} файлов. Начинаю генерацию...") wait_msg = await message.answer("🎨 Генерирую...")