feat: Add optional telegram_id field to Generation and GenerationRequest models.

This commit is contained in:
xds
2026-02-05 21:10:50 +03:00
parent 463e73fa1e
commit d4682b1418
10 changed files with 29 additions and 3 deletions

View File

@@ -9,6 +9,8 @@ from repos.dao import DAO
# ... ваши импорты ...
from aiogram import Bot
# Провайдеры "сырых" клиентов из состояния приложения
def get_mongo_client(request: Request) -> AsyncIOMotorClient:
return request.app.state.mongo_client
@@ -16,6 +18,9 @@ def get_mongo_client(request: Request) -> AsyncIOMotorClient:
def get_gemini_client(request: Request) -> GoogleAdapter:
return request.app.state.gemini_client
def get_bot_client(request: Request) -> Bot:
return request.app.state.bot
# Провайдер DAO (собирается из mongo_client)
def get_dao(mongo_client: AsyncIOMotorClient = Depends(get_mongo_client)) -> DAO:
# FastAPI кэширует результат Depends в рамках одного запроса,
@@ -26,5 +31,6 @@ def get_dao(mongo_client: AsyncIOMotorClient = Depends(get_mongo_client)) -> DAO
def get_generation_service(
dao: DAO = Depends(get_dao),
gemini: GoogleAdapter = Depends(get_gemini_client),
bot: Bot = Depends(get_bot_client),
) -> GenerationService:
return GenerationService(dao, gemini)
return GenerationService(dao, gemini, bot)

View File

@@ -13,6 +13,7 @@ class GenerationRequest(BaseModel):
aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN
quality: Quality = Quality.ONEK
prompt: str
telegram_id: Optional[int] = None
assets_list: List[str]

View File

@@ -5,6 +5,8 @@ from datetime import datetime, UTC
from typing import List, Optional, Tuple, Any, Dict
from io import BytesIO
from aiogram import Bot
from aiogram.types import BufferedInputFile
from adapters.Exception import GoogleGenerationException
from adapters.google_adapter import GoogleAdapter
from api.models.GenerationRequest import GenerationRequest, GenerationResponse
@@ -58,9 +60,10 @@ async def generate_image_task(
return images_bytes, metrics
class GenerationService:
def __init__(self, dao: DAO, gemini: GoogleAdapter):
def __init__(self, dao: DAO, gemini: GoogleAdapter, bot: Optional[Bot] = None):
self.dao = dao
self.gemini = gemini
self.bot = bot
async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str:
@@ -249,6 +252,20 @@ class GenerationService:
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")
# 6. Send to Telegram if telegram_id is provided
if generation.telegram_id and self.bot:
try:
for asset in created_assets:
if asset.data:
await self.bot.send_photo(
chat_id=generation.telegram_id,
photo=BufferedInputFile(asset.data, filename=f"{asset.name}.jpg"),
caption=f"Generated from prompt: {generation.prompt[:100]}..."
)
logger.info(f"Sent {len(created_assets)} assets to Telegram ID: {generation.telegram_id}")
except Exception as e:
logger.error(f"Failed to send assets to Telegram ID {generation.telegram_id}: {e}")
async def _simulate_progress(self, generation: Generation):
"""