310 lines
12 KiB
Python
310 lines
12 KiB
Python
import asyncio
|
||
import logging
|
||
from contextlib import asynccontextmanager
|
||
|
||
from aiogram import Bot, Dispatcher, Router, F
|
||
from aiogram.client.default import DefaultBotProperties
|
||
from aiogram.enums import ParseMode
|
||
from aiogram.filters import CommandStart, Command
|
||
from aiogram.types import Message
|
||
from aiogram.fsm.storage.mongo import MongoStorage
|
||
from fastapi import FastAPI
|
||
from motor.motor_asyncio import AsyncIOMotorClient
|
||
from prometheus_client import Info
|
||
from starlette.middleware.cors import CORSMiddleware
|
||
from prometheus_fastapi_instrumentator import Instrumentator
|
||
|
||
# --- ИМПОРТЫ ПРОЕКТА ---
|
||
from config import settings
|
||
from adapters.google_adapter import GoogleAdapter
|
||
from adapters.s3_adapter import S3Adapter
|
||
from api.service.generation_service import GenerationService
|
||
from api.service.album_service import AlbumService
|
||
from middlewares.album import AlbumMiddleware
|
||
from middlewares.auth import AuthMiddleware
|
||
from middlewares.dao import DaoMiddleware
|
||
from scheduler.daily_scheduler import DailyScheduler
|
||
from scheduler.telegram_admin_handler import create_daily_scheduler_router
|
||
|
||
# Репозитории и DAO
|
||
from repos.char_repo import CharacterRepo
|
||
from repos.user_repo import UsersRepo
|
||
from repos.dao import DAO
|
||
|
||
|
||
# Роутеры
|
||
from routers.auth_router import router as auth_router
|
||
from routers.gen_router import router as gen_router
|
||
from routers.char_router import router as char_router
|
||
from routers.assets_router import router as assets_router # Роутер бота для ассетов
|
||
from api.endpoints.assets_router import router as api_assets_router # Роутер FastAPI
|
||
from api.endpoints.character_router import router as api_char_router # Роутер FastAPI
|
||
from api.endpoints.generation_router import router as api_gen_router
|
||
from api.endpoints.auth import router as api_auth_router
|
||
from api.endpoints.admin import router as api_admin_router
|
||
from api.endpoints.album_router import router as api_album_router
|
||
from api.endpoints.project_router import router as project_api_router
|
||
from api.endpoints.idea_router import router as idea_api_router
|
||
from api.endpoints.post_router import router as post_api_router
|
||
from api.endpoints.environment_router import router as environment_api_router
|
||
from api.endpoints.inspiration_router import router as inspiration_api_router
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# --- КОНФИГУРАЦИЯ ---
|
||
# Настройки теперь берутся из config.py
|
||
BOT_TOKEN = settings.BOT_TOKEN
|
||
GEMINI_API_KEY = settings.GEMINI_API_KEY
|
||
|
||
MONGO_HOST = settings.MONGO_HOST
|
||
DB_NAME = settings.DB_NAME
|
||
ADMIN_ID = settings.ADMIN_ID
|
||
|
||
|
||
def setup_logging():
|
||
logging.basicConfig(level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(name)s (%(filename)s:%(lineno)d): %(message)s",
|
||
force=True)
|
||
|
||
|
||
# --- ИНИЦИАЛИЗАЦИЯ ЗАВИСИМОСТЕЙ ---
|
||
if BOT_TOKEN is None:
|
||
raise ValueError("BOT_TOKEN is not set")
|
||
bot = Bot(token=BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
|
||
|
||
# Клиент БД создаем глобально, чтобы он был доступен и боту (Storage), и API
|
||
mongo_client = AsyncIOMotorClient(MONGO_HOST)
|
||
|
||
# Репозитории
|
||
# Репозитории
|
||
users_repo = UsersRepo(mongo_client)
|
||
char_repo = CharacterRepo(mongo_client)
|
||
|
||
# S3 Adapter
|
||
s3_adapter = S3Adapter(
|
||
endpoint_url=settings.MINIO_ENDPOINT,
|
||
aws_access_key_id=settings.MINIO_ACCESS_KEY,
|
||
aws_secret_access_key=settings.MINIO_SECRET_KEY,
|
||
bucket_name=settings.MINIO_BUCKET
|
||
)
|
||
|
||
dao = DAO(mongo_client, s3_adapter) # Главный DAO для бота
|
||
if GEMINI_API_KEY is None:
|
||
raise ValueError("GEMINI_API_KEY is not set")
|
||
gemini = GoogleAdapter(api_key=GEMINI_API_KEY)
|
||
if bot is None:
|
||
raise ValueError("bot is not set")
|
||
generation_service = GenerationService(dao=dao, gemini=gemini, s3_adapter=s3_adapter, bot=bot)
|
||
album_service = AlbumService(dao)
|
||
|
||
# Dispatcher
|
||
dp = Dispatcher(storage=MongoStorage(mongo_client, db_name=DB_NAME))
|
||
|
||
# Внедрение зависимостей (глобально для бота)
|
||
dp["repo"] = users_repo
|
||
dp["admin_id"] = ADMIN_ID
|
||
dp["gemini"] = gemini
|
||
|
||
# --- НАСТРОЙКА РОУТЕРОВ БОТА ---
|
||
|
||
# 1. Роутеры без мидлварей (например, auth)
|
||
dp.include_router(auth_router)
|
||
|
||
# 2. Основные роутеры (daily_scheduler router добавляется в lifespan)
|
||
main_router = Router()
|
||
dp.include_router(main_router)
|
||
dp.include_router(assets_router)
|
||
dp.include_router(char_router)
|
||
dp.include_router(gen_router)
|
||
|
||
# --- НАСТРОЙКА MIDDLEWARES БОТА ---
|
||
|
||
# DaoMiddleware прокидывает объект 'dao' во все хендлеры
|
||
dp.update.middleware(DaoMiddleware(dao=dao))
|
||
|
||
# AuthMiddleware проверяет права доступа
|
||
main_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
|
||
gen_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
|
||
assets_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
|
||
|
||
# AlbumMiddleware для обработки групп фото
|
||
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
|
||
|
||
|
||
async def start_scheduler(service: GenerationService):
|
||
while True:
|
||
try:
|
||
logger.info("Running scheduler for stacked generation killing")
|
||
await service.cleanup_stale_generations()
|
||
await service.cleanup_old_data(days=14)
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception as e:
|
||
logger.error(f"Scheduler error: {e}")
|
||
await asyncio.sleep(60) # Check every 60 seconds
|
||
|
||
|
||
def _build_daily_scheduler() -> DailyScheduler:
|
||
"""Construct DailyScheduler; MetaAdapter is optional (needs env vars)."""
|
||
meta_adapter = None
|
||
if settings.META_ACCESS_TOKEN and settings.META_INSTAGRAM_ACCOUNT_ID:
|
||
from adapters.meta_adapter import MetaAdapter
|
||
meta_adapter = MetaAdapter(
|
||
access_token=settings.META_ACCESS_TOKEN,
|
||
instagram_account_id=settings.META_INSTAGRAM_ACCOUNT_ID,
|
||
)
|
||
logger.info("MetaAdapter initialized")
|
||
else:
|
||
logger.warning("META_ACCESS_TOKEN / META_INSTAGRAM_ACCOUNT_ID not set — Instagram publishing disabled")
|
||
|
||
if not settings.SCHEDULER_CHARACTER_ID:
|
||
logger.warning("SCHEDULER_CHARACTER_ID not set — daily scheduler will error at runtime")
|
||
|
||
return DailyScheduler(
|
||
dao=dao,
|
||
gemini=gemini,
|
||
s3_adapter=s3_adapter,
|
||
generation_service=generation_service,
|
||
bot=bot,
|
||
admin_id=ADMIN_ID,
|
||
character_id=settings.SCHEDULER_CHARACTER_ID or "",
|
||
meta_adapter=meta_adapter,
|
||
)
|
||
|
||
# --- LIFESPAN (Запуск FastAPI + Bot) ---
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
# --- STARTUP ---
|
||
setup_logging()
|
||
print("🚀 Starting up...")
|
||
|
||
# 1. Настройка DAO для FastAPI
|
||
# Используем уже созданный mongo_client
|
||
db = mongo_client[DB_NAME]
|
||
|
||
# Инициализируем DAO для ассетов и кладем в state приложения
|
||
# Теперь в эндпоинтах можно делать request.app.state.assets_dao
|
||
|
||
app.state.mongo_client = mongo_client
|
||
app.state.gemini_client = gemini
|
||
app.state.bot = bot
|
||
app.state.s3_adapter = s3_adapter
|
||
app.state.album_service = album_service
|
||
app.state.users_repo = users_repo # Добавляем репозиторий в state
|
||
|
||
print("✅ DB & DAO initialized")
|
||
|
||
# 2. Инициализация и регистрация daily_scheduler роутера
|
||
daily_scheduler = _build_daily_scheduler()
|
||
dp.include_router(create_daily_scheduler_router(daily_scheduler))
|
||
print("📅 Daily scheduler router registered")
|
||
|
||
# 3. ЗАПУСК БОТА (в фоне)
|
||
# handle_signals=False — бот не перехватывает сигналы остановки у uvicorn
|
||
polling_task = asyncio.create_task(
|
||
dp.start_polling(bot, handle_signals=False)
|
||
)
|
||
print("🤖 Bot polling started")
|
||
|
||
# 4. ЗАПУСК ШЕДУЛЕРОВ
|
||
scheduler_task = asyncio.create_task(start_scheduler(generation_service))
|
||
daily_scheduler_task = asyncio.create_task(daily_scheduler.run_loop())
|
||
print("⏰ Schedulers started")
|
||
|
||
yield
|
||
|
||
# --- SHUTDOWN ---
|
||
print("🛑 Shutting down...")
|
||
|
||
# Останавливаем все фоновые задачи
|
||
for task, name in [
|
||
(polling_task, "Bot polling"),
|
||
(scheduler_task, "Stale-gen scheduler"),
|
||
(daily_scheduler_task, "Daily scheduler"),
|
||
]:
|
||
task.cancel()
|
||
try:
|
||
await task
|
||
except asyncio.CancelledError:
|
||
print(f"⏹ {name} stopped")
|
||
|
||
# 4. Отключение БД
|
||
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается
|
||
# mongo_client.close()
|
||
print("🛑 DB Connection closed")
|
||
|
||
|
||
# --- НАСТРОЙКА FASTAPI ---
|
||
app = FastAPI(title="Assets API", lifespan=lifespan)
|
||
|
||
# CORS
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# Подключаем роутеры API
|
||
app.include_router(api_auth_router)
|
||
app.include_router(api_admin_router)
|
||
app.include_router(api_assets_router)
|
||
app.include_router(api_char_router)
|
||
app.include_router(api_gen_router)
|
||
app.include_router(api_album_router)
|
||
app.include_router(project_api_router)
|
||
app.include_router(idea_api_router)
|
||
app.include_router(post_api_router)
|
||
app.include_router(environment_api_router)
|
||
app.include_router(inspiration_api_router)
|
||
|
||
# Prometheus Metrics (Instrument after all routers are added)
|
||
Instrumentator(
|
||
should_group_status_codes=False, # 200/201/204 отдельно (по желанию)
|
||
should_ignore_untemplated=False, # НЕ игнорировать "сырые" пути
|
||
# should_group_untemplated=False, # (опционально) не схлопывать untemplated в "none"
|
||
).instrument(
|
||
app,
|
||
metric_namespace="ai_bot",
|
||
).expose(app, endpoint="/metrics", include_in_schema=False)
|
||
app_info = Info("fastapi_app_info", "FastAPI application info")
|
||
app_info.info({"app_name": "ai-bot"})
|
||
|
||
|
||
# --- ХЕНДЛЕРЫ БОТА (Main Router) ---
|
||
@main_router.message(Command("help"))
|
||
async def show_help(message: Message) -> None:
|
||
await message.answer("ℹ️ <b>Справка:</b>\n\n"
|
||
"📝 <b>Текст:</b> Просто отправь промпт.\n"
|
||
"🎨 <b>Фото:</b> /image {промпт} (или прикрепи фото с подписью).\n\n"
|
||
"⚠️ Диалоги не сохраняются (каждое сообщение — новый запрос).")
|
||
|
||
|
||
@main_router.message(CommandStart())
|
||
async def cmd_start(message: Message):
|
||
await message.answer("👋 Привет! Я готов к работе.\n\n"
|
||
"Напиши мне, что нужно сгенерировать, или используй /help.")
|
||
|
||
|
||
# --- ЗАПУСК ---
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
setup_logging()
|
||
|
||
async def main():
|
||
# Создаем конфигурацию uvicorn вручную
|
||
# loop="asyncio" заставляет использовать стандартный цикл
|
||
config = uvicorn.Config(app, host="0.0.0.0", port=8090, loop="asyncio", timeout_keep_alive=120)
|
||
server = uvicorn.Server(config)
|
||
|
||
# Запускаем сервер (lifespan запустится внутри)
|
||
await server.serve()
|
||
|
||
|
||
try:
|
||
# Сами запускаем цикл, контролируя аргументы
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
# Корректно обрабатываем выход
|
||
pass |