feat: Add optional telegram_id field to Generation and GenerationRequest models.
This commit is contained in:
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):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user