feat: Add optional telegram_id field to Generation and GenerationRequest models.
This commit is contained in:
Binary file not shown.
Binary file not shown.
@@ -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)
|
||||
@@ -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]
|
||||
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -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):
|
||||
"""
|
||||
|
||||
3
main.py
3
main.py
@@ -66,7 +66,7 @@ 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)
|
||||
generation_service = GenerationService(dao, gemini, bot)
|
||||
|
||||
# Dispatcher
|
||||
dp = Dispatcher(storage=MongoStorage(mongo_client, db_name=DB_NAME))
|
||||
@@ -118,6 +118,7 @@ async def lifespan(app: FastAPI):
|
||||
app.state.mongo_client = mongo_client
|
||||
app.state.mongo_client = mongo_client
|
||||
app.state.gemini_client = gemini
|
||||
app.state.bot = bot
|
||||
|
||||
print("✅ DB & DAO initialized")
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ class Generation(BaseModel):
|
||||
status: GenerationStatus = GenerationStatus.RUNNING
|
||||
failed_reason: Optional[str] = None
|
||||
linked_character_id: Optional[str] = None
|
||||
telegram_id: Optional[int] = None
|
||||
aspect_ratio: AspectRatios
|
||||
quality: Quality
|
||||
prompt: str
|
||||
|
||||
Binary file not shown.
Reference in New Issue
Block a user