+ api
This commit is contained in:
165
main.py
165
main.py
@@ -1,107 +1,194 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
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, CommandObject
|
||||
from aiogram.types import Message, BufferedInputFile
|
||||
from aiogram.filters import CommandStart, Command
|
||||
from aiogram.types import Message
|
||||
from aiogram.fsm.storage.mongo import MongoStorage
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import FastAPI
|
||||
from motor.motor_asyncio import AsyncIOMotorClient
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
||||
# Импорты
|
||||
# --- ИМПОРТЫ ПРОЕКТА ---
|
||||
from adapters.google_adapter import GoogleAdapter
|
||||
from middlewares.album import AlbumMiddleware
|
||||
from middlewares.auth import AuthMiddleware
|
||||
from middlewares.dao import DaoMiddleware
|
||||
|
||||
# Репозитории и DAO
|
||||
from repos.char_repo import CharacterRepo
|
||||
from repos.dao import DAO
|
||||
from repos.user_repo import UsersRepo
|
||||
from routers import char_router
|
||||
# ВАЖНО: Импортируем роутер с логикой кнопок, а не создаем пустой
|
||||
from repos.dao import DAO
|
||||
# Предполагаю, что AssetsDAO лежит тут или в repos.assets_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 routers.assets_router import router as assets_router # Роутер бота для ассетов
|
||||
from api.endpoints.assets import router as api_assets_router # Роутер FastAPI
|
||||
from api.endpoints.character_router import router as api_char_router # Роутер FastAPI
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Настройки
|
||||
# --- КОНФИГУРАЦИЯ ---
|
||||
BOT_TOKEN = os.getenv("BOT_TOKEN")
|
||||
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
||||
MONGO_HOST = os.getenv("MONGO_HOST")
|
||||
ADMIN_ID = int(os.getenv("ADMIN_ID")) # Сразу преобразуем в int
|
||||
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))
|
||||
|
||||
# Инициализация
|
||||
|
||||
def setup_logging():
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
|
||||
|
||||
|
||||
# --- ИНИЦИАЛИЗАЦИЯ ЗАВИСИМОСТЕЙ ---
|
||||
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)
|
||||
dao = DAO(mongo_client) # Главный DAO для бота
|
||||
|
||||
# Dispatcher
|
||||
# Если MongoStorage пока не настроен на authSource=admin, можно временно убрать storage=...
|
||||
dp = Dispatcher(storage=MongoStorage(mongo_client))
|
||||
dp = Dispatcher(storage=MongoStorage(mongo_client, db_name=DB_NAME))
|
||||
|
||||
# ВНЕДРЕНИЕ ЗАВИСИМОСТЕЙ (чтобы они были доступны в хендлерах)
|
||||
# Внедрение зависимостей (глобально для бота)
|
||||
dp["repo"] = users_repo
|
||||
dp["admin_id"] = ADMIN_ID
|
||||
dp["gemini"] = GoogleAdapter(api_key=GEMINI_API_KEY) # Инициализируем тут
|
||||
dp["gemini"] = GoogleAdapter(api_key=GEMINI_API_KEY)
|
||||
|
||||
# РОУТИНГ
|
||||
# --- НАСТРОЙКА РОУТЕРОВ БОТА ---
|
||||
|
||||
# 1. Роутер авторизации (кнопки) - ПОДКЛЮЧАЕМ ПЕРВЫМ И БЕЗ МИДЛВАРИ
|
||||
# 1. Роутеры без мидлварей (например, auth)
|
||||
dp.include_router(auth_router)
|
||||
|
||||
# 2. Основные роутеры
|
||||
main_router = Router()
|
||||
dp.include_router(main_router)
|
||||
dp.include_router(assets_router)
|
||||
dp.include_router(char_router)
|
||||
dp.include_router(gen_router)
|
||||
|
||||
# 2. Основной роутер (чат с ботом)
|
||||
# --- НАСТРОЙКА 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))
|
||||
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
|
||||
assets_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_ID))
|
||||
dp.update.middleware(DaoMiddleware(dao=DAO(client=mongo_client)))
|
||||
|
||||
# AlbumMiddleware для обработки групп фото
|
||||
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
|
||||
|
||||
|
||||
# --- LIFESPAN (Запуск FastAPI + Bot) ---
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# --- STARTUP ---
|
||||
print("🚀 Starting up...")
|
||||
|
||||
def setup_logging() -> None:
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
|
||||
# 1. Настройка DAO для FastAPI
|
||||
# Используем уже созданный mongo_client
|
||||
db = mongo_client[DB_NAME]
|
||||
|
||||
# Инициализируем DAO для ассетов и кладем в state приложения
|
||||
# Теперь в эндпоинтах можно делать request.app.state.assets_dao
|
||||
app.state.dao = dao
|
||||
|
||||
print("✅ DB & DAO initialized")
|
||||
|
||||
# 2. ЗАПУСК БОТА (в фоне)
|
||||
# Важно: handle_signals=False, чтобы бот не перехватывал сигналы остановки у uvicorn
|
||||
# Мы НЕ передаем сюда dao=..., так как он уже подключен через Middleware выше
|
||||
polling_task = asyncio.create_task(
|
||||
dp.start_polling(bot, handle_signals=False)
|
||||
)
|
||||
print("🤖 Bot polling started")
|
||||
|
||||
yield
|
||||
|
||||
# --- SHUTDOWN ---
|
||||
print("🛑 Shutting down...")
|
||||
|
||||
# 3. Остановка бота
|
||||
polling_task.cancel()
|
||||
try:
|
||||
await polling_task
|
||||
except asyncio.CancelledError:
|
||||
print("🤖 Bot polling stopped")
|
||||
|
||||
# 4. Отключение БД
|
||||
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается
|
||||
# mongo_client.close()
|
||||
print("🛑 DB Connection closed")
|
||||
|
||||
|
||||
# --- ХЕНДЛЕРЫ ОСНОВНОГО РОУТЕРА ---
|
||||
# Переносим их прямо сюда или в отдельный файл routers/chat_router.py
|
||||
# --- НАСТРОЙКА FASTAPI ---
|
||||
app = FastAPI(title="Assets API", lifespan=lifespan)
|
||||
|
||||
# CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Подключаем роутер API
|
||||
app.include_router(api_assets_router)
|
||||
app.include_router(api_char_router)
|
||||
|
||||
|
||||
# --- ХЕНДЛЕРЫ БОТА (Main Router) ---
|
||||
@main_router.message(Command("help"))
|
||||
async def show_help(message: Message) -> None:
|
||||
await message.answer("Для того, чтобы обратиться для текстовой генерации - просто отправь промпт.\n\n"
|
||||
"Для генерации фото - /image {prompt}\n\n"
|
||||
"Можно отправить фото и команду /image {prompt}\n\n"
|
||||
"Диалоги не поддерживаются!!!! <b>Каждое новое сообщение - новый диалог</b>")
|
||||
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"
|
||||
"Для того, чтобы обратиться для текстовой генерации - просто отправь промпт.\n\n"
|
||||
"Для генерации фото - /image {prompt}\n\n"
|
||||
"Можно отправить фото и команду /image {prompt}\n\n"
|
||||
"Диалоги не поддерживаются!!!! <b>Каждое новое сообщение - новый диалог</b>"
|
||||
)
|
||||
"Напиши мне, что нужно сгенерировать, или используй /help.")
|
||||
|
||||
|
||||
# --- ЗАПУСК ---
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
setup_logging()
|
||||
|
||||
|
||||
async def main():
|
||||
# Создаем конфигурацию uvicorn вручную
|
||||
# loop="asyncio" заставляет использовать стандартный цикл
|
||||
config = uvicorn.Config(app, host="0.0.0.0", port=8000, loop="asyncio")
|
||||
server = uvicorn.Server(config)
|
||||
|
||||
# Запускаем сервер (lifespan запустится внутри)
|
||||
await server.serve()
|
||||
|
||||
|
||||
try:
|
||||
asyncio.run(dp.start_polling(bot))
|
||||
# Сами запускаем цикл, контролируя аргументы
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
print("Bot stopped")
|
||||
# Корректно обрабатываем выход
|
||||
pass
|
||||
Reference in New Issue
Block a user