diff --git a/api/__init__.py b/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/endpoints/__init__.py b/api/endpoints/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/endpoints/assets.py b/api/endpoints/assets.py new file mode 100644 index 0000000..f00fa44 --- /dev/null +++ b/api/endpoints/assets.py @@ -0,0 +1,29 @@ +from fastapi import APIRouter +from fastapi.openapi.models import MediaType +from starlette.exceptions import HTTPException +from starlette.requests import Request +from starlette.responses import Response, JSONResponse + +from repos.dao import DAO + +router = APIRouter(prefix="/api/assets", tags=["Assets"]) + +@router.get("/{asset_id}") +async def get_asset(asset_id: str, request: Request) -> Response: + dao = request.app.state.dao + asset = await dao.assets.get_asset(asset_id) + # 2. Проверка на существование + if not asset: + raise HTTPException(status_code=404, detail="Asset not found") + return Response(content=asset.data, media_type="image/png") + + +@router.get("") +async def get_assets(request: Request) -> JSONResponse: + dao: DAO = request.app.state.dao + assets = await dao.assets.get_assets() + assets_links = [] + for asset in assets: + assets_links.append("/api/assets/{}".format(asset.id)) + return JSONResponse(content=assets_links) + diff --git a/api/endpoints/character_router.py b/api/endpoints/character_router.py new file mode 100644 index 0000000..d7c54b6 --- /dev/null +++ b/api/endpoints/character_router.py @@ -0,0 +1,35 @@ +from typing import List + +from fastapi import APIRouter +from starlette.exceptions import HTTPException +from starlette.requests import Request + +from models.Asset import Asset +from models.Character import Character +from repos.dao import DAO + +router = APIRouter(prefix="/api/characters", tags=["Characters"]) + + +@router.get("/", response_model=List[Character]) +async def get_characters(request: Request) -> List[Character]: + dao: DAO = request.app.state.dao + characters = await dao.chars.get_all_characters() + return characters + + +@router.get("/{character_id}/assets", response_model=List[Asset]) +async def get_character_assets(character_id: str, request: Request) -> List[Asset]: + dao: DAO = request.app.state.dao + character = await dao.chars.get_character(character_id) + if character is None: + raise HTTPException(status_code=404, detail="Character not found") + return await dao.assets.get_assets_by_char_id(character_id) + + +@router.get("/{character_id}", response_model=Character) +async def get_character_by_id(character_id: int, request: Request) -> Character: + dao: DAO = request.app.state.dao + character = await dao.chars.get_character_by_id(character_id) + return character + diff --git a/main.py b/main.py index c63e858..e6c2781 100644 --- a/main.py +++ b/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" - "Диалоги не поддерживаются!!!! Каждое новое сообщение - новый диалог") + await message.answer("ℹ️ Справка:\n\n" + "📝 Текст: Просто отправь промпт.\n" + "🎨 Фото: /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" - "Диалоги не поддерживаются!!!! Каждое новое сообщение - новый диалог" - ) + "Напиши мне, что нужно сгенерировать, или используй /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 \ No newline at end of file diff --git a/repos/assets_repo.py b/repos/assets_repo.py index 7a284b6..e1b693e 100644 --- a/repos/assets_repo.py +++ b/repos/assets_repo.py @@ -15,8 +15,8 @@ class AssetsRepo: asset.id = res.inserted_id return asset - async def get_assets(self, limit: int = 10, offset: int = 0) -> List[Asset]: - res = await self.collection.find({},{"data":0}).sort("created_at", -1).skip(offset).limit(limit).to_list(None) + async def get_assets(self, limit: int = 10, offset: int = 0) -> List[Asset]: + res = await self.collection.find({}, {"data": 0}).sort("created_at", -1).skip(offset).limit(limit).to_list(None) assets = [] for doc in res: # Конвертируем ObjectId в строку и кладем в поле id @@ -28,12 +28,26 @@ class AssetsRepo: return assets async def get_asset(self, asset_id: str, with_data: bool = True) -> Asset: - res = await self.collection.find_one({"_id": ObjectId(asset_id)}, {"data": 0 if not with_data else 1}) + res = await self.collection.find_one({"_id": ObjectId(asset_id)}, + {"_id": 1, "name": 1, "type": 1, "tg_doc_file_id": 1, + "data": 0 if not with_data else 1}) res["id"] = str(res.pop("_id")) return Asset(**res) - async def update_asset(self, asset_id: str, asset: Asset): if not asset.id: raise Exception(f"Asset ID not found: {asset_id}") await self.collection.update_one({"_id": ObjectId(asset_id)}, {"$set": asset.model_dump()}) + + async def set_tg_photo_file_id(self, asset_id: str, tg_photo_file_id: str): + await self.collection.update_one({"_id": ObjectId(asset_id)}, {"$set": {"tg_photo_file_id": tg_photo_file_id}}) + + async def get_assets_by_char_id(self, character_id: str, limit: int = 10, offset: int = 0) -> List[Asset]: + docs = await self.collection.find({"linked_char_id": character_id}, + {"data": 0}, sort=[("created_at", -1)]).limit(limit).skip(offset).to_list( + None) + assets = [] + for doc in docs: + doc["id"] = str(doc.pop("_id")) + assets.append(Asset(**doc)) + return assets diff --git a/requirements.txt b/requirements.txt index c0c6420..919c804 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,15 +3,18 @@ aiogram==3.24.0 aiohappyeyeballs==2.6.1 aiohttp==3.11.18 aiosignal==1.4.0 +annotated-doc==0.0.4 annotated-types==0.7.0 anyio==4.12.1 attrs==25.4.0 certifi==2026.1.4 cffi==2.0.0 charset-normalizer==3.4.4 +click==8.3.1 cryptography==46.0.4 distro==1.9.0 dnspython==2.8.0 +fastapi==0.128.0 frozenlist==1.8.0 google-auth==2.48.0 google-genai==1.61.0 @@ -34,8 +37,10 @@ python-dotenv==1.2.1 requests==2.32.5 rsa==4.9.1 sniffio==1.3.1 +starlette==0.50.0 tenacity==9.1.2 typing_extensions==4.15.0 urllib3==2.6.3 +uvicorn==0.40.0 websockets==15.0.1 yarl==1.22.0 diff --git a/routers/assets_router.py b/routers/assets_router.py index e773e8e..47e3443 100644 --- a/routers/assets_router.py +++ b/routers/assets_router.py @@ -18,18 +18,18 @@ async def assets_command(msg: Message, dao: DAO): media_group.append(InputMediaPhoto(media=asset.tg_photo_file_id)) elif asset.tg_doc_file_id: asset_full_info = await dao.assets.get_asset(asset.id) + asset = asset_full_info media_group.append(InputMediaPhoto(media=BufferedInputFile(asset_full_info.data, asset_full_info.name))) else: continue keyboard.append(InlineKeyboardButton(text=F"{index + 1}", callback_data=f"asset_doc_{asset.id}")) mg = await msg.answer_media_group(media_group, - reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboard])) + reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboard])) await msg.answer("Для запроса документов выбери фото и на кнопку ниже:", reply_markup=InlineKeyboardMarkup(inline_keyboard=[keyboard])) for media_index, media in enumerate(mg): if assets[media_index].tg_photo_file_id is None: - assets[media_index].tg_photo_file_id = media.photo[-1].file_id - await dao.assets.update_asset(assets[media_index].id, assets[media_index]) + await dao.assets.set_tg_photo_file_id(assets[media_index].id, media.photo[-1].file_id) @router.callback_query(F.data.startswith("asset_doc_"))