Compare commits
26 Commits
d1f67c773f
...
enviroment
| Author | SHA1 | Date | |
|---|---|---|---|
| 5aa6391dc8 | |||
| ffb0463fe0 | |||
| dd0f8a1cb6 | |||
| 4af5134726 | |||
| 7488665d04 | |||
| ecc88aca62 | |||
| 70f50170fc | |||
| f4207fc4c1 | |||
| c50d2c8ad9 | |||
| 4586daac38 | |||
| 198ac44960 | |||
| d820d9145b | |||
| c93e577bcf | |||
| c5d4849bff | |||
| 9abfbef871 | |||
| 68a3f529cb | |||
| e2c050515d | |||
| 5e7dc19bf3 | |||
| 97483b7030 | |||
| 2d3da59de9 | |||
| 279cb5c6f6 | |||
| 30138bab38 | |||
| 977cab92f8 | |||
| dcab238d3e | |||
| 9d2e4e47de | |||
| c6142715d9 |
15
.gitignore
vendored
15
.gitignore
vendored
@@ -9,3 +9,18 @@ minio_backup.tar.gz
|
|||||||
.idea
|
.idea
|
||||||
.venv
|
.venv
|
||||||
.vscode
|
.vscode
|
||||||
|
.vscode/launch.json
|
||||||
|
middlewares/__pycache__/
|
||||||
|
middlewares/*.pyc
|
||||||
|
api/__pycache__/
|
||||||
|
api/*.pyc
|
||||||
|
repos/__pycache__/
|
||||||
|
repos/*.pyc
|
||||||
|
adapters/__pycache__/
|
||||||
|
adapters/*.pyc
|
||||||
|
services/__pycache__/
|
||||||
|
services/*.pyc
|
||||||
|
utils/__pycache__/
|
||||||
|
utils/*.pyc
|
||||||
|
.vscode/launch.json
|
||||||
|
repos/__pycache__/assets_repo.cpython-313.pyc
|
||||||
|
|||||||
27
.vscode/launch.json
vendored
27
.vscode/launch.json
vendored
@@ -7,7 +7,7 @@
|
|||||||
"request": "launch",
|
"request": "launch",
|
||||||
"module": "uvicorn",
|
"module": "uvicorn",
|
||||||
"args": [
|
"args": [
|
||||||
"main:app",
|
"aiws:app",
|
||||||
"--reload",
|
"--reload",
|
||||||
"--port",
|
"--port",
|
||||||
"8090",
|
"8090",
|
||||||
@@ -16,31 +16,6 @@
|
|||||||
],
|
],
|
||||||
"jinja": true,
|
"jinja": true,
|
||||||
"justMyCode": true
|
"justMyCode": true
|
||||||
},
|
|
||||||
{
|
|
||||||
"name": "Python: Current File",
|
|
||||||
"type": "debugpy",
|
|
||||||
"request": "launch",
|
|
||||||
"program": "${file}",
|
|
||||||
"console": "integratedTerminal",
|
|
||||||
"justMyCode": true,
|
|
||||||
"env": {
|
|
||||||
"PYTHONPATH": "${workspaceFolder}"
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"name": "Debug Tests: Current File",
|
|
||||||
"type": "debugpy",
|
|
||||||
"request": "launch",
|
|
||||||
"module": "pytest",
|
|
||||||
"args": [
|
|
||||||
"${file}"
|
|
||||||
],
|
|
||||||
"console": "integratedTerminal",
|
|
||||||
"justMyCode": true,
|
|
||||||
"env": {
|
|
||||||
"PYTHONPATH": "${workspaceFolder}"
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -23,28 +23,30 @@ class GoogleAdapter:
|
|||||||
self.TEXT_MODEL = "gemini-3-pro-preview"
|
self.TEXT_MODEL = "gemini-3-pro-preview"
|
||||||
self.IMAGE_MODEL = "gemini-3-pro-image-preview"
|
self.IMAGE_MODEL = "gemini-3-pro-image-preview"
|
||||||
|
|
||||||
def _prepare_contents(self, prompt: str, images_list: List[bytes] = None) -> list:
|
def _prepare_contents(self, prompt: str, images_list: List[bytes] | None = None) -> tuple:
|
||||||
"""Вспомогательный метод для подготовки контента (текст + картинки)"""
|
"""Вспомогательный метод для подготовки контента (текст + картинки).
|
||||||
contents = [prompt]
|
Returns (contents, opened_images) — caller MUST close opened_images after use."""
|
||||||
|
contents : list [Any]= [prompt]
|
||||||
|
opened_images = []
|
||||||
if images_list:
|
if images_list:
|
||||||
logger.info(f"Preparing content with {len(images_list)} images")
|
logger.info(f"Preparing content with {len(images_list)} images")
|
||||||
for img_bytes in images_list:
|
for img_bytes in images_list:
|
||||||
try:
|
try:
|
||||||
# Gemini API требует PIL Image на входе
|
|
||||||
image = Image.open(io.BytesIO(img_bytes))
|
image = Image.open(io.BytesIO(img_bytes))
|
||||||
contents.append(image)
|
contents.append(image)
|
||||||
|
opened_images.append(image)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing input image: {e}")
|
logger.error(f"Error processing input image: {e}")
|
||||||
else:
|
else:
|
||||||
logger.info("Preparing content with no images")
|
logger.info("Preparing content with no images")
|
||||||
return contents
|
return contents, opened_images
|
||||||
|
|
||||||
def generate_text(self, prompt: str, images_list: List[bytes] = None) -> str:
|
def generate_text(self, prompt: str, images_list: List[bytes] | None = None) -> str:
|
||||||
"""
|
"""
|
||||||
Генерация текста (Чат или Vision).
|
Генерация текста (Чат или Vision).
|
||||||
Возвращает строку с ответом.
|
Возвращает строку с ответом.
|
||||||
"""
|
"""
|
||||||
contents = self._prepare_contents(prompt, images_list)
|
contents, opened_images = self._prepare_contents(prompt, images_list)
|
||||||
logger.info(f"Generating text: {prompt}")
|
logger.info(f"Generating text: {prompt}")
|
||||||
try:
|
try:
|
||||||
response = self.client.models.generate_content(
|
response = self.client.models.generate_content(
|
||||||
@@ -68,14 +70,17 @@ class GoogleAdapter:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Gemini Text API Error: {e}")
|
logger.error(f"Gemini Text API Error: {e}")
|
||||||
raise GoogleGenerationException(f"Gemini Text API Error: {e}")
|
raise GoogleGenerationException(f"Gemini Text API Error: {e}")
|
||||||
|
finally:
|
||||||
|
for img in opened_images:
|
||||||
|
img.close()
|
||||||
|
|
||||||
def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]:
|
def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] | None = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]:
|
||||||
"""
|
"""
|
||||||
Генерация изображений (Text-to-Image или Image-to-Image).
|
Генерация изображений (Text-to-Image или Image-to-Image).
|
||||||
Возвращает список байтовых потоков (готовых к отправке).
|
Возвращает список байтовых потоков (готовых к отправке).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
contents = self._prepare_contents(prompt, images_list)
|
contents, opened_images = self._prepare_contents(prompt, images_list)
|
||||||
logger.info(f"Generating image. Prompt length: {len(prompt)}, Ratio: {aspect_ratio}, Quality: {quality}")
|
logger.info(f"Generating image. Prompt length: {len(prompt)}, Ratio: {aspect_ratio}, Quality: {quality}")
|
||||||
|
|
||||||
start_time = datetime.now()
|
start_time = datetime.now()
|
||||||
@@ -101,8 +106,20 @@ class GoogleAdapter:
|
|||||||
if response.usage_metadata:
|
if response.usage_metadata:
|
||||||
token_usage = response.usage_metadata.total_token_count
|
token_usage = response.usage_metadata.total_token_count
|
||||||
|
|
||||||
if response.parts is None and response.candidates[0].finish_reason is not None:
|
# Check prompt-level block (e.g. PROHIBITED_CONTENT) — no candidates in this case
|
||||||
raise GoogleGenerationException(f"Generation blocked in cause of {response.candidates[0].finish_reason.value}")
|
if response.prompt_feedback and response.prompt_feedback.block_reason:
|
||||||
|
raise GoogleGenerationException(
|
||||||
|
f"Generation blocked at prompt level: {response.prompt_feedback.block_reason.value}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check candidate-level block
|
||||||
|
if response.parts is None:
|
||||||
|
response_reason = (
|
||||||
|
response.candidates[0].finish_reason
|
||||||
|
if response.candidates and len(response.candidates) > 0
|
||||||
|
else "Unknown"
|
||||||
|
)
|
||||||
|
raise GoogleGenerationException(f"Generation blocked: {response_reason}")
|
||||||
|
|
||||||
generated_images = []
|
generated_images = []
|
||||||
|
|
||||||
@@ -113,7 +130,9 @@ class GoogleAdapter:
|
|||||||
try:
|
try:
|
||||||
# 1. Берем сырые байты
|
# 1. Берем сырые байты
|
||||||
raw_data = part.inline_data.data
|
raw_data = part.inline_data.data
|
||||||
byte_arr = io.BytesIO(raw_data)
|
if raw_data is None:
|
||||||
|
raise GoogleGenerationException("Generation returned no data")
|
||||||
|
byte_arr : io.BytesIO = io.BytesIO(raw_data)
|
||||||
|
|
||||||
# 2. Нейминг (формально, для TG)
|
# 2. Нейминг (формально, для TG)
|
||||||
timestamp = datetime.now().timestamp()
|
timestamp = datetime.now().timestamp()
|
||||||
@@ -148,3 +167,7 @@ class GoogleAdapter:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Gemini Image API Error: {e}")
|
logger.error(f"Gemini Image API Error: {e}")
|
||||||
raise GoogleGenerationException(f"Gemini Image API Error: {e}")
|
raise GoogleGenerationException(f"Gemini Image API Error: {e}")
|
||||||
|
finally:
|
||||||
|
for img in opened_images:
|
||||||
|
img.close()
|
||||||
|
del contents
|
||||||
@@ -18,7 +18,7 @@ class S3Adapter:
|
|||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def _get_client(self):
|
async def _get_client(self):
|
||||||
async with self.session.client(
|
async with self.session.client( # type: ignore[reportGeneralTypeIssues]
|
||||||
"s3",
|
"s3",
|
||||||
endpoint_url=self.endpoint_url,
|
endpoint_url=self.endpoint_url,
|
||||||
aws_access_key_id=self.aws_access_key_id,
|
aws_access_key_id=self.aws_access_key_id,
|
||||||
@@ -56,6 +56,23 @@ class S3Adapter:
|
|||||||
print(f"Error downloading from S3: {e}")
|
print(f"Error downloading from S3: {e}")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
async def stream_file(self, object_name: str, chunk_size: int = 65536):
|
||||||
|
"""Streams a file from S3 yielding chunks. Memory-efficient for large files."""
|
||||||
|
try:
|
||||||
|
async with self._get_client() as client:
|
||||||
|
response = await client.get_object(Bucket=self.bucket_name, Key=object_name)
|
||||||
|
# aioboto3 Body is an aiohttp StreamReader wrapper
|
||||||
|
body = response['Body']
|
||||||
|
|
||||||
|
while True:
|
||||||
|
chunk = await body.read(chunk_size)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
yield chunk
|
||||||
|
except ClientError as e:
|
||||||
|
print(f"Error streaming from S3: {e}")
|
||||||
|
return
|
||||||
|
|
||||||
async def delete_file(self, object_name: str):
|
async def delete_file(self, object_name: str):
|
||||||
"""Deletes a file from S3."""
|
"""Deletes a file from S3."""
|
||||||
try:
|
try:
|
||||||
|
|||||||
70
aiws.py
70
aiws.py
@@ -1,6 +1,5 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import os
|
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
from aiogram import Bot, Dispatcher, Router, F
|
from aiogram import Bot, Dispatcher, Router, F
|
||||||
@@ -9,7 +8,6 @@ from aiogram.enums import ParseMode
|
|||||||
from aiogram.filters import CommandStart, Command
|
from aiogram.filters import CommandStart, Command
|
||||||
from aiogram.types import Message
|
from aiogram.types import Message
|
||||||
from aiogram.fsm.storage.mongo import MongoStorage
|
from aiogram.fsm.storage.mongo import MongoStorage
|
||||||
from dotenv import load_dotenv
|
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
from motor.motor_asyncio import AsyncIOMotorClient
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
from prometheus_client import Info
|
from prometheus_client import Info
|
||||||
@@ -17,6 +15,7 @@ from starlette.middleware.cors import CORSMiddleware
|
|||||||
from prometheus_fastapi_instrumentator import Instrumentator
|
from prometheus_fastapi_instrumentator import Instrumentator
|
||||||
|
|
||||||
# --- ИМПОРТЫ ПРОЕКТА ---
|
# --- ИМПОРТЫ ПРОЕКТА ---
|
||||||
|
from config import settings
|
||||||
from adapters.google_adapter import GoogleAdapter
|
from adapters.google_adapter import GoogleAdapter
|
||||||
from adapters.s3_adapter import S3Adapter
|
from adapters.s3_adapter import S3Adapter
|
||||||
from api.service.generation_service import GenerationService
|
from api.service.generation_service import GenerationService
|
||||||
@@ -43,17 +42,20 @@ from api.endpoints.auth import router as api_auth_router
|
|||||||
from api.endpoints.admin import router as api_admin_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.album_router import router as api_album_router
|
||||||
from api.endpoints.project_router import router as project_api_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
|
||||||
|
|
||||||
load_dotenv()
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# --- КОНФИГУРАЦИЯ ---
|
# --- КОНФИГУРАЦИЯ ---
|
||||||
BOT_TOKEN = os.getenv("BOT_TOKEN")
|
# Настройки теперь берутся из config.py
|
||||||
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
BOT_TOKEN = settings.BOT_TOKEN
|
||||||
|
GEMINI_API_KEY = settings.GEMINI_API_KEY
|
||||||
|
|
||||||
MONGO_HOST = os.getenv("MONGO_HOST") # Например: mongodb://localhost:27017
|
MONGO_HOST = settings.MONGO_HOST
|
||||||
DB_NAME = os.getenv("DB_NAME", "my_bot_db") # Имя базы данных
|
DB_NAME = settings.DB_NAME
|
||||||
ADMIN_ID = int(os.getenv("ADMIN_ID", 0))
|
ADMIN_ID = settings.ADMIN_ID
|
||||||
|
|
||||||
|
|
||||||
def setup_logging():
|
def setup_logging():
|
||||||
@@ -63,6 +65,8 @@ def setup_logging():
|
|||||||
|
|
||||||
|
|
||||||
# --- ИНИЦИАЛИЗАЦИЯ ЗАВИСИМОСТЕЙ ---
|
# --- ИНИЦИАЛИЗАЦИЯ ЗАВИСИМОСТЕЙ ---
|
||||||
|
if BOT_TOKEN is None:
|
||||||
|
raise ValueError("BOT_TOKEN is not set")
|
||||||
bot = Bot(token=BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
|
bot = Bot(token=BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
|
||||||
|
|
||||||
# Клиент БД создаем глобально, чтобы он был доступен и боту (Storage), и API
|
# Клиент БД создаем глобально, чтобы он был доступен и боту (Storage), и API
|
||||||
@@ -75,15 +79,19 @@ char_repo = CharacterRepo(mongo_client)
|
|||||||
|
|
||||||
# S3 Adapter
|
# S3 Adapter
|
||||||
s3_adapter = S3Adapter(
|
s3_adapter = S3Adapter(
|
||||||
endpoint_url=os.getenv("MINIO_ENDPOINT", "http://31.59.58.220:9000"),
|
endpoint_url=settings.MINIO_ENDPOINT,
|
||||||
aws_access_key_id=os.getenv("MINIO_ACCESS_KEY", "minioadmin"),
|
aws_access_key_id=settings.MINIO_ACCESS_KEY,
|
||||||
aws_secret_access_key=os.getenv("MINIO_SECRET_KEY", "minioadmin"),
|
aws_secret_access_key=settings.MINIO_SECRET_KEY,
|
||||||
bucket_name=os.getenv("MINIO_BUCKET", "ai-char")
|
bucket_name=settings.MINIO_BUCKET
|
||||||
)
|
)
|
||||||
|
|
||||||
dao = DAO(mongo_client, s3_adapter) # Главный DAO для бота
|
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)
|
gemini = GoogleAdapter(api_key=GEMINI_API_KEY)
|
||||||
generation_service = GenerationService(dao, gemini, bot)
|
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)
|
album_service = AlbumService(dao)
|
||||||
|
|
||||||
# Dispatcher
|
# Dispatcher
|
||||||
@@ -120,6 +128,18 @@ assets_router.message.middleware(AuthMiddleware(repo=users_repo, admin_id=ADMIN_
|
|||||||
gen_router.message.middleware(AlbumMiddleware(latency=0.8))
|
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=2)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
break
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Scheduler error: {e}")
|
||||||
|
await asyncio.sleep(60) # Check every 60 seconds
|
||||||
|
|
||||||
# --- LIFESPAN (Запуск FastAPI + Bot) ---
|
# --- LIFESPAN (Запуск FastAPI + Bot) ---
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
@@ -151,17 +171,28 @@ async def lifespan(app: FastAPI):
|
|||||||
# )
|
# )
|
||||||
# print("🤖 Bot polling started")
|
# print("🤖 Bot polling started")
|
||||||
|
|
||||||
|
# 3. ЗАПУСК ШЕДУЛЕРА
|
||||||
|
scheduler_task = asyncio.create_task(start_scheduler(generation_service))
|
||||||
|
print("⏰ Scheduler started")
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
# --- SHUTDOWN ---
|
# --- SHUTDOWN ---
|
||||||
print("🛑 Shutting down...")
|
print("🛑 Shutting down...")
|
||||||
|
|
||||||
# 3. Остановка бота
|
# 4. Остановка шедулера
|
||||||
polling_task.cancel()
|
scheduler_task.cancel()
|
||||||
try:
|
try:
|
||||||
await polling_task
|
await scheduler_task
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
print("🤖 Bot polling stopped")
|
print("⏰ Scheduler stopped")
|
||||||
|
|
||||||
|
# 3. Остановка бота
|
||||||
|
# polling_task.cancel()
|
||||||
|
# try:
|
||||||
|
# await polling_task
|
||||||
|
# except asyncio.CancelledError:
|
||||||
|
# print("🤖 Bot polling stopped")
|
||||||
|
|
||||||
# 4. Отключение БД
|
# 4. Отключение БД
|
||||||
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается
|
# Обычно Motor закрывать не обязательно при выходе, но хорошим тоном считается
|
||||||
@@ -188,6 +219,9 @@ app.include_router(api_char_router)
|
|||||||
app.include_router(api_gen_router)
|
app.include_router(api_gen_router)
|
||||||
app.include_router(api_album_router)
|
app.include_router(api_album_router)
|
||||||
app.include_router(project_api_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)
|
||||||
|
|
||||||
# Prometheus Metrics (Instrument after all routers are added)
|
# Prometheus Metrics (Instrument after all routers are added)
|
||||||
Instrumentator(
|
Instrumentator(
|
||||||
@@ -226,7 +260,7 @@ if __name__ == "__main__":
|
|||||||
async def main():
|
async def main():
|
||||||
# Создаем конфигурацию uvicorn вручную
|
# Создаем конфигурацию uvicorn вручную
|
||||||
# loop="asyncio" заставляет использовать стандартный цикл
|
# loop="asyncio" заставляет использовать стандартный цикл
|
||||||
config = uvicorn.Config(app, host="0.0.0.0", port=8090, loop="asyncio", timeout_keep_alive=120, env_file=".env.development")
|
config = uvicorn.Config(app, host="0.0.0.0", port=8090, loop="asyncio", timeout_keep_alive=120)
|
||||||
server = uvicorn.Server(config)
|
server = uvicorn.Server(config)
|
||||||
|
|
||||||
# Запускаем сервер (lifespan запустится внутри)
|
# Запускаем сервер (lifespan запустится внутри)
|
||||||
|
|||||||
Binary file not shown.
@@ -5,6 +5,7 @@ from motor.motor_asyncio import AsyncIOMotorClient
|
|||||||
from adapters.google_adapter import GoogleAdapter
|
from adapters.google_adapter import GoogleAdapter
|
||||||
from api.service.generation_service import GenerationService
|
from api.service.generation_service import GenerationService
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
|
from api.service.album_service import AlbumService
|
||||||
|
|
||||||
|
|
||||||
# ... ваши импорты ...
|
# ... ваши импорты ...
|
||||||
@@ -45,7 +46,20 @@ def get_generation_service(
|
|||||||
) -> GenerationService:
|
) -> GenerationService:
|
||||||
return GenerationService(dao, gemini, s3_adapter, bot)
|
return GenerationService(dao, gemini, s3_adapter, bot)
|
||||||
|
|
||||||
|
from api.service.idea_service import IdeaService
|
||||||
|
|
||||||
|
def get_idea_service(dao: DAO = Depends(get_dao)) -> IdeaService:
|
||||||
|
return IdeaService(dao)
|
||||||
|
|
||||||
from fastapi import Header
|
from fastapi import Header
|
||||||
|
|
||||||
async def get_project_id(x_project_id: Optional[str] = Header(None, alias="X-Project-ID")) -> Optional[str]:
|
async def get_project_id(x_project_id: Optional[str] = Header(None, alias="X-Project-ID")) -> Optional[str]:
|
||||||
return x_project_id
|
return x_project_id
|
||||||
|
|
||||||
|
async def get_album_service(dao: DAO = Depends(get_dao)) -> AlbumService:
|
||||||
|
return AlbumService(dao)
|
||||||
|
|
||||||
|
from api.service.post_service import PostService
|
||||||
|
|
||||||
|
def get_post_service(dao: DAO = Depends(get_dao)) -> PostService:
|
||||||
|
return PostService(dao)
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -5,6 +5,8 @@ from fastapi.security import OAuth2PasswordBearer
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from repos.user_repo import UsersRepo, UserStatus
|
from repos.user_repo import UsersRepo, UserStatus
|
||||||
|
from api.dependency import get_dao
|
||||||
|
from repos.dao import DAO
|
||||||
from utils.security import verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
|
from utils.security import verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
|
||||||
from jose import JWTError, jwt
|
from jose import JWTError, jwt
|
||||||
from starlette.requests import Request
|
from starlette.requests import Request
|
||||||
@@ -23,7 +25,7 @@ async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)], repo:
|
|||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
||||||
username: str = payload.get("sub")
|
username: str | None = payload.get("sub")
|
||||||
if username is None:
|
if username is None:
|
||||||
raise credentials_exception
|
raise credentials_exception
|
||||||
except JWTError:
|
except JWTError:
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
from fastapi import APIRouter, Depends, HTTPException, status, Request
|
from fastapi import APIRouter, HTTPException, status, Request
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from api.models.GenerationRequest import GenerationResponse
|
from api.models.GenerationRequest import GenerationResponse
|
||||||
from models.Album import Album
|
from models.Album import Album
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
|
from api.dependency import get_album_service
|
||||||
|
from api.service.album_service import AlbumService
|
||||||
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/api/albums", tags=["Albums"])
|
router = APIRouter(prefix="/api/albums", tags=["Albums"])
|
||||||
|
|
||||||
|
|||||||
@@ -9,10 +9,10 @@ from pymongo import MongoClient
|
|||||||
from starlette import status
|
from starlette import status
|
||||||
from starlette.exceptions import HTTPException
|
from starlette.exceptions import HTTPException
|
||||||
from starlette.requests import Request
|
from starlette.requests import Request
|
||||||
from starlette.responses import Response, JSONResponse
|
from starlette.responses import Response, JSONResponse, StreamingResponse
|
||||||
|
|
||||||
from adapters.s3_adapter import S3Adapter
|
from adapters.s3_adapter import S3Adapter
|
||||||
from api.models.AssetDTO import AssetsResponse, AssetResponse
|
from api.models import AssetsResponse, AssetResponse
|
||||||
from models.Asset import Asset, AssetType, AssetContentType
|
from models.Asset import Asset, AssetType, AssetContentType
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
from api.dependency import get_dao, get_mongo_client, get_s3_adapter
|
from api.dependency import get_dao, get_mongo_client, get_s3_adapter
|
||||||
@@ -33,27 +33,46 @@ async def get_asset(
|
|||||||
asset_id: str,
|
asset_id: str,
|
||||||
request: Request,
|
request: Request,
|
||||||
thumbnail: bool = False,
|
thumbnail: bool = False,
|
||||||
dao: DAO = Depends(get_dao)
|
dao: DAO = Depends(get_dao),
|
||||||
|
s3_adapter: S3Adapter = Depends(get_s3_adapter),
|
||||||
) -> Response:
|
) -> Response:
|
||||||
logger.debug(f"get_asset called for ID: {asset_id}, thumbnail={thumbnail}")
|
logger.debug(f"get_asset called for ID: {asset_id}, thumbnail={thumbnail}")
|
||||||
asset = await dao.assets.get_asset(asset_id)
|
# Загружаем только метаданные (без data/thumbnail bytes)
|
||||||
# 2. Проверка на существование
|
asset = await dao.assets.get_asset(asset_id, with_data=False)
|
||||||
if not asset:
|
if not asset:
|
||||||
raise HTTPException(status_code=404, detail="Asset not found")
|
raise HTTPException(status_code=404, detail="Asset not found")
|
||||||
|
|
||||||
headers = {
|
headers = {
|
||||||
# Кэшировать на 1 год (31536000 сек)
|
|
||||||
"Cache-Control": "public, max-age=31536000, immutable"
|
"Cache-Control": "public, max-age=31536000, immutable"
|
||||||
}
|
}
|
||||||
|
|
||||||
content = asset.data
|
# Thumbnail: маленький, можно грузить в RAM
|
||||||
media_type = "image/png" # Default, or detect
|
if thumbnail:
|
||||||
|
if asset.minio_thumbnail_object_name and s3_adapter:
|
||||||
|
thumb_bytes = await s3_adapter.get_file(asset.minio_thumbnail_object_name)
|
||||||
|
if thumb_bytes:
|
||||||
|
return Response(content=thumb_bytes, media_type="image/jpeg", headers=headers)
|
||||||
|
# Fallback: thumbnail in DB
|
||||||
|
if asset.thumbnail:
|
||||||
|
return Response(content=asset.thumbnail, media_type="image/jpeg", headers=headers)
|
||||||
|
# No thumbnail available — fall through to main content
|
||||||
|
|
||||||
if thumbnail and asset.thumbnail:
|
# Main content: стримим из S3 без загрузки в RAM
|
||||||
content = asset.thumbnail
|
if asset.minio_object_name and s3_adapter:
|
||||||
media_type = "image/jpeg"
|
content_type = "image/png"
|
||||||
|
# if asset.content_type == AssetContentType.VIDEO:
|
||||||
|
# content_type = "video/mp4"
|
||||||
|
return StreamingResponse(
|
||||||
|
s3_adapter.stream_file(asset.minio_object_name),
|
||||||
|
media_type=content_type,
|
||||||
|
headers=headers,
|
||||||
|
)
|
||||||
|
|
||||||
return Response(content=content, media_type=media_type, headers=headers)
|
# Fallback: data stored in DB (legacy)
|
||||||
|
if asset.data:
|
||||||
|
return Response(content=asset.data, media_type="image/png", headers=headers)
|
||||||
|
|
||||||
|
raise HTTPException(status_code=404, detail="Asset data not found")
|
||||||
|
|
||||||
@router.delete("/orphans", dependencies=[Depends(get_current_user)])
|
@router.delete("/orphans", dependencies=[Depends(get_current_user)])
|
||||||
async def delete_orphan_assets_from_minio(
|
async def delete_orphan_assets_from_minio(
|
||||||
@@ -259,8 +278,7 @@ async def upload_asset(
|
|||||||
type=asset.type.value if hasattr(asset.type, "value") else asset.type,
|
type=asset.type.value if hasattr(asset.type, "value") else asset.type,
|
||||||
content_type=asset.content_type.value if hasattr(asset.content_type, "value") else asset.content_type,
|
content_type=asset.content_type.value if hasattr(asset.content_type, "value") else asset.content_type,
|
||||||
linked_char_id=asset.linked_char_id,
|
linked_char_id=asset.linked_char_id,
|
||||||
created_at=asset.created_at,
|
created_at=asset.created_at
|
||||||
url=asset.url
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -5,11 +5,11 @@ from pydantic import BaseModel
|
|||||||
from starlette.exceptions import HTTPException
|
from starlette.exceptions import HTTPException
|
||||||
from starlette.requests import Request
|
from starlette.requests import Request
|
||||||
|
|
||||||
from api.models.AssetDTO import AssetsResponse, AssetResponse
|
from api.models import AssetsResponse, AssetResponse
|
||||||
from api.models.GenerationRequest import GenerationRequest, GenerationResponse
|
from api.models import GenerationRequest, GenerationResponse
|
||||||
from models.Asset import Asset
|
from models.Asset import Asset
|
||||||
from models.Character import Character
|
from models.Character import Character
|
||||||
from api.models.CharacterDTO import CharacterCreateRequest, CharacterUpdateRequest
|
from api.models import CharacterCreateRequest, CharacterUpdateRequest
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
from api.dependency import get_dao
|
from api.dependency import get_dao
|
||||||
|
|
||||||
@@ -24,8 +24,15 @@ router = APIRouter(prefix="/api/characters", tags=["Characters"], dependencies=[
|
|||||||
|
|
||||||
|
|
||||||
@router.get("/", response_model=List[Character])
|
@router.get("/", response_model=List[Character])
|
||||||
async def get_characters(request: Request, dao: DAO = Depends(get_dao), current_user: dict = Depends(get_current_user), project_id: Optional[str] = Depends(get_project_id)) -> List[Character]:
|
async def get_characters(
|
||||||
logger.info("get_characters called")
|
request: Request,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
limit: int = 100,
|
||||||
|
offset: int = 0
|
||||||
|
) -> List[Character]:
|
||||||
|
logger.info(f"get_characters called. Limit: {limit}, Offset: {offset}")
|
||||||
|
|
||||||
user_id_filter = str(current_user["_id"])
|
user_id_filter = str(current_user["_id"])
|
||||||
if project_id:
|
if project_id:
|
||||||
@@ -34,7 +41,12 @@ async def get_characters(request: Request, dao: DAO = Depends(get_dao), current_
|
|||||||
raise HTTPException(status_code=403, detail="Project access denied")
|
raise HTTPException(status_code=403, detail="Project access denied")
|
||||||
user_id_filter = None
|
user_id_filter = None
|
||||||
|
|
||||||
characters = await dao.chars.get_all_characters(created_by=user_id_filter, project_id=project_id)
|
characters = await dao.chars.get_all_characters(
|
||||||
|
created_by=user_id_filter,
|
||||||
|
project_id=project_id,
|
||||||
|
limit=limit,
|
||||||
|
offset=offset
|
||||||
|
)
|
||||||
return characters
|
return characters
|
||||||
|
|
||||||
|
|
||||||
@@ -178,10 +190,3 @@ async def delete_character(
|
|||||||
raise HTTPException(status_code=500, detail="Failed to delete character")
|
raise HTTPException(status_code=500, detail="Failed to delete character")
|
||||||
|
|
||||||
return
|
return
|
||||||
|
|
||||||
|
|
||||||
@router.post("/{character_id}/_run", response_model=GenerationResponse)
|
|
||||||
async def post_character_generation(character_id: str, generation: GenerationRequest,
|
|
||||||
request: Request) -> GenerationResponse:
|
|
||||||
logger.info(f"post_character_generation called. CharacterID: {character_id}")
|
|
||||||
generation_service = request.app.state.generation_service
|
|
||||||
|
|||||||
180
api/endpoints/environment_router.py
Normal file
180
api/endpoints/environment_router.py
Normal file
@@ -0,0 +1,180 @@
|
|||||||
|
import logging
|
||||||
|
from typing import List, Optional
|
||||||
|
from fastapi import APIRouter, Depends, HTTPException
|
||||||
|
from starlette import status
|
||||||
|
|
||||||
|
from api.dependency import get_dao
|
||||||
|
from api.endpoints.auth import get_current_user
|
||||||
|
from api.models.EnvironmentRequest import EnvironmentCreate, EnvironmentUpdate, AssetToEnvironment, AssetsToEnvironment
|
||||||
|
from models.Environment import Environment
|
||||||
|
from repos.dao import DAO
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/environments", tags=["Environments"], dependencies=[Depends(get_current_user)])
|
||||||
|
|
||||||
|
|
||||||
|
async def check_character_access(character_id: str, current_user: dict, dao: DAO):
|
||||||
|
character = await dao.chars.get_character(character_id)
|
||||||
|
if not character:
|
||||||
|
raise HTTPException(status_code=404, detail="Character not found")
|
||||||
|
|
||||||
|
is_creator = character.created_by == str(current_user["_id"])
|
||||||
|
is_project_member = False
|
||||||
|
if character.project_id and character.project_id in current_user.get("project_ids", []):
|
||||||
|
is_project_member = True
|
||||||
|
|
||||||
|
if not is_creator and not is_project_member:
|
||||||
|
raise HTTPException(status_code=403, detail="Access denied to character")
|
||||||
|
return character
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/", response_model=Environment)
|
||||||
|
async def create_environment(
|
||||||
|
env_req: EnvironmentCreate,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
logger.info(f"Creating environment '{env_req.name}' for character {env_req.character_id}")
|
||||||
|
await check_character_access(env_req.character_id, current_user, dao)
|
||||||
|
|
||||||
|
# Verify assets exist if provided
|
||||||
|
if env_req.asset_ids:
|
||||||
|
for aid in env_req.asset_ids:
|
||||||
|
asset = await dao.assets.get_asset(aid)
|
||||||
|
if not asset:
|
||||||
|
raise HTTPException(status_code=400, detail=f"Asset {aid} not found")
|
||||||
|
|
||||||
|
new_env = Environment(**env_req.model_dump())
|
||||||
|
created_env = await dao.environments.create_env(new_env)
|
||||||
|
return created_env
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/character/{character_id}", response_model=List[Environment])
|
||||||
|
async def get_character_environments(
|
||||||
|
character_id: str,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
logger.info(f"Getting environments for character {character_id}")
|
||||||
|
await check_character_access(character_id, current_user, dao)
|
||||||
|
return await dao.environments.get_character_envs(character_id)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/{env_id}", response_model=Environment)
|
||||||
|
async def get_environment(
|
||||||
|
env_id: str,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
return env
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/{env_id}", response_model=Environment)
|
||||||
|
async def update_environment(
|
||||||
|
env_id: str,
|
||||||
|
env_update: EnvironmentUpdate,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
|
||||||
|
update_data = env_update.model_dump(exclude_unset=True)
|
||||||
|
if not update_data:
|
||||||
|
return env
|
||||||
|
|
||||||
|
success = await dao.environments.update_env(env_id, update_data)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=500, detail="Failed to update environment")
|
||||||
|
|
||||||
|
return await dao.environments.get_env(env_id)
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{env_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||||||
|
async def delete_environment(
|
||||||
|
env_id: str,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
|
||||||
|
success = await dao.environments.delete_env(env_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=500, detail="Failed to delete environment")
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/{env_id}/assets", status_code=status.HTTP_200_OK)
|
||||||
|
async def add_asset_to_environment(
|
||||||
|
env_id: str,
|
||||||
|
req: AssetToEnvironment,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
|
||||||
|
# Verify asset exists
|
||||||
|
asset = await dao.assets.get_asset(req.asset_id)
|
||||||
|
if not asset:
|
||||||
|
raise HTTPException(status_code=404, detail="Asset not found")
|
||||||
|
|
||||||
|
success = await dao.environments.add_asset(env_id, req.asset_id)
|
||||||
|
return {"success": success}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/{env_id}/assets/batch", status_code=status.HTTP_200_OK)
|
||||||
|
async def add_assets_batch_to_environment(
|
||||||
|
env_id: str,
|
||||||
|
req: AssetsToEnvironment,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
|
||||||
|
# Verify all assets exist
|
||||||
|
assets = await dao.assets.get_assets_by_ids(req.asset_ids)
|
||||||
|
if len(assets) != len(req.asset_ids):
|
||||||
|
found_ids = {a.id for a in assets}
|
||||||
|
missing_ids = [aid for aid in req.asset_ids if aid not in found_ids]
|
||||||
|
raise HTTPException(status_code=404, detail=f"Some assets not found: {missing_ids}")
|
||||||
|
|
||||||
|
success = await dao.environments.add_assets(env_id, req.asset_ids)
|
||||||
|
return {"success": success}
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{env_id}/assets/{asset_id}", status_code=status.HTTP_200_OK)
|
||||||
|
async def remove_asset_from_environment(
|
||||||
|
env_id: str,
|
||||||
|
asset_id: str,
|
||||||
|
dao: DAO = Depends(get_dao),
|
||||||
|
current_user: dict = Depends(get_current_user)
|
||||||
|
):
|
||||||
|
env = await dao.environments.get_env(env_id)
|
||||||
|
if not env:
|
||||||
|
raise HTTPException(status_code=404, detail="Environment not found")
|
||||||
|
|
||||||
|
await check_character_access(env.character_id, current_user, dao)
|
||||||
|
|
||||||
|
success = await dao.environments.remove_asset(env_id, asset_id)
|
||||||
|
return {"success": success}
|
||||||
@@ -1,25 +1,32 @@
|
|||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import json
|
||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
|
|
||||||
from fastapi import APIRouter, UploadFile, File, Form, Header, HTTPException
|
from fastapi import APIRouter, UploadFile, File, Form, Header, HTTPException
|
||||||
from fastapi.params import Depends
|
from fastapi.params import Depends
|
||||||
|
from starlette import status
|
||||||
from starlette.requests import Request
|
from starlette.requests import Request
|
||||||
|
|
||||||
from api import service
|
from config import settings
|
||||||
from api.dependency import get_generation_service, get_project_id, get_dao
|
from api.dependency import get_generation_service, get_project_id, get_dao
|
||||||
from repos.dao import DAO
|
from api.endpoints.auth import get_current_user
|
||||||
|
from api.models import (
|
||||||
from api.models.GenerationRequest import GenerationResponse, GenerationRequest, GenerationsResponse, PromptResponse, PromptRequest
|
GenerationResponse,
|
||||||
|
GenerationRequest,
|
||||||
|
GenerationsResponse,
|
||||||
|
PromptResponse,
|
||||||
|
PromptRequest,
|
||||||
|
GenerationGroupResponse,
|
||||||
|
FinancialReport,
|
||||||
|
ExternalGenerationRequest
|
||||||
|
)
|
||||||
from api.service.generation_service import GenerationService
|
from api.service.generation_service import GenerationService
|
||||||
from models.Generation import Generation
|
from repos.dao import DAO
|
||||||
|
from utils.external_auth import verify_signature
|
||||||
from starlette import status
|
|
||||||
|
|
||||||
import logging
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
from api.endpoints.auth import get_current_user
|
|
||||||
|
|
||||||
router = APIRouter(prefix='/api/generations', tags=["Generation"])
|
router = APIRouter(prefix='/api/generations', tags=["Generation"])
|
||||||
|
|
||||||
|
|
||||||
@@ -68,12 +75,53 @@ async def get_generations(character_id: Optional[str] = None, limit: int = 10, o
|
|||||||
return await generation_service.get_generations(character_id, limit=limit, offset=offset, user_id=user_id_filter, project_id=project_id)
|
return await generation_service.get_generations(character_id, limit=limit, offset=offset, user_id=user_id_filter, project_id=project_id)
|
||||||
|
|
||||||
|
|
||||||
@router.post("/_run", response_model=GenerationResponse)
|
@router.get("/usage", response_model=FinancialReport)
|
||||||
|
async def get_usage_report(
|
||||||
|
breakdown: Optional[str] = None, # "user" or "project"
|
||||||
|
generation_service: GenerationService = Depends(get_generation_service),
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
dao: DAO = Depends(get_dao)
|
||||||
|
) -> FinancialReport:
|
||||||
|
"""
|
||||||
|
Returns usage statistics (runs, tokens, cost) for the current user or project.
|
||||||
|
If project_id is provided, returns stats for that project.
|
||||||
|
Otherwise, returns stats for the current user.
|
||||||
|
"""
|
||||||
|
user_id_filter = str(current_user["_id"])
|
||||||
|
breakdown_by = None
|
||||||
|
|
||||||
|
if project_id:
|
||||||
|
# Permission check
|
||||||
|
project = await dao.projects.get_project(project_id)
|
||||||
|
if not project or str(current_user["_id"]) not in project.members:
|
||||||
|
raise HTTPException(status_code=403, detail="Project access denied")
|
||||||
|
user_id_filter = None # If we are in project, we see stats for the WHOLE project by default
|
||||||
|
if breakdown == "user":
|
||||||
|
breakdown_by = "created_by"
|
||||||
|
elif breakdown == "project":
|
||||||
|
breakdown_by = "project_id"
|
||||||
|
else:
|
||||||
|
# Default: Stats for current user
|
||||||
|
if breakdown == "project":
|
||||||
|
breakdown_by = "project_id"
|
||||||
|
elif breakdown == "user":
|
||||||
|
# This would breakdown personal usage by user (yourself), but could be useful if it included collaborators?
|
||||||
|
# No, if project_id is None, it's personal.
|
||||||
|
breakdown_by = "created_by"
|
||||||
|
|
||||||
|
return await generation_service.get_financial_report(
|
||||||
|
user_id=user_id_filter,
|
||||||
|
project_id=project_id,
|
||||||
|
breakdown_by=breakdown_by
|
||||||
|
)
|
||||||
|
|
||||||
|
@router.post("/_run", response_model=GenerationGroupResponse)
|
||||||
async def post_generation(generation: GenerationRequest, request: Request,
|
async def post_generation(generation: GenerationRequest, request: Request,
|
||||||
generation_service: GenerationService = Depends(get_generation_service),
|
generation_service: GenerationService = Depends(get_generation_service),
|
||||||
current_user: dict = Depends(get_current_user),
|
current_user: dict = Depends(get_current_user),
|
||||||
project_id: Optional[str] = Depends(get_project_id),
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
dao: DAO = Depends(get_dao)) -> GenerationResponse:
|
dao: DAO = Depends(get_dao)) -> GenerationGroupResponse:
|
||||||
logger.info(f"post_generation (run) called. LinkedCharId: {generation.linked_character_id}, PromptLength: {len(generation.prompt)}")
|
logger.info(f"post_generation (run) called. LinkedCharId: {generation.linked_character_id}, PromptLength: {len(generation.prompt)}")
|
||||||
|
|
||||||
if project_id:
|
if project_id:
|
||||||
@@ -85,16 +133,6 @@ async def post_generation(generation: GenerationRequest, request: Request,
|
|||||||
return await generation_service.create_generation_task(generation, user_id=str(current_user.get("_id")))
|
return await generation_service.create_generation_task(generation, user_id=str(current_user.get("_id")))
|
||||||
|
|
||||||
|
|
||||||
@router.get("/{generation_id}", response_model=GenerationResponse)
|
|
||||||
async def get_generation(generation_id: str,
|
|
||||||
generation_service: GenerationService = Depends(get_generation_service),
|
|
||||||
current_user: dict = Depends(get_current_user)) -> GenerationResponse:
|
|
||||||
logger.debug(f"get_generation called for ID: {generation_id}")
|
|
||||||
gen = await generation_service.get_generation(generation_id)
|
|
||||||
if gen and gen.created_by != str(current_user["_id"]):
|
|
||||||
raise HTTPException(status_code=403, detail="Access denied")
|
|
||||||
return gen
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/running")
|
@router.get("/running")
|
||||||
async def get_running_generations(request: Request,
|
async def get_running_generations(request: Request,
|
||||||
@@ -113,6 +151,35 @@ async def get_running_generations(request: Request,
|
|||||||
return await generation_service.get_running_generations(user_id=user_id_filter, project_id=project_id)
|
return await generation_service.get_running_generations(user_id=user_id_filter, project_id=project_id)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/group/{group_id}", response_model=GenerationGroupResponse)
|
||||||
|
async def get_generation_group(group_id: str,
|
||||||
|
generation_service: GenerationService = Depends(get_generation_service),
|
||||||
|
current_user: dict = Depends(get_current_user)):
|
||||||
|
logger.info(f"get_generation_group called for group_id: {group_id}")
|
||||||
|
generations = await generation_service.dao.generations.get_generations_by_group(group_id)
|
||||||
|
gen_responses = [GenerationResponse(**gen.model_dump()) for gen in generations]
|
||||||
|
return GenerationGroupResponse(generation_group_id=group_id, generations=gen_responses)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/{generation_id}", response_model=GenerationResponse)
|
||||||
|
async def get_generation(generation_id: str,
|
||||||
|
generation_service: GenerationService = Depends(get_generation_service),
|
||||||
|
current_user: dict = Depends(get_current_user)) -> GenerationResponse:
|
||||||
|
logger.debug(f"get_generation called for ID: {generation_id}")
|
||||||
|
gen = await generation_service.get_generation(generation_id)
|
||||||
|
if gen and gen.created_by != str(current_user["_id"]):
|
||||||
|
# Check project membership
|
||||||
|
is_member = False
|
||||||
|
if gen.project_id:
|
||||||
|
project = await generation_service.dao.projects.get_project(gen.project_id)
|
||||||
|
if project and str(current_user["_id"]) in project.members:
|
||||||
|
is_member = True
|
||||||
|
|
||||||
|
if not is_member:
|
||||||
|
raise HTTPException(status_code=403, detail="Access denied")
|
||||||
|
return gen
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/import", response_model=GenerationResponse)
|
@router.post("/import", response_model=GenerationResponse)
|
||||||
@@ -125,17 +192,13 @@ async def import_external_generation(
|
|||||||
Import a generation from an external source.
|
Import a generation from an external source.
|
||||||
Requires server-to-server authentication via HMAC signature.
|
Requires server-to-server authentication via HMAC signature.
|
||||||
"""
|
"""
|
||||||
import os
|
|
||||||
from utils.external_auth import verify_signature
|
|
||||||
from api.models.ExternalGenerationDTO import ExternalGenerationRequest
|
|
||||||
|
|
||||||
logger.info("import_external_generation called")
|
logger.info("import_external_generation called")
|
||||||
|
|
||||||
# Get raw request body for signature verification
|
# Get raw request body for signature verification
|
||||||
body = await request.body()
|
body = await request.body()
|
||||||
|
|
||||||
# Verify signature
|
# Verify signature
|
||||||
secret = os.getenv("EXTERNAL_API_SECRET")
|
secret = settings.EXTERNAL_API_SECRET
|
||||||
if not secret:
|
if not secret:
|
||||||
logger.error("EXTERNAL_API_SECRET not configured")
|
logger.error("EXTERNAL_API_SECRET not configured")
|
||||||
raise HTTPException(status_code=500, detail="Server configuration error")
|
raise HTTPException(status_code=500, detail="Server configuration error")
|
||||||
@@ -145,7 +208,6 @@ async def import_external_generation(
|
|||||||
raise HTTPException(status_code=401, detail="Invalid signature")
|
raise HTTPException(status_code=401, detail="Invalid signature")
|
||||||
|
|
||||||
# Parse request body
|
# Parse request body
|
||||||
import json
|
|
||||||
try:
|
try:
|
||||||
data = json.loads(body.decode('utf-8'))
|
data = json.loads(body.decode('utf-8'))
|
||||||
external_gen = ExternalGenerationRequest(**data)
|
external_gen = ExternalGenerationRequest(**data)
|
||||||
|
|||||||
104
api/endpoints/idea_router.py
Normal file
104
api/endpoints/idea_router.py
Normal file
@@ -0,0 +1,104 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from fastapi import APIRouter, Depends, HTTPException, Query, Body
|
||||||
|
from api.dependency import get_idea_service, get_project_id, get_generation_service
|
||||||
|
from api.endpoints.auth import get_current_user
|
||||||
|
from api.service.idea_service import IdeaService
|
||||||
|
from api.service.generation_service import GenerationService
|
||||||
|
from models.Idea import Idea
|
||||||
|
from api.models import GenerationResponse, GenerationsResponse
|
||||||
|
from api.models import IdeaRequest, PostRequest # Adjusting for general model usage
|
||||||
|
from api.models.IdeaRequest import IdeaCreateRequest, IdeaUpdateRequest, IdeaResponse
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/ideas", tags=["ideas"])
|
||||||
|
|
||||||
|
@router.post("", response_model=Idea)
|
||||||
|
async def create_idea(
|
||||||
|
request: IdeaCreateRequest,
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
pid = project_id or request.project_id
|
||||||
|
|
||||||
|
return await idea_service.create_idea(request.name, request.description, pid, str(current_user["_id"]))
|
||||||
|
|
||||||
|
@router.get("", response_model=List[IdeaResponse])
|
||||||
|
async def get_ideas(
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
limit: int = 20,
|
||||||
|
offset: int = 0,
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
return await idea_service.get_ideas(project_id, str(current_user["_id"]), limit, offset)
|
||||||
|
|
||||||
|
@router.get("/{idea_id}", response_model=Idea)
|
||||||
|
async def get_idea(
|
||||||
|
idea_id: str,
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
idea = await idea_service.get_idea(idea_id)
|
||||||
|
if not idea:
|
||||||
|
raise HTTPException(status_code=404, detail="Idea not found")
|
||||||
|
return idea
|
||||||
|
|
||||||
|
@router.put("/{idea_id}", response_model=Idea)
|
||||||
|
async def update_idea(
|
||||||
|
idea_id: str,
|
||||||
|
request: IdeaUpdateRequest,
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
idea = await idea_service.update_idea(idea_id, request.name, request.description)
|
||||||
|
if not idea:
|
||||||
|
raise HTTPException(status_code=404, detail="Idea not found")
|
||||||
|
return idea
|
||||||
|
|
||||||
|
@router.delete("/{idea_id}")
|
||||||
|
async def delete_idea(
|
||||||
|
idea_id: str,
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
success = await idea_service.delete_idea(idea_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Idea not found or could not be deleted")
|
||||||
|
return {"status": "success"}
|
||||||
|
|
||||||
|
@router.get("/{idea_id}/generations", response_model=GenerationsResponse)
|
||||||
|
async def get_idea_generations(
|
||||||
|
idea_id: str,
|
||||||
|
limit: int = 50,
|
||||||
|
offset: int = 0,
|
||||||
|
generation_service: GenerationService = Depends(get_generation_service)
|
||||||
|
):
|
||||||
|
# Depending on how generation service implements filtering by idea_id.
|
||||||
|
# We might need to update generation_service to support getting by idea_id directly
|
||||||
|
# or ensure generic get_generations supports it.
|
||||||
|
# Looking at generation_router.py, get_generations doesn't have idea_id arg?
|
||||||
|
# Let's check generation_service.get_generations signature again.
|
||||||
|
# It has: (character_id, limit, offset, user_id, project_id). NO IDEA_ID.
|
||||||
|
# I need to update GenerationService.get_generations too!
|
||||||
|
|
||||||
|
# For now, let's assume I will update it.
|
||||||
|
return await generation_service.get_generations(idea_id=idea_id, limit=limit, offset=offset)
|
||||||
|
|
||||||
|
@router.post("/{idea_id}/generations/{generation_id}")
|
||||||
|
async def add_generation_to_idea(
|
||||||
|
idea_id: str,
|
||||||
|
generation_id: str,
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
success = await idea_service.add_generation_to_idea(idea_id, generation_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Idea or Generation not found")
|
||||||
|
return {"status": "success"}
|
||||||
|
|
||||||
|
@router.delete("/{idea_id}/generations/{generation_id}")
|
||||||
|
async def remove_generation_from_idea(
|
||||||
|
idea_id: str,
|
||||||
|
generation_id: str,
|
||||||
|
idea_service: IdeaService = Depends(get_idea_service)
|
||||||
|
):
|
||||||
|
success = await idea_service.remove_generation_from_idea(idea_id, generation_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Idea or Generation not found")
|
||||||
|
return {"status": "success"}
|
||||||
99
api/endpoints/post_router.py
Normal file
99
api/endpoints/post_router.py
Normal file
@@ -0,0 +1,99 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from datetime import datetime
|
||||||
|
from fastapi import APIRouter, Depends, HTTPException
|
||||||
|
|
||||||
|
from api.dependency import get_post_service, get_project_id
|
||||||
|
from api.endpoints.auth import get_current_user
|
||||||
|
from api.service.post_service import PostService
|
||||||
|
from api.models import PostRequest, PostCreateRequest, PostUpdateRequest, AddGenerationsRequest
|
||||||
|
from models.Post import Post
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/posts", tags=["posts"])
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("", response_model=Post)
|
||||||
|
async def create_post(
|
||||||
|
request: PostCreateRequest,
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
pid = project_id or request.project_id
|
||||||
|
return await post_service.create_post(
|
||||||
|
date=request.date,
|
||||||
|
topic=request.topic,
|
||||||
|
generation_ids=request.generation_ids,
|
||||||
|
project_id=pid,
|
||||||
|
user_id=str(current_user["_id"]),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("", response_model=List[Post])
|
||||||
|
async def get_posts(
|
||||||
|
project_id: Optional[str] = Depends(get_project_id),
|
||||||
|
limit: int = 200,
|
||||||
|
offset: int = 0,
|
||||||
|
date_from: Optional[datetime] = None,
|
||||||
|
date_to: Optional[datetime] = None,
|
||||||
|
current_user: dict = Depends(get_current_user),
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
return await post_service.get_posts(project_id, str(current_user["_id"]), limit, offset, date_from, date_to)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/{post_id}", response_model=Post)
|
||||||
|
async def get_post(
|
||||||
|
post_id: str,
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
post = await post_service.get_post(post_id)
|
||||||
|
if not post:
|
||||||
|
raise HTTPException(status_code=404, detail="Post not found")
|
||||||
|
return post
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/{post_id}", response_model=Post)
|
||||||
|
async def update_post(
|
||||||
|
post_id: str,
|
||||||
|
request: PostUpdateRequest,
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
post = await post_service.update_post(post_id, date=request.date, topic=request.topic)
|
||||||
|
if not post:
|
||||||
|
raise HTTPException(status_code=404, detail="Post not found")
|
||||||
|
return post
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{post_id}")
|
||||||
|
async def delete_post(
|
||||||
|
post_id: str,
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
success = await post_service.delete_post(post_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Post not found or could not be deleted")
|
||||||
|
return {"status": "success"}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/{post_id}/generations")
|
||||||
|
async def add_generations(
|
||||||
|
post_id: str,
|
||||||
|
request: AddGenerationsRequest,
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
success = await post_service.add_generations(post_id, request.generation_ids)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Post not found")
|
||||||
|
return {"status": "success"}
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{post_id}/generations/{generation_id}")
|
||||||
|
async def remove_generation(
|
||||||
|
post_id: str,
|
||||||
|
generation_id: str,
|
||||||
|
post_service: PostService = Depends(get_post_service),
|
||||||
|
):
|
||||||
|
success = await post_service.remove_generation(post_id, generation_id)
|
||||||
|
if not success:
|
||||||
|
raise HTTPException(status_code=404, detail="Post not found or generation not linked")
|
||||||
|
return {"status": "success"}
|
||||||
@@ -1,4 +1,6 @@
|
|||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
|
|
||||||
|
from bson import ObjectId
|
||||||
from fastapi import APIRouter, Depends, HTTPException, status
|
from fastapi import APIRouter, Depends, HTTPException, status
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from api.dependency import get_dao
|
from api.dependency import get_dao
|
||||||
@@ -12,14 +14,46 @@ class ProjectCreate(BaseModel):
|
|||||||
name: str
|
name: str
|
||||||
description: Optional[str] = None
|
description: Optional[str] = None
|
||||||
|
|
||||||
|
class ProjectMemberResponse(BaseModel):
|
||||||
|
id: str
|
||||||
|
username: str
|
||||||
|
|
||||||
class ProjectResponse(BaseModel):
|
class ProjectResponse(BaseModel):
|
||||||
id: str
|
id: str
|
||||||
name: str
|
name: str
|
||||||
description: Optional[str] = None
|
description: Optional[str] = None
|
||||||
owner_id: str
|
owner_id: str
|
||||||
members: List[str]
|
members: List[ProjectMemberResponse]
|
||||||
is_owner: bool = False
|
is_owner: bool = False
|
||||||
|
|
||||||
|
async def _get_project_response(project: Project, current_user_id: str, dao: DAO) -> ProjectResponse:
|
||||||
|
member_responses = []
|
||||||
|
for member_id in project.members:
|
||||||
|
# We need a way to get user by ID. Let's check UsersRepo for get_user by ObjectId or string.
|
||||||
|
# Currently UsersRepo has get_user(user_id: int) for Telegram IDs.
|
||||||
|
# But for Web users we might need to search by _id.
|
||||||
|
# Let's try to get user info.
|
||||||
|
# Since project.members contains strings (ObjectIds as strings), we search by _id.
|
||||||
|
user_doc = await dao.users.collection.find_one({"_id": ObjectId(member_id)})
|
||||||
|
if not user_doc and member_id.isdigit():
|
||||||
|
# Fallback for telegram IDs if they are stored as strings of digits
|
||||||
|
user_doc = await dao.users.get_user(int(member_id))
|
||||||
|
|
||||||
|
username = "unknown"
|
||||||
|
if user_doc:
|
||||||
|
username = user_doc.get("username", "unknown")
|
||||||
|
|
||||||
|
member_responses.append(ProjectMemberResponse(id=member_id, username=username))
|
||||||
|
|
||||||
|
return ProjectResponse(
|
||||||
|
id=project.id,
|
||||||
|
name=project.name,
|
||||||
|
description=project.description,
|
||||||
|
owner_id=project.owner_id,
|
||||||
|
members=member_responses,
|
||||||
|
is_owner=(project.owner_id == current_user_id)
|
||||||
|
)
|
||||||
|
|
||||||
@router.post("", response_model=ProjectResponse)
|
@router.post("", response_model=ProjectResponse)
|
||||||
async def create_project(
|
async def create_project(
|
||||||
project_data: ProjectCreate,
|
project_data: ProjectCreate,
|
||||||
@@ -34,27 +68,15 @@ async def create_project(
|
|||||||
members=[user_id]
|
members=[user_id]
|
||||||
)
|
)
|
||||||
project_id = await dao.projects.create_project(new_project)
|
project_id = await dao.projects.create_project(new_project)
|
||||||
|
new_project.id = project_id
|
||||||
|
|
||||||
# Add project to user's project list
|
# Add project to user's project list
|
||||||
# Assuming user_repo has a method to add project or we do it directly?
|
|
||||||
# UserRepo doesn't have add_project method yet.
|
|
||||||
# But since UserRepo is just a wrapper around collection, lets add it here or update UserRepo later?
|
|
||||||
# Better to update UserRepo. For now, let's just return success.
|
|
||||||
# But user needs to see it in list.
|
|
||||||
# Update user in DB
|
|
||||||
await dao.users.collection.update_one(
|
await dao.users.collection.update_one(
|
||||||
{"_id": current_user["_id"]},
|
{"_id": current_user["_id"]},
|
||||||
{"$addToSet": {"project_ids": project_id}}
|
{"$addToSet": {"project_ids": project_id}}
|
||||||
)
|
)
|
||||||
|
|
||||||
return ProjectResponse(
|
return await _get_project_response(new_project, user_id, dao)
|
||||||
id=project_id,
|
|
||||||
name=new_project.name,
|
|
||||||
description=new_project.description,
|
|
||||||
owner_id=new_project.owner_id,
|
|
||||||
members=new_project.members,
|
|
||||||
is_owner=True
|
|
||||||
)
|
|
||||||
|
|
||||||
@router.get("", response_model=List[ProjectResponse])
|
@router.get("", response_model=List[ProjectResponse])
|
||||||
async def get_my_projects(
|
async def get_my_projects(
|
||||||
@@ -66,14 +88,7 @@ async def get_my_projects(
|
|||||||
|
|
||||||
responses = []
|
responses = []
|
||||||
for p in projects:
|
for p in projects:
|
||||||
responses.append(ProjectResponse(
|
responses.append(await _get_project_response(p, user_id, dao))
|
||||||
id=p.id,
|
|
||||||
name=p.name,
|
|
||||||
description=p.description,
|
|
||||||
owner_id=p.owner_id,
|
|
||||||
members=p.members,
|
|
||||||
is_owner=(p.owner_id == user_id)
|
|
||||||
))
|
|
||||||
return responses
|
return responses
|
||||||
|
|
||||||
class MemberAdd(BaseModel):
|
class MemberAdd(BaseModel):
|
||||||
|
|||||||
22
api/models/EnvironmentRequest.py
Normal file
22
api/models/EnvironmentRequest.py
Normal file
@@ -0,0 +1,22 @@
|
|||||||
|
from typing import Optional, List
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
|
||||||
|
class EnvironmentCreate(BaseModel):
|
||||||
|
character_id: str
|
||||||
|
name: str = Field(..., min_length=1)
|
||||||
|
description: Optional[str] = None
|
||||||
|
asset_ids: Optional[List[str]] = []
|
||||||
|
|
||||||
|
|
||||||
|
class EnvironmentUpdate(BaseModel):
|
||||||
|
name: Optional[str] = Field(None, min_length=1)
|
||||||
|
description: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class AssetToEnvironment(BaseModel):
|
||||||
|
asset_id: str
|
||||||
|
|
||||||
|
|
||||||
|
class AssetsToEnvironment(BaseModel):
|
||||||
|
asset_ids: List[str]
|
||||||
18
api/models/FinancialUsageDTO.py
Normal file
18
api/models/FinancialUsageDTO.py
Normal file
@@ -0,0 +1,18 @@
|
|||||||
|
from pydantic import BaseModel
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
|
class UsageStats(BaseModel):
|
||||||
|
total_runs: int
|
||||||
|
total_tokens: int
|
||||||
|
total_input_tokens: int
|
||||||
|
total_output_tokens: int
|
||||||
|
total_cost: float
|
||||||
|
|
||||||
|
class UsageByEntity(BaseModel):
|
||||||
|
entity_id: Optional[str] = None
|
||||||
|
stats: UsageStats
|
||||||
|
|
||||||
|
class FinancialReport(BaseModel):
|
||||||
|
summary: UsageStats
|
||||||
|
by_user: Optional[List[UsageByEntity]] = None
|
||||||
|
by_project: Optional[List[UsageByEntity]] = None
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
from datetime import datetime, UTC
|
from datetime import datetime, UTC
|
||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
from models.Asset import Asset
|
from models.Asset import Asset
|
||||||
from models.Generation import GenerationStatus
|
from models.Generation import GenerationStatus
|
||||||
@@ -16,7 +16,10 @@ class GenerationRequest(BaseModel):
|
|||||||
telegram_id: Optional[int] = None
|
telegram_id: Optional[int] = None
|
||||||
use_profile_image: bool = True
|
use_profile_image: bool = True
|
||||||
assets_list: List[str]
|
assets_list: List[str]
|
||||||
|
environment_id: Optional[str] = None
|
||||||
project_id: Optional[str] = None
|
project_id: Optional[str] = None
|
||||||
|
idea_id: Optional[str] = None
|
||||||
|
count: int = Field(default=1, ge=1, le=10)
|
||||||
|
|
||||||
|
|
||||||
class GenerationsResponse(BaseModel):
|
class GenerationsResponse(BaseModel):
|
||||||
@@ -45,10 +48,16 @@ class GenerationResponse(BaseModel):
|
|||||||
progress: int = 0
|
progress: int = 0
|
||||||
cost: Optional[float] = None
|
cost: Optional[float] = None
|
||||||
created_by: Optional[str] = None
|
created_by: Optional[str] = None
|
||||||
|
generation_group_id: Optional[str] = None
|
||||||
|
idea_id: Optional[str] = None
|
||||||
created_at: datetime = datetime.now(UTC)
|
created_at: datetime = datetime.now(UTC)
|
||||||
updated_at: datetime = datetime.now(UTC)
|
updated_at: datetime = datetime.now(UTC)
|
||||||
|
|
||||||
|
|
||||||
|
class GenerationGroupResponse(BaseModel):
|
||||||
|
generation_group_id: str
|
||||||
|
generations: List[GenerationResponse]
|
||||||
|
|
||||||
|
|
||||||
class PromptRequest(BaseModel):
|
class PromptRequest(BaseModel):
|
||||||
prompt: str
|
prompt: str
|
||||||
|
|||||||
16
api/models/IdeaRequest.py
Normal file
16
api/models/IdeaRequest.py
Normal file
@@ -0,0 +1,16 @@
|
|||||||
|
from typing import Optional
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from models.Idea import Idea
|
||||||
|
from api.models.GenerationRequest import GenerationResponse
|
||||||
|
|
||||||
|
class IdeaCreateRequest(BaseModel):
|
||||||
|
name: str
|
||||||
|
description: Optional[str] = None
|
||||||
|
project_id: Optional[str] = None # Optional in body if passed via header/dependency
|
||||||
|
|
||||||
|
class IdeaUpdateRequest(BaseModel):
|
||||||
|
name: Optional[str] = None
|
||||||
|
description: Optional[str] = None
|
||||||
|
|
||||||
|
class IdeaResponse(Idea):
|
||||||
|
last_generation: Optional[GenerationResponse] = None
|
||||||
19
api/models/PostRequest.py
Normal file
19
api/models/PostRequest.py
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
from datetime import datetime
|
||||||
|
from typing import Optional, List
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
|
class PostCreateRequest(BaseModel):
|
||||||
|
date: datetime
|
||||||
|
topic: str
|
||||||
|
generation_ids: List[str] = []
|
||||||
|
project_id: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class PostUpdateRequest(BaseModel):
|
||||||
|
date: Optional[datetime] = None
|
||||||
|
topic: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class AddGenerationsRequest(BaseModel):
|
||||||
|
generation_ids: List[str]
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
from .AssetDTO import AssetResponse, AssetsResponse
|
||||||
|
from .CharacterDTO import CharacterCreateRequest, CharacterUpdateRequest
|
||||||
|
from .ExternalGenerationDTO import ExternalGenerationRequest
|
||||||
|
from .FinancialUsageDTO import FinancialReport, UsageStats, UsageByEntity
|
||||||
|
from .GenerationRequest import GenerationRequest, GenerationResponse, GenerationsResponse, GenerationGroupResponse, PromptRequest, PromptResponse
|
||||||
|
from .IdeaRequest import IdeaCreateRequest, IdeaUpdateRequest, IdeaResponse
|
||||||
|
from .PostRequest import PostCreateRequest, PostUpdateRequest, AddGenerationsRequest
|
||||||
|
|||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,26 +1,32 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import base64
|
||||||
import logging
|
import logging
|
||||||
import random
|
import random
|
||||||
import base64
|
|
||||||
from datetime import datetime, UTC
|
from datetime import datetime, UTC
|
||||||
from typing import List, Optional, Tuple, Any, Dict
|
from typing import List, Optional, Tuple, Any, Dict
|
||||||
from io import BytesIO
|
from uuid import uuid4
|
||||||
import httpx
|
|
||||||
|
|
||||||
|
import httpx
|
||||||
from aiogram import Bot
|
from aiogram import Bot
|
||||||
from aiogram.types import BufferedInputFile
|
from aiogram.types import BufferedInputFile
|
||||||
|
from fastapi import HTTPException
|
||||||
|
|
||||||
from adapters.Exception import GoogleGenerationException
|
from adapters.Exception import GoogleGenerationException
|
||||||
from adapters.google_adapter import GoogleAdapter
|
from adapters.google_adapter import GoogleAdapter
|
||||||
from api.models.GenerationRequest import GenerationRequest, GenerationResponse, GenerationsResponse
|
from adapters.s3_adapter import S3Adapter
|
||||||
|
from api.models import FinancialReport, UsageStats, UsageByEntity
|
||||||
|
from api.models import GenerationRequest, GenerationResponse, GenerationsResponse, GenerationGroupResponse
|
||||||
# Импортируйте ваши модели DAO, Asset, Generation корректно
|
# Импортируйте ваши модели DAO, Asset, Generation корректно
|
||||||
from models.Asset import Asset, AssetType, AssetContentType
|
from models.Asset import Asset, AssetType, AssetContentType
|
||||||
from models.Generation import Generation, GenerationStatus
|
from models.Generation import Generation, GenerationStatus
|
||||||
from models.enums import AspectRatios, Quality, GenType
|
from models.enums import AspectRatios, Quality
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
from adapters.s3_adapter import S3Adapter
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Limit concurrent generations to 4
|
||||||
|
generation_semaphore = asyncio.Semaphore(4)
|
||||||
|
|
||||||
|
|
||||||
# --- Вспомогательная функция генерации ---
|
# --- Вспомогательная функция генерации ---
|
||||||
async def generate_image_task(
|
async def generate_image_task(
|
||||||
@@ -50,16 +56,18 @@ async def generate_image_task(
|
|||||||
logger.info(f"generate_image_task completed, received {len(generated_images_io) if generated_images_io else 0} images")
|
logger.info(f"generate_image_task completed, received {len(generated_images_io) if generated_images_io else 0} images")
|
||||||
except GoogleGenerationException as e:
|
except GoogleGenerationException as e:
|
||||||
raise e
|
raise e
|
||||||
|
finally:
|
||||||
|
# Освобождаем входные данные — они больше не нужны
|
||||||
|
del media_group_bytes
|
||||||
|
|
||||||
images_bytes = []
|
images_bytes = []
|
||||||
if generated_images_io:
|
if generated_images_io:
|
||||||
for img_io in generated_images_io:
|
for img_io in generated_images_io:
|
||||||
# Читаем байты из BytesIO
|
|
||||||
img_io.seek(0)
|
img_io.seek(0)
|
||||||
content = img_io.read()
|
images_bytes.append(img_io.read())
|
||||||
images_bytes.append(content)
|
|
||||||
|
|
||||||
# Закрываем поток
|
|
||||||
img_io.close()
|
img_io.close()
|
||||||
|
# Освобождаем список BytesIO сразу
|
||||||
|
del generated_images_io
|
||||||
|
|
||||||
return images_bytes, metrics
|
return images_bytes, metrics
|
||||||
|
|
||||||
@@ -71,7 +79,7 @@ class GenerationService:
|
|||||||
self.bot = bot
|
self.bot = bot
|
||||||
|
|
||||||
|
|
||||||
async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str:
|
async def ask_prompt_assistant(self, prompt: str, assets: list[str] | None = None) -> str:
|
||||||
future_prompt = """You are an prompt-assistant. You improving user-entered prompts for image generation. User may upload reference image too.
|
future_prompt = """You are an prompt-assistant. You improving user-entered prompts for image generation. User may upload reference image too.
|
||||||
I will provide sources prompt entered by user. Understand user needs and generate best variation of prompt.
|
I will provide sources prompt entered by user. Understand user needs and generate best variation of prompt.
|
||||||
ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT: """
|
ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT: """
|
||||||
@@ -94,10 +102,9 @@ class GenerationService:
|
|||||||
|
|
||||||
return await asyncio.to_thread(self.gemini.generate_text, prompt=technical_prompt, images_list=images)
|
return await asyncio.to_thread(self.gemini.generate_text, prompt=technical_prompt, images_list=images)
|
||||||
|
|
||||||
async def get_generations(self, character_id: Optional[str] = None, limit: int = 10, offset: int = 0, user_id: Optional[str] = None, project_id: Optional[str] = None) -> List[
|
async def get_generations(self, character_id: Optional[str] = None, limit: int = 10, offset: int = 0, user_id: Optional[str] = None, project_id: Optional[str] = None, idea_id: Optional[str] = None) -> GenerationsResponse:
|
||||||
Generation]:
|
generations = await self.dao.generations.get_generations(character_id = character_id,limit=limit, offset=offset, created_by=user_id, project_id=project_id, idea_id=idea_id)
|
||||||
generations = await self.dao.generations.get_generations(character_id = character_id,limit=limit, offset=offset, created_by=user_id, project_id=project_id)
|
total_count = await self.dao.generations.count_generations(character_id = character_id, created_by=user_id, project_id=project_id, idea_id=idea_id)
|
||||||
total_count = await self.dao.generations.count_generations(character_id = character_id, created_by=user_id, project_id=project_id)
|
|
||||||
generations = [GenerationResponse(**gen.model_dump()) for gen in generations]
|
generations = [GenerationResponse(**gen.model_dump()) for gen in generations]
|
||||||
return GenerationsResponse(generations=generations, total_count=total_count)
|
return GenerationsResponse(generations=generations, total_count=total_count)
|
||||||
|
|
||||||
@@ -111,27 +118,51 @@ class GenerationService:
|
|||||||
async def get_running_generations(self, user_id: Optional[str] = None, project_id: Optional[str] = None) -> List[Generation]:
|
async def get_running_generations(self, user_id: Optional[str] = None, project_id: Optional[str] = None) -> List[Generation]:
|
||||||
return await self.dao.generations.get_generations(status=GenerationStatus.RUNNING, created_by=user_id, project_id=project_id)
|
return await self.dao.generations.get_generations(status=GenerationStatus.RUNNING, created_by=user_id, project_id=project_id)
|
||||||
|
|
||||||
async def create_generation_task(self, generation_request: GenerationRequest, user_id: Optional[str] = None) -> GenerationResponse:
|
async def create_generation_task(self, generation_request: GenerationRequest, user_id: Optional[str] = None, generation_group_id: Optional[str] = None) -> GenerationGroupResponse:
|
||||||
|
count = generation_request.count
|
||||||
|
|
||||||
|
if generation_group_id is None:
|
||||||
|
generation_group_id = str(uuid4())
|
||||||
|
|
||||||
|
results = []
|
||||||
|
for _ in range(count):
|
||||||
|
gen_response = await self._create_single_generation(generation_request, user_id, generation_group_id)
|
||||||
|
results.append(gen_response)
|
||||||
|
return GenerationGroupResponse(generation_group_id=generation_group_id, generations=results)
|
||||||
|
|
||||||
|
async def _create_single_generation(self, generation_request: GenerationRequest, user_id: Optional[str] = None, generation_group_id: Optional[str] = None) -> GenerationResponse:
|
||||||
gen_id = None
|
gen_id = None
|
||||||
generation_model = None
|
generation_model = None
|
||||||
|
|
||||||
|
if generation_request.environment_id and not generation_request.linked_character_id:
|
||||||
|
raise HTTPException(status_code=400, detail="environment_id can only be used when linked_character_id is provided")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
generation_model = Generation(**generation_request.model_dump())
|
generation_model = Generation(**generation_request.model_dump(exclude={'count'}))
|
||||||
if user_id:
|
if user_id:
|
||||||
generation_model.created_by = user_id
|
generation_model.created_by = user_id
|
||||||
|
if generation_group_id:
|
||||||
|
generation_model.generation_group_id = generation_group_id
|
||||||
|
|
||||||
|
# Explicitly set idea_id from request if present (already in model_dump, but ensuring clarity)
|
||||||
|
if generation_request.idea_id:
|
||||||
|
generation_model.idea_id = generation_request.idea_id
|
||||||
|
|
||||||
gen_id = await self.dao.generations.create_generation(generation_model)
|
gen_id = await self.dao.generations.create_generation(generation_model)
|
||||||
generation_model.id = gen_id
|
generation_model.id = gen_id
|
||||||
|
|
||||||
async def runner(gen):
|
async def runner(gen):
|
||||||
logger.info(f"Starting background generation task for ID: {gen.id}")
|
logger.info(f"Generation {gen.id} entered queue (waiting for slot)...")
|
||||||
try:
|
try:
|
||||||
|
async with generation_semaphore:
|
||||||
|
logger.info(f"Starting background generation task for ID: {gen.id}")
|
||||||
await self.create_generation(gen)
|
await self.create_generation(gen)
|
||||||
logger.info(f"Background generation task finished for ID: {gen.id}")
|
logger.info(f"Background generation task finished for ID: {gen.id}")
|
||||||
except Exception:
|
except Exception:
|
||||||
# если генерация уже пошла и упала — пометим FAILED
|
# если генерация уже пошла и упала — пометим FAILED
|
||||||
try:
|
try:
|
||||||
db_gen = await self.dao.generations.get_generation(gen.id)
|
db_gen = await self.dao.generations.get_generation(gen.id)
|
||||||
|
if db_gen is not None:
|
||||||
db_gen.status = GenerationStatus.FAILED
|
db_gen.status = GenerationStatus.FAILED
|
||||||
await self.dao.generations.update_generation(db_gen)
|
await self.dao.generations.update_generation(db_gen)
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -147,6 +178,7 @@ class GenerationService:
|
|||||||
if gen_id is not None:
|
if gen_id is not None:
|
||||||
try:
|
try:
|
||||||
gen = await self.dao.generations.get_generation(gen_id)
|
gen = await self.dao.generations.get_generation(gen_id)
|
||||||
|
if gen is not None:
|
||||||
gen.status = GenerationStatus.FAILED
|
gen.status = GenerationStatus.FAILED
|
||||||
await self.dao.generations.update_generation(gen)
|
await self.dao.generations.update_generation(gen)
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -158,42 +190,38 @@ class GenerationService:
|
|||||||
logger.info(f"Processing generation {generation.id}. Character ID: {generation.linked_character_id}")
|
logger.info(f"Processing generation {generation.id}. Character ID: {generation.linked_character_id}")
|
||||||
|
|
||||||
# 2. Получаем ассеты-референсы (если они есть)
|
# 2. Получаем ассеты-референсы (если они есть)
|
||||||
reference_assets: List[Asset] = []
|
|
||||||
media_group_bytes: List[bytes] = []
|
media_group_bytes: List[bytes] = []
|
||||||
generation_prompt = generation.prompt
|
generation_prompt = generation.prompt
|
||||||
# generation_prompt = f"""
|
|
||||||
|
|
||||||
# Create detailed image of character in scene.
|
# 2.1 Аватар персонажа (всегда первый, если включен)
|
||||||
|
|
||||||
# SCENE DESCRIPTION: {generation.prompt}
|
|
||||||
|
|
||||||
# Rules:
|
|
||||||
# - Integrate the character's appearance naturally into the scene description.
|
|
||||||
# - Focus on lighting, texture, and composition.
|
|
||||||
# """
|
|
||||||
if generation.linked_character_id is not None:
|
if generation.linked_character_id is not None:
|
||||||
char_info = await self.dao.chars.get_character(generation.linked_character_id)
|
char_info = await self.dao.chars.get_character(generation.linked_character_id)
|
||||||
if char_info is None:
|
if char_info is None:
|
||||||
raise Exception(f"Character ID {generation.linked_character_id} not found")
|
raise Exception(f"Character ID {generation.linked_character_id} not found")
|
||||||
if generation.use_profile_image:
|
|
||||||
|
if generation.use_profile_image and char_info.avatar_asset_id:
|
||||||
avatar_asset = await self.dao.assets.get_asset(char_info.avatar_asset_id)
|
avatar_asset = await self.dao.assets.get_asset(char_info.avatar_asset_id)
|
||||||
if avatar_asset:
|
if avatar_asset:
|
||||||
media_group_bytes.append(avatar_asset.data)
|
img_data = await self._get_asset_data(avatar_asset)
|
||||||
# generation_prompt = generation_prompt.replace("$char_bio_inserted", f"1. CHARACTER BIO (Must be strictly followed): {char_info.character_bio}")
|
if img_data:
|
||||||
|
media_group_bytes.append(img_data)
|
||||||
|
|
||||||
reference_assets = await self.dao.assets.get_assets_by_ids(generation.assets_list)
|
# 2.2 Явно указанные ассеты
|
||||||
|
if generation.assets_list:
|
||||||
# Извлекаем данные (bytes) из ассетов для отправки в Gemini
|
explicit_assets = await self.dao.assets.get_assets_by_ids(generation.assets_list)
|
||||||
for asset in reference_assets:
|
for asset in explicit_assets:
|
||||||
if asset.content_type != AssetContentType.IMAGE:
|
ref_asset_data = await self._get_asset_data(asset)
|
||||||
continue
|
if ref_asset_data:
|
||||||
|
media_group_bytes.append(ref_asset_data)
|
||||||
img_data = None
|
|
||||||
if asset.minio_object_name:
|
|
||||||
img_data = await self.s3_adapter.get_file(asset.minio_object_name)
|
|
||||||
elif asset.data:
|
|
||||||
img_data = asset.data
|
|
||||||
|
|
||||||
|
# 2.3 Ассеты из окружения (в самый конец)
|
||||||
|
if generation.environment_id:
|
||||||
|
env = await self.dao.environments.get_env(generation.environment_id)
|
||||||
|
if env and env.asset_ids:
|
||||||
|
logger.info(f"Loading {len(env.asset_ids)} assets from environment {env.name} ({env.id})")
|
||||||
|
env_assets = await self.dao.assets.get_assets_by_ids(env.asset_ids)
|
||||||
|
for asset in env_assets:
|
||||||
|
img_data = await self._get_asset_data(asset)
|
||||||
if img_data:
|
if img_data:
|
||||||
media_group_bytes.append(img_data)
|
media_group_bytes.append(img_data)
|
||||||
|
|
||||||
@@ -279,7 +307,9 @@ class GenerationService:
|
|||||||
|
|
||||||
# 5. (Опционально) Обновляем запись генерации ссылками на результаты
|
# 5. (Опционально) Обновляем запись генерации ссылками на результаты
|
||||||
# Предполагаем, что у модели Generation есть поле result_asset_ids
|
# Предполагаем, что у модели Generation есть поле result_asset_ids
|
||||||
result_ids = [a.id for a in created_assets]
|
result_ids = []
|
||||||
|
for a in created_assets:
|
||||||
|
result_ids.append(a.id)
|
||||||
|
|
||||||
generation.result_list = result_ids
|
generation.result_list = result_ids
|
||||||
generation.status = GenerationStatus.DONE
|
generation.status = GenerationStatus.DONE
|
||||||
@@ -310,6 +340,14 @@ class GenerationService:
|
|||||||
logger.error(f"Failed to send assets to Telegram ID {generation.telegram_id}: {e}")
|
logger.error(f"Failed to send assets to Telegram ID {generation.telegram_id}: {e}")
|
||||||
|
|
||||||
|
|
||||||
|
async def _get_asset_data(self, asset: Asset) -> Optional[bytes]:
|
||||||
|
if asset.content_type != AssetContentType.IMAGE:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if asset.minio_object_name:
|
||||||
|
return await self.s3_adapter.get_file(asset.minio_object_name)
|
||||||
|
return asset.data
|
||||||
|
|
||||||
async def _simulate_progress(self, generation: Generation):
|
async def _simulate_progress(self, generation: Generation):
|
||||||
"""
|
"""
|
||||||
Increments progress from 0 to 90 over ~20 seconds.
|
Increments progress from 0 to 90 over ~20 seconds.
|
||||||
@@ -347,7 +385,6 @@ class GenerationService:
|
|||||||
Returns:
|
Returns:
|
||||||
Created Generation object
|
Created Generation object
|
||||||
"""
|
"""
|
||||||
from api.models.ExternalGenerationDTO import ExternalGenerationRequest
|
|
||||||
|
|
||||||
# Validate image source
|
# Validate image source
|
||||||
external_gen.validate_image_source()
|
external_gen.validate_image_source()
|
||||||
@@ -444,3 +481,61 @@ class GenerationService:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error deleting generation {generation_id}: {e}")
|
logger.error(f"Error deleting generation {generation_id}: {e}")
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
async def cleanup_stale_generations(self):
|
||||||
|
"""
|
||||||
|
Cancels generations that have been running for more than 1 hour.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
count = await self.dao.generations.cancel_stale_generations(timeout_minutes=60)
|
||||||
|
if count > 0:
|
||||||
|
logger.info(f"Cleaned up {count} stale generations (timeout)")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error cleaning up stale generations: {e}")
|
||||||
|
|
||||||
|
async def cleanup_old_data(self, days: int = 2):
|
||||||
|
"""
|
||||||
|
Очистка старых данных:
|
||||||
|
1. Мягко удаляет генерации старше N дней
|
||||||
|
2. Мягко удаляет связанные ассеты + жёстко удаляет файлы из S3
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# 1. Мягко удаляем генерации и собираем asset IDs
|
||||||
|
gen_count, asset_ids = await self.dao.generations.soft_delete_old_generations(days=days)
|
||||||
|
|
||||||
|
if gen_count > 0:
|
||||||
|
logger.info(f"Soft-deleted {gen_count} generations older than {days} days. "
|
||||||
|
f"Found {len(asset_ids)} associated asset IDs.")
|
||||||
|
|
||||||
|
# 2. Мягко удаляем ассеты + жёстко удаляем файлы из S3
|
||||||
|
if asset_ids:
|
||||||
|
purged = await self.dao.assets.soft_delete_and_purge_assets(asset_ids)
|
||||||
|
logger.info(f"Purged {purged} assets (soft-deleted + S3 files removed).")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error during old data cleanup: {e}")
|
||||||
|
|
||||||
|
async def get_financial_report(self, user_id: Optional[str] = None, project_id: Optional[str] = None, breakdown_by: Optional[str] = None) -> FinancialReport:
|
||||||
|
"""
|
||||||
|
Generates a financial usage report for a specific user or project.
|
||||||
|
'breakdown_by' can be 'created_by' or 'project_id'.
|
||||||
|
"""
|
||||||
|
summary_data = await self.dao.generations.get_usage_stats(created_by=user_id, project_id=project_id)
|
||||||
|
summary = UsageStats(**summary_data)
|
||||||
|
|
||||||
|
by_user = None
|
||||||
|
by_project = None
|
||||||
|
|
||||||
|
if breakdown_by == "created_by":
|
||||||
|
res = await self.dao.generations.get_usage_breakdown(group_by="created_by", project_id=project_id, created_by=user_id)
|
||||||
|
by_user = [UsageByEntity(**item) for item in res]
|
||||||
|
|
||||||
|
if breakdown_by == "project_id":
|
||||||
|
res = await self.dao.generations.get_usage_breakdown(group_by="project_id", project_id=project_id, created_by=user_id)
|
||||||
|
by_project = [UsageByEntity(**item) for item in res]
|
||||||
|
|
||||||
|
return FinancialReport(
|
||||||
|
summary=summary,
|
||||||
|
by_user=by_user,
|
||||||
|
by_project=by_project
|
||||||
|
)
|
||||||
75
api/service/idea_service.py
Normal file
75
api/service/idea_service.py
Normal file
@@ -0,0 +1,75 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from datetime import datetime
|
||||||
|
from repos.dao import DAO
|
||||||
|
from models.Idea import Idea
|
||||||
|
|
||||||
|
class IdeaService:
|
||||||
|
def __init__(self, dao: DAO):
|
||||||
|
self.dao = dao
|
||||||
|
|
||||||
|
async def create_idea(self, name: str, description: Optional[str], project_id: Optional[str], user_id: str) -> Idea:
|
||||||
|
idea = Idea(name=name, description=description, project_id=project_id, created_by=user_id)
|
||||||
|
idea_id = await self.dao.ideas.create_idea(idea)
|
||||||
|
idea.id = idea_id
|
||||||
|
return idea
|
||||||
|
|
||||||
|
async def get_ideas(self, project_id: Optional[str], user_id: str, limit: int = 20, offset: int = 0) -> List[dict]:
|
||||||
|
return await self.dao.ideas.get_ideas(project_id, user_id, limit, offset)
|
||||||
|
|
||||||
|
async def get_idea(self, idea_id: str) -> Optional[Idea]:
|
||||||
|
return await self.dao.ideas.get_idea(idea_id)
|
||||||
|
|
||||||
|
async def update_idea(self, idea_id: str, name: Optional[str] = None, description: Optional[str] = None) -> Optional[Idea]:
|
||||||
|
idea = await self.dao.ideas.get_idea(idea_id)
|
||||||
|
if not idea:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if name is not None:
|
||||||
|
idea.name = name
|
||||||
|
if description is not None:
|
||||||
|
idea.description = description
|
||||||
|
|
||||||
|
idea.updated_at = datetime.now()
|
||||||
|
await self.dao.ideas.update_idea(idea)
|
||||||
|
return idea
|
||||||
|
|
||||||
|
async def delete_idea(self, idea_id: str) -> bool:
|
||||||
|
return await self.dao.ideas.delete_idea(idea_id)
|
||||||
|
|
||||||
|
async def add_generation_to_idea(self, idea_id: str, generation_id: str) -> bool:
|
||||||
|
# Verify idea exists
|
||||||
|
idea = await self.dao.ideas.get_idea(idea_id)
|
||||||
|
if not idea:
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Get generation
|
||||||
|
gen = await self.dao.generations.get_generation(generation_id)
|
||||||
|
if not gen:
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Link
|
||||||
|
gen.idea_id = idea_id
|
||||||
|
gen.updated_at = datetime.now()
|
||||||
|
await self.dao.generations.update_generation(gen)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def remove_generation_from_idea(self, idea_id: str, generation_id: str) -> bool:
|
||||||
|
# Verify idea exists (optional, but good for validation)
|
||||||
|
idea = await self.dao.ideas.get_idea(idea_id)
|
||||||
|
if not idea:
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Get generation
|
||||||
|
gen = await self.dao.generations.get_generation(generation_id)
|
||||||
|
if not gen:
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Unlink only if currently linked to this idea
|
||||||
|
if gen.idea_id == idea_id:
|
||||||
|
gen.idea_id = None
|
||||||
|
gen.updated_at = datetime.now()
|
||||||
|
await self.dao.generations.update_generation(gen)
|
||||||
|
return True
|
||||||
|
|
||||||
|
return False
|
||||||
|
|
||||||
79
api/service/post_service.py
Normal file
79
api/service/post_service.py
Normal file
@@ -0,0 +1,79 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from datetime import datetime, UTC
|
||||||
|
|
||||||
|
from repos.dao import DAO
|
||||||
|
from models.Post import Post
|
||||||
|
|
||||||
|
|
||||||
|
class PostService:
|
||||||
|
def __init__(self, dao: DAO):
|
||||||
|
self.dao = dao
|
||||||
|
|
||||||
|
async def create_post(
|
||||||
|
self,
|
||||||
|
date: datetime,
|
||||||
|
topic: str,
|
||||||
|
generation_ids: List[str],
|
||||||
|
project_id: Optional[str],
|
||||||
|
user_id: str,
|
||||||
|
) -> Post:
|
||||||
|
post = Post(
|
||||||
|
date=date,
|
||||||
|
topic=topic,
|
||||||
|
generation_ids=generation_ids,
|
||||||
|
project_id=project_id,
|
||||||
|
created_by=user_id,
|
||||||
|
)
|
||||||
|
post_id = await self.dao.posts.create_post(post)
|
||||||
|
post.id = post_id
|
||||||
|
return post
|
||||||
|
|
||||||
|
async def get_post(self, post_id: str) -> Optional[Post]:
|
||||||
|
return await self.dao.posts.get_post(post_id)
|
||||||
|
|
||||||
|
async def get_posts(
|
||||||
|
self,
|
||||||
|
project_id: Optional[str],
|
||||||
|
user_id: str,
|
||||||
|
limit: int = 20,
|
||||||
|
offset: int = 0,
|
||||||
|
date_from: Optional[datetime] = None,
|
||||||
|
date_to: Optional[datetime] = None,
|
||||||
|
) -> List[Post]:
|
||||||
|
return await self.dao.posts.get_posts(project_id, user_id, limit, offset, date_from, date_to)
|
||||||
|
|
||||||
|
async def update_post(
|
||||||
|
self,
|
||||||
|
post_id: str,
|
||||||
|
date: Optional[datetime] = None,
|
||||||
|
topic: Optional[str] = None,
|
||||||
|
) -> Optional[Post]:
|
||||||
|
post = await self.dao.posts.get_post(post_id)
|
||||||
|
if not post:
|
||||||
|
return None
|
||||||
|
|
||||||
|
updates: dict = {"updated_at": datetime.now(UTC)}
|
||||||
|
if date is not None:
|
||||||
|
updates["date"] = date
|
||||||
|
if topic is not None:
|
||||||
|
updates["topic"] = topic
|
||||||
|
|
||||||
|
await self.dao.posts.update_post(post_id, updates)
|
||||||
|
|
||||||
|
# Return refreshed post
|
||||||
|
return await self.dao.posts.get_post(post_id)
|
||||||
|
|
||||||
|
async def delete_post(self, post_id: str) -> bool:
|
||||||
|
return await self.dao.posts.delete_post(post_id)
|
||||||
|
|
||||||
|
async def add_generations(self, post_id: str, generation_ids: List[str]) -> bool:
|
||||||
|
post = await self.dao.posts.get_post(post_id)
|
||||||
|
if not post:
|
||||||
|
return False
|
||||||
|
return await self.dao.posts.add_generations(post_id, generation_ids)
|
||||||
|
|
||||||
|
async def remove_generation(self, post_id: str, generation_id: str) -> bool:
|
||||||
|
post = await self.dao.posts.get_post(post_id)
|
||||||
|
if not post:
|
||||||
|
return False
|
||||||
|
return await self.dao.posts.remove_generation(post_id, generation_id)
|
||||||
39
config.py
Normal file
39
config.py
Normal file
@@ -0,0 +1,39 @@
|
|||||||
|
import os
|
||||||
|
from typing import Optional
|
||||||
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||||
|
|
||||||
|
|
||||||
|
class Settings(BaseSettings):
|
||||||
|
# Telegram Bot
|
||||||
|
BOT_TOKEN: str
|
||||||
|
ADMIN_ID: int = 0
|
||||||
|
|
||||||
|
# AI Service
|
||||||
|
GEMINI_API_KEY: str
|
||||||
|
|
||||||
|
# Database
|
||||||
|
MONGO_HOST: str = "mongodb://localhost:27017"
|
||||||
|
DB_NAME: str = "my_bot_db"
|
||||||
|
|
||||||
|
# S3 Storage (Minio)
|
||||||
|
MINIO_ENDPOINT: str = "http://localhost:9000"
|
||||||
|
MINIO_ACCESS_KEY: str = "minioadmin"
|
||||||
|
MINIO_SECRET_KEY: str = "minioadmin"
|
||||||
|
MINIO_BUCKET: str = "ai-char"
|
||||||
|
|
||||||
|
# External API
|
||||||
|
EXTERNAL_API_SECRET: Optional[str] = None
|
||||||
|
|
||||||
|
# JWT Security
|
||||||
|
SECRET_KEY: str = "CHANGE_ME_TO_A_SUPER_SECRET_KEY"
|
||||||
|
ALGORITHM: str = "HS256"
|
||||||
|
ACCESS_TOKEN_EXPIRE_MINUTES: int = 30 * 24 * 60 # 30 days
|
||||||
|
|
||||||
|
model_config = SettingsConfigDict(
|
||||||
|
env_file=os.getenv("ENV_FILE", ".env"),
|
||||||
|
env_file_encoding="utf-8",
|
||||||
|
extra="ignore"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
settings = Settings()
|
||||||
@@ -27,19 +27,19 @@ class AlbumMiddleware(BaseMiddleware):
|
|||||||
# Ждем сбора остальных частей
|
# Ждем сбора остальных частей
|
||||||
await asyncio.sleep(self.latency)
|
await asyncio.sleep(self.latency)
|
||||||
|
|
||||||
# Проверяем, что ключ все еще существует (на всякий случай)
|
# Проверяем, что ключ все еще существует
|
||||||
if group_id in self.album_data:
|
if group_id in self.album_data:
|
||||||
# Передаем собранный альбом в хендлер
|
# Передаем собранный альбом в хендлер
|
||||||
# Сортируем по message_id, чтобы порядок был верным
|
# Сортируем по message_id, чтобы порядок был верным
|
||||||
self.album_data[group_id].sort(key=lambda x: x.message_id)
|
current_album = self.album_data[group_id]
|
||||||
data["album"] = self.album_data[group_id]
|
current_album.sort(key=lambda x: x.message_id)
|
||||||
|
data["album"] = current_album
|
||||||
return await handler(event, data)
|
return await handler(event, data)
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
# ЧИСТКА: Удаляем всегда, если это "головной" поток, который создал запись
|
# ЧИСТКА: Удаляем запись после обработки или таймаута
|
||||||
# Проверяем, что мы удаляем именно то, что создали, и ключ существует
|
# Используем pop() с дефолтом, чтобы избежать KeyError
|
||||||
if group_id in self.album_data and self.album_data[group_id][0] == event:
|
self.album_data.pop(group_id, None)
|
||||||
del self.album_data[group_id]
|
|
||||||
|
|
||||||
else:
|
else:
|
||||||
# Если группа уже собирается - просто добавляем и выходим
|
# Если группа уже собирается - просто добавляем и выходим
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ class Asset(BaseModel):
|
|||||||
tags: List[str] = []
|
tags: List[str] = []
|
||||||
created_by: Optional[str] = None
|
created_by: Optional[str] = None
|
||||||
project_id: Optional[str] = None
|
project_id: Optional[str] = None
|
||||||
|
is_deleted: bool = False
|
||||||
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
|
|
||||||
@@ -62,6 +63,7 @@ class Asset(BaseModel):
|
|||||||
|
|
||||||
# --- CALCULATED FIELD ---
|
# --- CALCULATED FIELD ---
|
||||||
@computed_field
|
@computed_field
|
||||||
|
@property
|
||||||
def url(self) -> str:
|
def url(self) -> str:
|
||||||
"""
|
"""
|
||||||
Это поле автоматически вычислится и попадет в model_dump() / .json()
|
Это поле автоматически вычислится и попадет в model_dump() / .json()
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ class Character(BaseModel):
|
|||||||
name: str
|
name: str
|
||||||
avatar_asset_id: Optional[str] = None
|
avatar_asset_id: Optional[str] = None
|
||||||
avatar_image: Optional[str] = None
|
avatar_image: Optional[str] = None
|
||||||
character_image_data: Optional[bytes] = None
|
|
||||||
character_image_doc_tg_id: Optional[str] = None
|
character_image_doc_tg_id: Optional[str] = None
|
||||||
character_image_tg_id: Optional[str] = None
|
character_image_tg_id: Optional[str] = None
|
||||||
character_bio: Optional[str] = None
|
character_bio: Optional[str] = None
|
||||||
|
|||||||
20
models/Environment.py
Normal file
20
models/Environment.py
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from pydantic import BaseModel, Field, ConfigDict
|
||||||
|
from datetime import datetime
|
||||||
|
from bson import ObjectId
|
||||||
|
|
||||||
|
|
||||||
|
class Environment(BaseModel):
|
||||||
|
id: Optional[str] = Field(None, alias="_id")
|
||||||
|
character_id: str
|
||||||
|
name: str = Field(..., min_length=1)
|
||||||
|
description: Optional[str] = None
|
||||||
|
asset_ids: List[str] = Field(default_factory=list)
|
||||||
|
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
updated_at: datetime = Field(default_factory=datetime.utcnow)
|
||||||
|
|
||||||
|
model_config = ConfigDict(
|
||||||
|
populate_by_name=True,
|
||||||
|
json_encoders={ObjectId: str},
|
||||||
|
arbitrary_types_allowed=True
|
||||||
|
)
|
||||||
@@ -35,8 +35,11 @@ class Generation(BaseModel):
|
|||||||
output_token_usage: Optional[int] = None
|
output_token_usage: Optional[int] = None
|
||||||
is_deleted: bool = False
|
is_deleted: bool = False
|
||||||
album_id: Optional[str] = None
|
album_id: Optional[str] = None
|
||||||
|
environment_id: Optional[str] = None
|
||||||
|
generation_group_id: Optional[str] = None
|
||||||
created_by: Optional[str] = None # Stores User ID (Telegram ID or Web User ObjectId)
|
created_by: Optional[str] = None # Stores User ID (Telegram ID or Web User ObjectId)
|
||||||
project_id: Optional[str] = None
|
project_id: Optional[str] = None
|
||||||
|
idea_id: Optional[str] = None
|
||||||
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
|
|
||||||
|
|||||||
13
models/Idea.py
Normal file
13
models/Idea.py
Normal file
@@ -0,0 +1,13 @@
|
|||||||
|
from datetime import datetime
|
||||||
|
from typing import Optional, List
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
class Idea(BaseModel):
|
||||||
|
id: Optional[str] = None
|
||||||
|
name: str = "New Idea"
|
||||||
|
description: Optional[str] = None
|
||||||
|
project_id: Optional[str] = None
|
||||||
|
created_by: str # User ID
|
||||||
|
is_deleted: bool = False
|
||||||
|
created_at: datetime = Field(default_factory=datetime.now)
|
||||||
|
updated_at: datetime = Field(default_factory=datetime.now)
|
||||||
23
models/Post.py
Normal file
23
models/Post.py
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
from datetime import datetime, timezone, UTC
|
||||||
|
from typing import Optional, List
|
||||||
|
from pydantic import BaseModel, Field, model_validator
|
||||||
|
|
||||||
|
|
||||||
|
class Post(BaseModel):
|
||||||
|
id: Optional[str] = None
|
||||||
|
date: datetime
|
||||||
|
topic: str
|
||||||
|
generation_ids: List[str] = Field(default_factory=list)
|
||||||
|
project_id: Optional[str] = None
|
||||||
|
created_by: str
|
||||||
|
is_deleted: bool = False
|
||||||
|
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
|
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def ensure_tz_aware(self):
|
||||||
|
for field in ("date", "created_at", "updated_at"):
|
||||||
|
val = getattr(self, field)
|
||||||
|
if val is not None and val.tzinfo is None:
|
||||||
|
setattr(self, field, val.replace(tzinfo=timezone.utc))
|
||||||
|
return self
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,6 +1,8 @@
|
|||||||
from typing import List, Optional
|
from typing import Any, List, Optional
|
||||||
import logging
|
import logging
|
||||||
|
from datetime import datetime, UTC
|
||||||
from bson import ObjectId
|
from bson import ObjectId
|
||||||
|
from uuid import uuid4
|
||||||
from motor.motor_asyncio import AsyncIOMotorClient
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
|
||||||
from models.Asset import Asset
|
from models.Asset import Asset
|
||||||
@@ -19,7 +21,8 @@ class AssetsRepo:
|
|||||||
# Main data
|
# Main data
|
||||||
if asset.data:
|
if asset.data:
|
||||||
ts = int(asset.created_at.timestamp())
|
ts = int(asset.created_at.timestamp())
|
||||||
object_name = f"{asset.type.value}/{ts}_{asset.name}"
|
uid = uuid4().hex[:8]
|
||||||
|
object_name = f"{asset.type.value}/{ts}_{uid}_{asset.name}"
|
||||||
|
|
||||||
uploaded = await self.s3.upload_file(object_name, asset.data)
|
uploaded = await self.s3.upload_file(object_name, asset.data)
|
||||||
if uploaded:
|
if uploaded:
|
||||||
@@ -32,7 +35,8 @@ class AssetsRepo:
|
|||||||
# Thumbnail
|
# Thumbnail
|
||||||
if asset.thumbnail:
|
if asset.thumbnail:
|
||||||
ts = int(asset.created_at.timestamp())
|
ts = int(asset.created_at.timestamp())
|
||||||
thumb_name = f"{asset.type.value}/thumbs/{ts}_{asset.name}_thumb.jpg"
|
uid = uuid4().hex[:8]
|
||||||
|
thumb_name = f"{asset.type.value}/thumbs/{ts}_{uid}_{asset.name}_thumb.jpg"
|
||||||
|
|
||||||
uploaded_thumb = await self.s3.upload_file(thumb_name, asset.thumbnail)
|
uploaded_thumb = await self.s3.upload_file(thumb_name, asset.thumbnail)
|
||||||
if uploaded_thumb:
|
if uploaded_thumb:
|
||||||
@@ -47,7 +51,7 @@ class AssetsRepo:
|
|||||||
return str(res.inserted_id)
|
return str(res.inserted_id)
|
||||||
|
|
||||||
async def get_assets(self, asset_type: Optional[str] = None, limit: int = 10, offset: int = 0, with_data: bool = False, created_by: Optional[str] = None, project_id: Optional[str] = None) -> List[Asset]:
|
async def get_assets(self, asset_type: Optional[str] = None, limit: int = 10, offset: int = 0, with_data: bool = False, created_by: Optional[str] = None, project_id: Optional[str] = None) -> List[Asset]:
|
||||||
filter = {}
|
filter: dict[str, Any]= {"is_deleted": {"$ne": True}}
|
||||||
if asset_type:
|
if asset_type:
|
||||||
filter["type"] = asset_type
|
filter["type"] = asset_type
|
||||||
args = {}
|
args = {}
|
||||||
@@ -134,7 +138,8 @@ class AssetsRepo:
|
|||||||
if self.s3:
|
if self.s3:
|
||||||
if asset.data:
|
if asset.data:
|
||||||
ts = int(asset.created_at.timestamp())
|
ts = int(asset.created_at.timestamp())
|
||||||
object_name = f"{asset.type.value}/{ts}_{asset.name}"
|
uid = uuid4().hex[:8]
|
||||||
|
object_name = f"{asset.type.value}/{ts}_{uid}_{asset.name}"
|
||||||
if await self.s3.upload_file(object_name, asset.data):
|
if await self.s3.upload_file(object_name, asset.data):
|
||||||
asset.minio_object_name = object_name
|
asset.minio_object_name = object_name
|
||||||
asset.minio_bucket = self.s3.bucket_name
|
asset.minio_bucket = self.s3.bucket_name
|
||||||
@@ -142,7 +147,8 @@ class AssetsRepo:
|
|||||||
|
|
||||||
if asset.thumbnail:
|
if asset.thumbnail:
|
||||||
ts = int(asset.created_at.timestamp())
|
ts = int(asset.created_at.timestamp())
|
||||||
thumb_name = f"{asset.type.value}/thumbs/{ts}_{asset.name}_thumb.jpg"
|
uid = uuid4().hex[:8]
|
||||||
|
thumb_name = f"{asset.type.value}/thumbs/{ts}_{uid}_{asset.name}_thumb.jpg"
|
||||||
if await self.s3.upload_file(thumb_name, asset.thumbnail):
|
if await self.s3.upload_file(thumb_name, asset.thumbnail):
|
||||||
asset.minio_thumbnail_object_name = thumb_name
|
asset.minio_thumbnail_object_name = thumb_name
|
||||||
asset.thumbnail = None
|
asset.thumbnail = None
|
||||||
@@ -169,6 +175,8 @@ class AssetsRepo:
|
|||||||
filter["linked_char_id"] = character_id
|
filter["linked_char_id"] = character_id
|
||||||
if created_by:
|
if created_by:
|
||||||
filter["created_by"] = created_by
|
filter["created_by"] = created_by
|
||||||
|
if project_id is None:
|
||||||
|
filter["project_id"] = None
|
||||||
if project_id:
|
if project_id:
|
||||||
filter["project_id"] = project_id
|
filter["project_id"] = project_id
|
||||||
return await self.collection.count_documents(filter)
|
return await self.collection.count_documents(filter)
|
||||||
@@ -197,6 +205,61 @@ class AssetsRepo:
|
|||||||
res = await self.collection.delete_one({"_id": ObjectId(asset_id)})
|
res = await self.collection.delete_one({"_id": ObjectId(asset_id)})
|
||||||
return res.deleted_count > 0
|
return res.deleted_count > 0
|
||||||
|
|
||||||
|
async def soft_delete_and_purge_assets(self, asset_ids: List[str]) -> int:
|
||||||
|
"""
|
||||||
|
Мягко удаляет ассеты и жёстко удаляет их файлы из S3.
|
||||||
|
Возвращает количество обработанных ассетов.
|
||||||
|
"""
|
||||||
|
if not asset_ids:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
object_ids = [ObjectId(aid) for aid in asset_ids if ObjectId.is_valid(aid)]
|
||||||
|
if not object_ids:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
# Находим ассеты, которые ещё не удалены
|
||||||
|
cursor = self.collection.find(
|
||||||
|
{"_id": {"$in": object_ids}, "is_deleted": {"$ne": True}},
|
||||||
|
{"minio_object_name": 1, "minio_thumbnail_object_name": 1}
|
||||||
|
)
|
||||||
|
|
||||||
|
purged_count = 0
|
||||||
|
ids_to_update = []
|
||||||
|
|
||||||
|
async for doc in cursor:
|
||||||
|
ids_to_update.append(doc["_id"])
|
||||||
|
|
||||||
|
# Жёсткое удаление файлов из S3
|
||||||
|
if self.s3:
|
||||||
|
if doc.get("minio_object_name"):
|
||||||
|
try:
|
||||||
|
await self.s3.delete_file(doc["minio_object_name"])
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to delete S3 object {doc['minio_object_name']}: {e}")
|
||||||
|
if doc.get("minio_thumbnail_object_name"):
|
||||||
|
try:
|
||||||
|
await self.s3.delete_file(doc["minio_thumbnail_object_name"])
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to delete S3 thumbnail {doc['minio_thumbnail_object_name']}: {e}")
|
||||||
|
|
||||||
|
purged_count += 1
|
||||||
|
|
||||||
|
# Мягкое удаление + очистка ссылок на S3
|
||||||
|
if ids_to_update:
|
||||||
|
await self.collection.update_many(
|
||||||
|
{"_id": {"$in": ids_to_update}},
|
||||||
|
{
|
||||||
|
"$set": {
|
||||||
|
"is_deleted": True,
|
||||||
|
"minio_object_name": None,
|
||||||
|
"minio_thumbnail_object_name": None,
|
||||||
|
"updated_at": datetime.now(UTC)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return purged_count
|
||||||
|
|
||||||
async def migrate_to_minio(self) -> dict:
|
async def migrate_to_minio(self) -> dict:
|
||||||
"""Переносит данные и thumbnails из Mongo в MinIO."""
|
"""Переносит данные и thumbnails из Mongo в MinIO."""
|
||||||
if not self.s3:
|
if not self.s3:
|
||||||
@@ -216,7 +279,8 @@ class AssetsRepo:
|
|||||||
created_at = doc.get("created_at")
|
created_at = doc.get("created_at")
|
||||||
ts = int(created_at.timestamp()) if created_at else 0
|
ts = int(created_at.timestamp()) if created_at else 0
|
||||||
|
|
||||||
object_name = f"{type_}/{ts}_{asset_id}_{name}"
|
uid = uuid4().hex[:8]
|
||||||
|
object_name = f"{type_}/{ts}_{uid}_{asset_id}_{name}"
|
||||||
if await self.s3.upload_file(object_name, data):
|
if await self.s3.upload_file(object_name, data):
|
||||||
await self.collection.update_one(
|
await self.collection.update_one(
|
||||||
{"_id": asset_id},
|
{"_id": asset_id},
|
||||||
@@ -243,7 +307,8 @@ class AssetsRepo:
|
|||||||
created_at = doc.get("created_at")
|
created_at = doc.get("created_at")
|
||||||
ts = int(created_at.timestamp()) if created_at else 0
|
ts = int(created_at.timestamp()) if created_at else 0
|
||||||
|
|
||||||
thumb_name = f"{type_}/thumbs/{ts}_{asset_id}_{name}_thumb.jpg"
|
uid = uuid4().hex[:8]
|
||||||
|
thumb_name = f"{type_}/thumbs/{ts}_{uid}_{asset_id}_{name}_thumb.jpg"
|
||||||
if await self.s3.upload_file(thumb_name, thumb):
|
if await self.s3.upload_file(thumb_name, thumb):
|
||||||
await self.collection.update_one(
|
await self.collection.update_one(
|
||||||
{"_id": asset_id},
|
{"_id": asset_id},
|
||||||
|
|||||||
@@ -15,26 +15,24 @@ class CharacterRepo:
|
|||||||
character.id = str(op.inserted_id)
|
character.id = str(op.inserted_id)
|
||||||
return character
|
return character
|
||||||
|
|
||||||
async def get_character(self, character_id: str, with_image_data: bool = False) -> Character | None:
|
async def get_character(self, character_id: str) -> Character | None:
|
||||||
args = {}
|
res = await self.collection.find_one({"_id": ObjectId(character_id)})
|
||||||
if not with_image_data:
|
|
||||||
args["character_image_data"] = 0
|
|
||||||
res = await self.collection.find_one({"_id": ObjectId(character_id)}, args)
|
|
||||||
if res is None:
|
if res is None:
|
||||||
return None
|
return None
|
||||||
else:
|
else:
|
||||||
res["id"] = str(res.pop("_id"))
|
res["id"] = str(res.pop("_id"))
|
||||||
return Character(**res)
|
return Character(**res)
|
||||||
|
|
||||||
async def get_all_characters(self, created_by: Optional[str] = None, project_id: Optional[str] = None) -> List[Character]:
|
async def get_all_characters(self, created_by: Optional[str] = None, project_id: Optional[str] = None, limit: int = 100, offset: int = 0) -> List[Character]:
|
||||||
filter = {}
|
filter = {}
|
||||||
if created_by:
|
if created_by:
|
||||||
filter["created_by"] = created_by
|
filter["created_by"] = created_by
|
||||||
|
if project_id is None:
|
||||||
|
filter["project_id"] = None
|
||||||
if project_id:
|
if project_id:
|
||||||
filter["project_id"] = project_id
|
filter["project_id"] = project_id
|
||||||
|
|
||||||
args = {"character_image_data": 0} # don't return image data for list
|
res = await self.collection.find(filter).skip(offset).limit(limit).to_list(None)
|
||||||
res = await self.collection.find(filter, args).to_list(None)
|
|
||||||
chars = []
|
chars = []
|
||||||
for doc in res:
|
for doc in res:
|
||||||
doc["id"] = str(doc.pop("_id"))
|
doc["id"] = str(doc.pop("_id"))
|
||||||
|
|||||||
@@ -6,6 +6,9 @@ from repos.generation_repo import GenerationRepo
|
|||||||
from repos.user_repo import UsersRepo
|
from repos.user_repo import UsersRepo
|
||||||
from repos.albums_repo import AlbumsRepo
|
from repos.albums_repo import AlbumsRepo
|
||||||
from repos.project_repo import ProjectRepo
|
from repos.project_repo import ProjectRepo
|
||||||
|
from repos.idea_repo import IdeaRepo
|
||||||
|
from repos.post_repo import PostRepo
|
||||||
|
from repos.environment_repo import EnvironmentRepo
|
||||||
|
|
||||||
|
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
@@ -19,3 +22,6 @@ class DAO:
|
|||||||
self.albums = AlbumsRepo(client, db_name)
|
self.albums = AlbumsRepo(client, db_name)
|
||||||
self.projects = ProjectRepo(client, db_name)
|
self.projects = ProjectRepo(client, db_name)
|
||||||
self.users = UsersRepo(client, db_name)
|
self.users = UsersRepo(client, db_name)
|
||||||
|
self.ideas = IdeaRepo(client, db_name)
|
||||||
|
self.posts = PostRepo(client, db_name)
|
||||||
|
self.environments = EnvironmentRepo(client, db_name)
|
||||||
|
|||||||
73
repos/environment_repo.py
Normal file
73
repos/environment_repo.py
Normal file
@@ -0,0 +1,73 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from datetime import datetime
|
||||||
|
from bson import ObjectId
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
from models.Environment import Environment
|
||||||
|
|
||||||
|
|
||||||
|
class EnvironmentRepo:
|
||||||
|
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
||||||
|
self.collection = client[db_name]["environments"]
|
||||||
|
|
||||||
|
async def create_env(self, env: Environment) -> Environment:
|
||||||
|
env_dict = env.model_dump(exclude={"id"})
|
||||||
|
res = await self.collection.insert_one(env_dict)
|
||||||
|
env.id = str(res.inserted_id)
|
||||||
|
return env
|
||||||
|
|
||||||
|
async def get_env(self, env_id: str) -> Optional[Environment]:
|
||||||
|
res = await self.collection.find_one({"_id": ObjectId(env_id)})
|
||||||
|
if not res:
|
||||||
|
return None
|
||||||
|
res["id"] = str(res.pop("_id"))
|
||||||
|
return Environment(**res)
|
||||||
|
|
||||||
|
async def get_character_envs(self, character_id: str) -> List[Environment]:
|
||||||
|
cursor = self.collection.find({"character_id": character_id})
|
||||||
|
envs = []
|
||||||
|
async for doc in cursor:
|
||||||
|
doc["id"] = str(doc.pop("_id"))
|
||||||
|
envs.append(Environment(**doc))
|
||||||
|
return envs
|
||||||
|
|
||||||
|
async def update_env(self, env_id: str, update_data: dict) -> bool:
|
||||||
|
update_data["updated_at"] = datetime.utcnow()
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(env_id)},
|
||||||
|
{"$set": update_data}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def delete_env(self, env_id: str) -> bool:
|
||||||
|
res = await self.collection.delete_one({"_id": ObjectId(env_id)})
|
||||||
|
return res.deleted_count > 0
|
||||||
|
|
||||||
|
async def add_asset(self, env_id: str, asset_id: str) -> bool:
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(env_id)},
|
||||||
|
{
|
||||||
|
"$addToSet": {"asset_ids": asset_id},
|
||||||
|
"$set": {"updated_at": datetime.utcnow()}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def add_assets(self, env_id: str, asset_ids: List[str]) -> bool:
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(env_id)},
|
||||||
|
{
|
||||||
|
"$addToSet": {"asset_ids": {"$each": asset_ids}},
|
||||||
|
"$set": {"updated_at": datetime.utcnow()}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def remove_asset(self, env_id: str, asset_id: str) -> bool:
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(env_id)},
|
||||||
|
{
|
||||||
|
"$pull": {"asset_ids": asset_id},
|
||||||
|
"$set": {"updated_at": datetime.utcnow()}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
@@ -1,4 +1,5 @@
|
|||||||
from typing import Optional, List
|
from typing import Any, Optional, List
|
||||||
|
from datetime import datetime, timedelta, UTC
|
||||||
|
|
||||||
from PIL.ImageChops import offset
|
from PIL.ImageChops import offset
|
||||||
from bson import ObjectId
|
from bson import ObjectId
|
||||||
@@ -16,7 +17,7 @@ class GenerationRepo:
|
|||||||
res = await self.collection.insert_one(generation.model_dump())
|
res = await self.collection.insert_one(generation.model_dump())
|
||||||
return str(res.inserted_id)
|
return str(res.inserted_id)
|
||||||
|
|
||||||
async def get_generation(self, generation_id: str) -> Optional[Generation]:
|
async def get_generation(self, generation_id: str) -> Generation | None:
|
||||||
res = await self.collection.find_one({"_id": ObjectId(generation_id)})
|
res = await self.collection.find_one({"_id": ObjectId(generation_id)})
|
||||||
if res is None:
|
if res is None:
|
||||||
return None
|
return None
|
||||||
@@ -25,20 +26,29 @@ class GenerationRepo:
|
|||||||
return Generation(**res)
|
return Generation(**res)
|
||||||
|
|
||||||
async def get_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None,
|
async def get_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None,
|
||||||
limit: int = 10, offset: int = 10, created_by: Optional[str] = None, project_id: Optional[str] = None) -> List[Generation]:
|
limit: int = 10, offset: int = 0, created_by: Optional[str] = None, project_id: Optional[str] = None, idea_id: Optional[str] = None) -> List[Generation]:
|
||||||
|
|
||||||
filter = {"is_deleted": False}
|
filter: dict[str, Any] = {"is_deleted": False}
|
||||||
if character_id is not None:
|
if character_id is not None:
|
||||||
filter["linked_character_id"] = character_id
|
filter["linked_character_id"] = character_id
|
||||||
if status is not None:
|
if status is not None:
|
||||||
filter["status"] = status
|
filter["status"] = status
|
||||||
if created_by is not None:
|
if created_by is not None:
|
||||||
filter["created_by"] = created_by
|
filter["created_by"] = created_by
|
||||||
|
# If filtering by created_by user (e.g. "My Generations"), we typically imply personal scope if project_id is None.
|
||||||
|
# But if project_id is passed, we filter by that.
|
||||||
|
if project_id is None:
|
||||||
filter["project_id"] = None
|
filter["project_id"] = None
|
||||||
if project_id is not None:
|
if project_id is not None:
|
||||||
filter["project_id"] = project_id
|
filter["project_id"] = project_id
|
||||||
|
if idea_id is not None:
|
||||||
|
filter["idea_id"] = idea_id
|
||||||
|
|
||||||
res = await self.collection.find(filter).sort("created_at", -1).skip(
|
# If fetching for an idea, sort by created_at ascending (cronological)
|
||||||
|
# Otherwise typically descending (newest first)
|
||||||
|
sort_order = 1 if idea_id else -1
|
||||||
|
|
||||||
|
res = await self.collection.find(filter).sort("created_at", sort_order).skip(
|
||||||
offset).limit(limit).to_list(None)
|
offset).limit(limit).to_list(None)
|
||||||
generations: List[Generation] = []
|
generations: List[Generation] = []
|
||||||
for generation in res:
|
for generation in res:
|
||||||
@@ -47,7 +57,7 @@ class GenerationRepo:
|
|||||||
return generations
|
return generations
|
||||||
|
|
||||||
async def count_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None,
|
async def count_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None,
|
||||||
album_id: Optional[str] = None, created_by: Optional[str] = None, project_id: Optional[str] = None) -> int:
|
album_id: Optional[str] = None, created_by: Optional[str] = None, project_id: Optional[str] = None, idea_id: Optional[str] = None) -> int:
|
||||||
args = {}
|
args = {}
|
||||||
if character_id is not None:
|
if character_id is not None:
|
||||||
args["linked_character_id"] = character_id
|
args["linked_character_id"] = character_id
|
||||||
@@ -55,8 +65,14 @@ class GenerationRepo:
|
|||||||
args["status"] = status
|
args["status"] = status
|
||||||
if created_by is not None:
|
if created_by is not None:
|
||||||
args["created_by"] = created_by
|
args["created_by"] = created_by
|
||||||
|
if project_id is None:
|
||||||
|
args["project_id"] = None
|
||||||
if project_id is not None:
|
if project_id is not None:
|
||||||
args["project_id"] = project_id
|
args["project_id"] = project_id
|
||||||
|
if idea_id is not None:
|
||||||
|
args["idea_id"] = idea_id
|
||||||
|
if album_id is not None:
|
||||||
|
args["album_id"] = album_id
|
||||||
return await self.collection.count_documents(args)
|
return await self.collection.count_documents(args)
|
||||||
|
|
||||||
async def get_generations_by_ids(self, generation_ids: List[str]) -> List[Generation]:
|
async def get_generations_by_ids(self, generation_ids: List[str]) -> List[Generation]:
|
||||||
@@ -77,3 +93,179 @@ class GenerationRepo:
|
|||||||
|
|
||||||
async def update_generation(self, generation: Generation, ):
|
async def update_generation(self, generation: Generation, ):
|
||||||
res = await self.collection.update_one({"_id": ObjectId(generation.id)}, {"$set": generation.model_dump()})
|
res = await self.collection.update_one({"_id": ObjectId(generation.id)}, {"$set": generation.model_dump()})
|
||||||
|
|
||||||
|
async def get_usage_stats(self, created_by: Optional[str] = None, project_id: Optional[str] = None) -> dict:
|
||||||
|
"""
|
||||||
|
Calculates usage statistics (runs, tokens, cost) using MongoDB aggregation.
|
||||||
|
Includes even soft-deleted generations to reflect actual expenditure.
|
||||||
|
"""
|
||||||
|
pipeline = []
|
||||||
|
|
||||||
|
# 1. Match all done generations (including soft-deleted)
|
||||||
|
match_stage = {"status": GenerationStatus.DONE}
|
||||||
|
if created_by:
|
||||||
|
match_stage["created_by"] = created_by
|
||||||
|
if project_id:
|
||||||
|
match_stage["project_id"] = project_id
|
||||||
|
|
||||||
|
pipeline.append({"$match": match_stage})
|
||||||
|
|
||||||
|
# 2. Group by null (total)
|
||||||
|
pipeline.append({
|
||||||
|
"$group": {
|
||||||
|
"_id": None,
|
||||||
|
"total_runs": {"$sum": 1},
|
||||||
|
"total_tokens": {
|
||||||
|
"$sum": {
|
||||||
|
"$cond": [
|
||||||
|
{"$and": [{"$gt": ["$input_token_usage", 0]}, {"$gt": ["$output_token_usage", 0]}]},
|
||||||
|
{"$add": ["$input_token_usage", "$output_token_usage"]},
|
||||||
|
{"$ifNull": ["$token_usage", 0]}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"total_input_tokens": {"$sum": {"$ifNull": ["$input_token_usage", 0]}},
|
||||||
|
"total_output_tokens": {"$sum": {"$ifNull": ["$output_token_usage", 0]}},
|
||||||
|
"total_cost": {
|
||||||
|
"$sum": {
|
||||||
|
"$add": [
|
||||||
|
{"$multiply": [{"$ifNull": ["$input_token_usage", 0]}, 0.000002]},
|
||||||
|
{"$multiply": [{"$ifNull": ["$output_token_usage", 0]}, 0.00012]}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
cursor = self.collection.aggregate(pipeline)
|
||||||
|
res = await cursor.to_list(1)
|
||||||
|
|
||||||
|
if not res:
|
||||||
|
return {
|
||||||
|
"total_runs": 0,
|
||||||
|
"total_tokens": 0,
|
||||||
|
"total_input_tokens": 0,
|
||||||
|
"total_output_tokens": 0,
|
||||||
|
"total_cost": 0.0
|
||||||
|
}
|
||||||
|
|
||||||
|
result = res[0]
|
||||||
|
result.pop("_id")
|
||||||
|
result["total_cost"] = round(result["total_cost"], 4)
|
||||||
|
return result
|
||||||
|
|
||||||
|
async def get_usage_breakdown(self, group_by: str = "created_by", project_id: Optional[str] = None, created_by: Optional[str] = None) -> List[dict]:
|
||||||
|
"""
|
||||||
|
Returns usage statistics grouped by user or project.
|
||||||
|
Includes even soft-deleted generations to reflect actual expenditure.
|
||||||
|
"""
|
||||||
|
pipeline = []
|
||||||
|
|
||||||
|
match_stage = {"status": GenerationStatus.DONE}
|
||||||
|
if project_id:
|
||||||
|
match_stage["project_id"] = project_id
|
||||||
|
if created_by:
|
||||||
|
match_stage["created_by"] = created_by
|
||||||
|
|
||||||
|
pipeline.append({"$match": match_stage})
|
||||||
|
|
||||||
|
pipeline.append({
|
||||||
|
"$group": {
|
||||||
|
"_id": f"${group_by}",
|
||||||
|
"total_runs": {"$sum": 1},
|
||||||
|
"total_tokens": {
|
||||||
|
"$sum": {
|
||||||
|
"$cond": [
|
||||||
|
{"$and": [{"$gt": ["$input_token_usage", 0]}, {"$gt": ["$output_token_usage", 0]}]},
|
||||||
|
{"$add": ["$input_token_usage", "$output_token_usage"]},
|
||||||
|
{"$ifNull": ["$token_usage", 0]}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"total_input_tokens": {"$sum": {"$ifNull": ["$input_token_usage", 0]}},
|
||||||
|
"total_output_tokens": {"$sum": {"$ifNull": ["$output_token_usage", 0]}},
|
||||||
|
"total_cost": {
|
||||||
|
"$sum": {
|
||||||
|
"$add": [
|
||||||
|
{"$multiply": [{"$ifNull": ["$input_token_usage", 0]}, 0.000002]},
|
||||||
|
{"$multiply": [{"$ifNull": ["$output_token_usage", 0]}, 0.00012]}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
pipeline.append({"$sort": {"total_cost": -1}})
|
||||||
|
|
||||||
|
cursor = self.collection.aggregate(pipeline)
|
||||||
|
res = await cursor.to_list(None)
|
||||||
|
|
||||||
|
results = []
|
||||||
|
for item in res:
|
||||||
|
entity_id = item.pop("_id")
|
||||||
|
item["total_cost"] = round(item["total_cost"], 4)
|
||||||
|
results.append({
|
||||||
|
"entity_id": str(entity_id) if entity_id else "unknown",
|
||||||
|
"stats": item
|
||||||
|
})
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
async def get_generations_by_group(self, group_id: str) -> List[Generation]:
|
||||||
|
res = await self.collection.find({"generation_group_id": group_id, "is_deleted": False}).sort("created_at", 1).to_list(None)
|
||||||
|
generations: List[Generation] = []
|
||||||
|
for generation in res:
|
||||||
|
generation["id"] = str(generation.pop("_id"))
|
||||||
|
generations.append(Generation(**generation))
|
||||||
|
return generations
|
||||||
|
|
||||||
|
async def cancel_stale_generations(self, timeout_minutes: int = 5) -> int:
|
||||||
|
cutoff_time = datetime.now(UTC) - timedelta(minutes=timeout_minutes)
|
||||||
|
res = await self.collection.update_many(
|
||||||
|
{
|
||||||
|
"status": GenerationStatus.RUNNING,
|
||||||
|
"created_at": {"$lt": cutoff_time}
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$set": {
|
||||||
|
"status": GenerationStatus.FAILED,
|
||||||
|
"failed_reason": "Timeout: Execution time limit exceeded",
|
||||||
|
"updated_at": datetime.now(UTC)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return res.modified_count
|
||||||
|
|
||||||
|
async def soft_delete_old_generations(self, days: int = 2) -> tuple[int, List[str]]:
|
||||||
|
"""
|
||||||
|
Мягко удаляет генерации старше N дней.
|
||||||
|
Возвращает (количество удалённых, список asset IDs для очистки).
|
||||||
|
"""
|
||||||
|
cutoff_time = datetime.now(UTC) - timedelta(days=days)
|
||||||
|
filter_query = {
|
||||||
|
"is_deleted": False,
|
||||||
|
"status": {"$in": [GenerationStatus.DONE, GenerationStatus.FAILED]},
|
||||||
|
"created_at": {"$lt": cutoff_time}
|
||||||
|
}
|
||||||
|
|
||||||
|
# Сначала собираем asset IDs из удаляемых генераций
|
||||||
|
asset_ids: List[str] = []
|
||||||
|
cursor = self.collection.find(filter_query, {"result_list": 1, "assets_list": 1})
|
||||||
|
async for doc in cursor:
|
||||||
|
asset_ids.extend(doc.get("result_list", []))
|
||||||
|
asset_ids.extend(doc.get("assets_list", []))
|
||||||
|
|
||||||
|
# Мягкое удаление
|
||||||
|
res = await self.collection.update_many(
|
||||||
|
filter_query,
|
||||||
|
{
|
||||||
|
"$set": {
|
||||||
|
"is_deleted": True,
|
||||||
|
"updated_at": datetime.now(UTC)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Убираем дубликаты
|
||||||
|
unique_asset_ids = list(set(asset_ids))
|
||||||
|
return res.modified_count, unique_asset_ids
|
||||||
|
|||||||
91
repos/idea_repo.py
Normal file
91
repos/idea_repo.py
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
from typing import Optional, List
|
||||||
|
from bson import ObjectId
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
from models.Idea import Idea
|
||||||
|
|
||||||
|
class IdeaRepo:
|
||||||
|
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
||||||
|
self.collection = client[db_name]["ideas"]
|
||||||
|
|
||||||
|
async def create_idea(self, idea: Idea) -> str:
|
||||||
|
res = await self.collection.insert_one(idea.model_dump())
|
||||||
|
return str(res.inserted_id)
|
||||||
|
|
||||||
|
async def get_idea(self, idea_id: str) -> Optional[Idea]:
|
||||||
|
if not ObjectId.is_valid(idea_id):
|
||||||
|
return None
|
||||||
|
res = await self.collection.find_one({"_id": ObjectId(idea_id)})
|
||||||
|
if res:
|
||||||
|
res["id"] = str(res.pop("_id"))
|
||||||
|
return Idea(**res)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def get_ideas(self, project_id: Optional[str], user_id: str, limit: int = 20, offset: int = 0) -> List[dict]:
|
||||||
|
if project_id:
|
||||||
|
match_stage = {"project_id": project_id, "is_deleted": False}
|
||||||
|
else:
|
||||||
|
match_stage = {"created_by": user_id, "project_id": None, "is_deleted": False}
|
||||||
|
|
||||||
|
pipeline = [
|
||||||
|
{"$match": match_stage},
|
||||||
|
{"$sort": {"updated_at": -1}},
|
||||||
|
{"$skip": offset},
|
||||||
|
{"$limit": limit},
|
||||||
|
# Add string id field for lookup
|
||||||
|
{"$addFields": {"str_id": {"$toString": "$_id"}}},
|
||||||
|
# Lookup generations
|
||||||
|
{
|
||||||
|
"$lookup": {
|
||||||
|
"from": "generations",
|
||||||
|
"let": {"idea_id": "$str_id"},
|
||||||
|
"pipeline": [
|
||||||
|
{
|
||||||
|
"$match": {
|
||||||
|
"$and": [
|
||||||
|
{"$expr": {"$eq": ["$idea_id", "$$idea_id"]}},
|
||||||
|
{"status": "done"},
|
||||||
|
{"result_list": {"$exists": True, "$not": {"$size": 0}}},
|
||||||
|
{"is_deleted": False}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{"$sort": {"created_at": -1}}, # Ensure we get the latest successful
|
||||||
|
{"$limit": 1}
|
||||||
|
],
|
||||||
|
"as": "generations"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
# Unwind generations array (preserve ideas without generations)
|
||||||
|
{"$unwind": {"path": "$generations", "preserveNullAndEmptyArrays": True}},
|
||||||
|
# Rename for clarity
|
||||||
|
{"$addFields": {
|
||||||
|
"last_generation": "$generations",
|
||||||
|
"id": "$str_id"
|
||||||
|
}},
|
||||||
|
{"$project": {"generations": 0, "str_id": 0, "_id": 0}}
|
||||||
|
]
|
||||||
|
|
||||||
|
return await self.collection.aggregate(pipeline).to_list(None)
|
||||||
|
|
||||||
|
async def delete_idea(self, idea_id: str) -> bool:
|
||||||
|
if not ObjectId.is_valid(idea_id):
|
||||||
|
return False
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(idea_id)},
|
||||||
|
{"$set": {"is_deleted": True}}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def update_idea(self, idea: Idea) -> bool:
|
||||||
|
if not idea.id or not ObjectId.is_valid(idea.id):
|
||||||
|
return False
|
||||||
|
|
||||||
|
idea_dict = idea.model_dump()
|
||||||
|
if "id" in idea_dict:
|
||||||
|
del idea_dict["id"]
|
||||||
|
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(idea.id)},
|
||||||
|
{"$set": idea_dict}
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
97
repos/post_repo.py
Normal file
97
repos/post_repo.py
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from datetime import datetime
|
||||||
|
import logging
|
||||||
|
from bson import ObjectId
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
|
||||||
|
from models.Post import Post
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PostRepo:
|
||||||
|
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
|
||||||
|
self.collection = client[db_name]["posts"]
|
||||||
|
|
||||||
|
async def create_post(self, post: Post) -> str:
|
||||||
|
res = await self.collection.insert_one(post.model_dump())
|
||||||
|
return str(res.inserted_id)
|
||||||
|
|
||||||
|
async def get_post(self, post_id: str) -> Optional[Post]:
|
||||||
|
if not ObjectId.is_valid(post_id):
|
||||||
|
return None
|
||||||
|
res = await self.collection.find_one({"_id": ObjectId(post_id), "is_deleted": False})
|
||||||
|
if res:
|
||||||
|
res["id"] = str(res.pop("_id"))
|
||||||
|
return Post(**res)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def get_posts(
|
||||||
|
self,
|
||||||
|
project_id: Optional[str],
|
||||||
|
user_id: str,
|
||||||
|
limit: int = 20,
|
||||||
|
offset: int = 0,
|
||||||
|
date_from: Optional[datetime] = None,
|
||||||
|
date_to: Optional[datetime] = None,
|
||||||
|
) -> List[Post]:
|
||||||
|
if project_id:
|
||||||
|
match = {"project_id": project_id, "is_deleted": False}
|
||||||
|
else:
|
||||||
|
match = {"created_by": user_id, "project_id": None, "is_deleted": False}
|
||||||
|
|
||||||
|
if date_from or date_to:
|
||||||
|
date_filter = {}
|
||||||
|
if date_from:
|
||||||
|
date_filter["$gte"] = date_from
|
||||||
|
if date_to:
|
||||||
|
date_filter["$lte"] = date_to
|
||||||
|
match["date"] = date_filter
|
||||||
|
|
||||||
|
cursor = (
|
||||||
|
self.collection.find(match)
|
||||||
|
.sort("date", -1)
|
||||||
|
.skip(offset)
|
||||||
|
.limit(limit)
|
||||||
|
)
|
||||||
|
posts = []
|
||||||
|
async for doc in cursor:
|
||||||
|
doc["id"] = str(doc.pop("_id"))
|
||||||
|
posts.append(Post(**doc))
|
||||||
|
return posts
|
||||||
|
|
||||||
|
async def update_post(self, post_id: str, data: dict) -> bool:
|
||||||
|
if not ObjectId.is_valid(post_id):
|
||||||
|
return False
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(post_id)},
|
||||||
|
{"$set": data},
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def delete_post(self, post_id: str) -> bool:
|
||||||
|
if not ObjectId.is_valid(post_id):
|
||||||
|
return False
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(post_id)},
|
||||||
|
{"$set": {"is_deleted": True}},
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def add_generations(self, post_id: str, generation_ids: List[str]) -> bool:
|
||||||
|
if not ObjectId.is_valid(post_id):
|
||||||
|
return False
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(post_id)},
|
||||||
|
{"$addToSet": {"generation_ids": {"$each": generation_ids}}},
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
|
|
||||||
|
async def remove_generation(self, post_id: str, generation_id: str) -> bool:
|
||||||
|
if not ObjectId.is_valid(post_id):
|
||||||
|
return False
|
||||||
|
res = await self.collection.update_one(
|
||||||
|
{"_id": ObjectId(post_id)},
|
||||||
|
{"$pull": {"generation_ids": generation_id}},
|
||||||
|
)
|
||||||
|
return res.modified_count > 0
|
||||||
@@ -51,3 +51,4 @@ python-jose[cryptography]==3.3.0
|
|||||||
python-multipart==0.0.22
|
python-multipart==0.0.22
|
||||||
email-validator
|
email-validator
|
||||||
prometheus-fastapi-instrumentator
|
prometheus-fastapi-instrumentator
|
||||||
|
pydantic-settings==2.13.0
|
||||||
|
|||||||
Binary file not shown.
Binary file not shown.
@@ -51,57 +51,66 @@ async def new_char_bio(message: Message, state: FSMContext, dao: DAO, bot: Bot):
|
|||||||
wait_msg = await message.answer("💾 Сохраняю персонажа...")
|
wait_msg = await message.answer("💾 Сохраняю персонажа...")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# ВОТ ТУТ скачиваем файл (прямо перед сохранением)
|
# 1. Скачиваем файл (один раз)
|
||||||
|
# TODO: Для больших файлов лучше использовать streaming или сохранять во временный файл
|
||||||
file_io = await bot.download(file_id)
|
file_io = await bot.download(file_id)
|
||||||
# photo_bytes = file_io.getvalue() # Получаем байты
|
file_bytes = file_io.read()
|
||||||
|
|
||||||
|
# 2. Создаем Character (сначала без ассета, чтобы получить ID)
|
||||||
# Создаем модель
|
|
||||||
char = Character(
|
char = Character(
|
||||||
id=None,
|
id=None,
|
||||||
name=name,
|
name=name,
|
||||||
character_image_data=file_io.read(),
|
|
||||||
character_image_tg_id=None,
|
character_image_tg_id=None,
|
||||||
character_image_doc_tg_id=file_id,
|
character_image_doc_tg_id=file_id,
|
||||||
character_bio=bio,
|
character_bio=bio,
|
||||||
created_by=str(message.from_user.id)
|
created_by=str(message.from_user.id)
|
||||||
)
|
)
|
||||||
file_io.close()
|
|
||||||
|
|
||||||
# Сохраняем через DAO
|
|
||||||
|
|
||||||
|
# Сохраняем, чтобы получить ID
|
||||||
await dao.chars.add_character(char)
|
await dao.chars.add_character(char)
|
||||||
file_info = await bot.get_file(char.character_image_doc_tg_id)
|
|
||||||
file_bytes = await bot.download_file(file_info.file_path)
|
# 3. Создаем Asset (связанный с персонажем)
|
||||||
file_io = file_bytes.read()
|
avatar_asset_id = await dao.assets.create_asset(
|
||||||
avatar_asset = await dao.assets.create_asset(
|
Asset(
|
||||||
Asset(name="avatar.png", type=AssetType.UPLOADED, content_type=AssetContentType.IMAGE, linked_char_id=str(char.id), data=file_io,
|
name="avatar.png",
|
||||||
tg_doc_file_id=file_id))
|
type=AssetType.UPLOADED,
|
||||||
char.avatar_image = avatar_asset.link
|
content_type=AssetContentType.IMAGE,
|
||||||
|
linked_char_id=str(char.id),
|
||||||
|
data=file_bytes,
|
||||||
|
tg_doc_file_id=file_id
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# 4. Обновляем персонажа ссылками на ассет
|
||||||
|
char.avatar_asset_id = avatar_asset_id
|
||||||
|
char.avatar_image = f"/api/assets/{avatar_asset_id}" # Формируем ссылку вручную или используем метод, если появится
|
||||||
|
|
||||||
# Отправляем подтверждение
|
# Отправляем подтверждение
|
||||||
# Используем байты для отправки обратно
|
|
||||||
photo_msg = await message.answer_photo(
|
photo_msg = await message.answer_photo(
|
||||||
photo=BufferedInputFile(file_io,
|
photo=BufferedInputFile(file_bytes, filename="char.jpg"),
|
||||||
filename="char.jpg") if not char.character_image_tg_id else char.character_image_tg_id,
|
|
||||||
caption=(
|
caption=(
|
||||||
"🎉 <b>Персонаж создан!</b>\n\n"
|
"🎉 <b>Персонаж создан!</b>\n\n"
|
||||||
f"👤 <b>Имя:</b> {char.name}\n"
|
f"👤 <b>Имя:</b> {char.name}\n"
|
||||||
f"📝 <b>Био:</b> {char.character_bio}"
|
f"📝 <b>Био:</b> {char.character_bio}"
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
file_bytes.close()
|
|
||||||
char.character_image_tg_id = photo_msg.photo[0].file_id
|
|
||||||
|
|
||||||
|
# Сохраняем TG ID фото (которое отправили как фото, а не документ)
|
||||||
|
char.character_image_tg_id = photo_msg.photo[-1].file_id
|
||||||
|
|
||||||
|
# Финальное обновление персонажа
|
||||||
await dao.chars.update_char(char.id, char)
|
await dao.chars.update_char(char.id, char)
|
||||||
|
|
||||||
await wait_msg.delete()
|
await wait_msg.delete()
|
||||||
|
file_io.close()
|
||||||
|
|
||||||
# Сбрасываем состояние
|
# Сбрасываем состояние
|
||||||
await state.clear()
|
await state.clear()
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logging.error(e)
|
logger.error(f"Error creating character: {e}")
|
||||||
|
traceback.print_exc()
|
||||||
await wait_msg.edit_text(f"❌ Ошибка при сохранении: {e}")
|
await wait_msg.edit_text(f"❌ Ошибка при сохранении: {e}")
|
||||||
# Не сбрасываем стейт, даем возможность попробовать ввести био снова или начать заново
|
|
||||||
|
|
||||||
|
|
||||||
@router.message(Command("chars"))
|
@router.message(Command("chars"))
|
||||||
|
|||||||
@@ -3,17 +3,17 @@ import pytest
|
|||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
from motor.motor_asyncio import AsyncIOMotorClient
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
import os
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
from config import settings
|
||||||
|
|
||||||
from main import app
|
from aiws import app
|
||||||
from api.endpoints.auth import get_current_user
|
from api.endpoints.auth import get_current_user
|
||||||
from api.dependency import get_dao
|
from api.dependency import get_dao
|
||||||
from repos.dao import DAO
|
from repos.dao import DAO
|
||||||
from models.Character import Character
|
from models.Character import Character
|
||||||
|
|
||||||
# Config for test DB
|
# Config for test DB
|
||||||
MONGO_HOST = os.getenv("MONGO_HOST", "mongodb://admin:super_secure_password@31.59.58.220:27017")
|
MONGO_HOST = settings.MONGO_HOST
|
||||||
DB_NAME = "bot_db_test_chars"
|
DB_NAME = "bot_db_test_chars"
|
||||||
|
|
||||||
# Mock User
|
# Mock User
|
||||||
|
|||||||
@@ -10,13 +10,13 @@ import json
|
|||||||
import requests
|
import requests
|
||||||
import base64
|
import base64
|
||||||
import os
|
import os
|
||||||
from dotenv import load_dotenv
|
from config import settings
|
||||||
|
|
||||||
load_dotenv()
|
# Load env is not needed as settings handles it
|
||||||
|
|
||||||
# Configuration
|
# Configuration
|
||||||
API_URL = "http://localhost:8090/api/generations/import"
|
API_URL = "http://localhost:8090/api/generations/import"
|
||||||
SECRET = os.getenv("EXTERNAL_API_SECRET", "your_super_secret_key_change_this_in_production")
|
SECRET = settings.EXTERNAL_API_SECRET or "your_super_secret_key_change_this_in_production"
|
||||||
|
|
||||||
# Sample generation data
|
# Sample generation data
|
||||||
generation_data = {
|
generation_data = {
|
||||||
|
|||||||
96
tests/test_idea.py
Normal file
96
tests/test_idea.py
Normal file
@@ -0,0 +1,96 @@
|
|||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
from bson import ObjectId
|
||||||
|
|
||||||
|
# Import from project root (requires PYTHONPATH=.)
|
||||||
|
from api.service.idea_service import IdeaService
|
||||||
|
from repos.dao import DAO
|
||||||
|
from models.Idea import Idea
|
||||||
|
from models.Generation import Generation, GenerationStatus
|
||||||
|
from models.enums import AspectRatios, Quality
|
||||||
|
from config import settings
|
||||||
|
|
||||||
|
MONGO_HOST = settings.MONGO_HOST
|
||||||
|
DB_NAME = settings.DB_NAME
|
||||||
|
|
||||||
|
print(f"Connecting to MongoDB: {MONGO_HOST}, DB: {DB_NAME}")
|
||||||
|
|
||||||
|
async def test_idea_flow():
|
||||||
|
client = AsyncIOMotorClient(MONGO_HOST)
|
||||||
|
dao = DAO(client, db_name=DB_NAME)
|
||||||
|
service = IdeaService(dao)
|
||||||
|
|
||||||
|
# 1. Create an Idea
|
||||||
|
print("Creating idea...")
|
||||||
|
user_id = "test_user_123"
|
||||||
|
project_id = "test_project_abc"
|
||||||
|
idea = await service.create_idea("My Test Idea", "Initial Description", project_id, user_id)
|
||||||
|
print(f"Idea created: {idea.id} - {idea.name}")
|
||||||
|
|
||||||
|
# 2. Update Idea
|
||||||
|
print("Updating idea...")
|
||||||
|
updated_idea = await service.update_idea(idea.id, description="Updated description")
|
||||||
|
print(f"Idea updated: {updated_idea.description}")
|
||||||
|
if updated_idea.description == "Updated description":
|
||||||
|
print("✅ Idea update successful")
|
||||||
|
else:
|
||||||
|
print("❌ Idea update FAILED")
|
||||||
|
|
||||||
|
# 3. Add Generation linked to Idea
|
||||||
|
print("Creating generation linked to idea...")
|
||||||
|
gen = Generation(
|
||||||
|
prompt="idea generation 1",
|
||||||
|
# idea_id=idea.id, <-- Intentionally NOT linking initially to test linking method
|
||||||
|
project_id=project_id,
|
||||||
|
created_by=user_id,
|
||||||
|
aspect_ratio=AspectRatios.NINESIXTEEN,
|
||||||
|
quality=Quality.ONEK,
|
||||||
|
assets_list=[]
|
||||||
|
)
|
||||||
|
gen_id = await dao.generations.create_generation(gen)
|
||||||
|
print(f"Created generation: {gen_id}")
|
||||||
|
|
||||||
|
# Link generation to idea
|
||||||
|
print("Linking generation to idea...")
|
||||||
|
success = await service.add_generation_to_idea(idea.id, gen_id)
|
||||||
|
if success:
|
||||||
|
print("✅ Linking successful")
|
||||||
|
else:
|
||||||
|
print("❌ Linking FAILED")
|
||||||
|
|
||||||
|
# Debug: Check if generation was saved with idea_id
|
||||||
|
saved_gen = await dao.generations.collection.find_one({"_id": ObjectId(gen_id)})
|
||||||
|
print(f"DEBUG: Saved Generation in DB idea_id: {saved_gen.get('idea_id')}")
|
||||||
|
|
||||||
|
# 4. Fetch Generations for Idea (Verify filtering and ordering)
|
||||||
|
print("Fetching generations for idea...")
|
||||||
|
gens = await service.dao.generations.get_generations(idea_id=idea.id) # using repo directly as service might return wrapper
|
||||||
|
print(f"Found {len(gens)} generations in idea")
|
||||||
|
|
||||||
|
if len(gens) == 1 and gens[0].id == gen_id:
|
||||||
|
print("✅ Generation retrieval successful")
|
||||||
|
else:
|
||||||
|
print("❌ Generation retrieval FAILED")
|
||||||
|
|
||||||
|
# 5. Fetch Ideas for Project
|
||||||
|
ideas = await service.get_ideas(project_id)
|
||||||
|
print(f"Found {len(ideas)} ideas for project")
|
||||||
|
|
||||||
|
# Cleaning up
|
||||||
|
print("Cleaning up...")
|
||||||
|
await service.delete_idea(idea.id)
|
||||||
|
await dao.generations.collection.delete_one({"_id": ObjectId(gen_id)})
|
||||||
|
|
||||||
|
# Verify deletion
|
||||||
|
deleted_idea = await service.get_idea(idea.id)
|
||||||
|
# IdeaRepo.delete_idea logic sets is_deleted=True
|
||||||
|
if deleted_idea and deleted_idea.is_deleted:
|
||||||
|
print(f"✅ Idea deleted successfully")
|
||||||
|
|
||||||
|
# Hard delete for cleanup
|
||||||
|
await dao.ideas.collection.delete_one({"_id": ObjectId(idea.id)})
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(test_idea_flow())
|
||||||
@@ -1,15 +1,14 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
from dotenv import load_dotenv
|
from config import settings
|
||||||
from adapters.s3_adapter import S3Adapter
|
from adapters.s3_adapter import S3Adapter
|
||||||
|
|
||||||
async def test_s3():
|
async def test_s3():
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
endpoint = os.getenv("MINIO_ENDPOINT", "http://localhost:9000")
|
endpoint = settings.MINIO_ENDPOINT
|
||||||
access_key = os.getenv("MINIO_ACCESS_KEY")
|
access_key = settings.MINIO_ACCESS_KEY
|
||||||
secret_key = os.getenv("MINIO_SECRET_KEY")
|
secret_key = settings.MINIO_SECRET_KEY
|
||||||
bucket = os.getenv("MINIO_BUCKET")
|
bucket = settings.MINIO_BUCKET
|
||||||
|
|
||||||
print(f"Connecting to {endpoint}, bucket: {bucket}")
|
print(f"Connecting to {endpoint}, bucket: {bucket}")
|
||||||
|
|
||||||
|
|||||||
50
tests/test_scheduler.py
Normal file
50
tests/test_scheduler.py
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
from datetime import datetime, timedelta, UTC
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
from models.Generation import Generation, GenerationStatus
|
||||||
|
from repos.generation_repo import GenerationRepo
|
||||||
|
from config import settings
|
||||||
|
|
||||||
|
# Mock configs if not present in env
|
||||||
|
MONGO_HOST = settings.MONGO_HOST
|
||||||
|
DB_NAME = settings.DB_NAME
|
||||||
|
|
||||||
|
print(f"Connecting to MongoDB: {MONGO_HOST}, DB: {DB_NAME}")
|
||||||
|
|
||||||
|
async def test_scheduler():
|
||||||
|
client = AsyncIOMotorClient(MONGO_HOST)
|
||||||
|
repo = GenerationRepo(client, db_name=DB_NAME)
|
||||||
|
|
||||||
|
# 1. Create a "stale" generation (2 hours ago)
|
||||||
|
stale_gen = Generation(
|
||||||
|
prompt="stale test",
|
||||||
|
status=GenerationStatus.RUNNING,
|
||||||
|
created_at=datetime.now(UTC) - timedelta(minutes=120),
|
||||||
|
assets_list=[],
|
||||||
|
aspect_ratio="NINESIXTEEN",
|
||||||
|
quality="ONEK"
|
||||||
|
)
|
||||||
|
gen_id = await repo.create_generation(stale_gen)
|
||||||
|
print(f"Created stale generation: {gen_id}")
|
||||||
|
|
||||||
|
# 2. Run cleanup
|
||||||
|
print("Running cleanup...")
|
||||||
|
count = await repo.cancel_stale_generations(timeout_minutes=60)
|
||||||
|
print(f"Cleaned up {count} generations")
|
||||||
|
|
||||||
|
# 3. Verify status
|
||||||
|
updated_gen = await repo.get_generation(gen_id)
|
||||||
|
print(f"Generation status: {updated_gen.status}")
|
||||||
|
print(f"Failed reason: {updated_gen.failed_reason}")
|
||||||
|
|
||||||
|
if updated_gen.status == GenerationStatus.FAILED:
|
||||||
|
print("✅ SUCCESS: Generation marked as FAILED")
|
||||||
|
else:
|
||||||
|
print("❌ FAILURE: Generation status not updated")
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
await repo.collection.delete_one({"_id": updated_gen.id}) # Remove test data
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(test_scheduler())
|
||||||
@@ -10,10 +10,11 @@ from repos.dao import DAO
|
|||||||
from models.Album import Album
|
from models.Album import Album
|
||||||
from models.Generation import Generation, GenerationStatus
|
from models.Generation import Generation, GenerationStatus
|
||||||
from models.enums import AspectRatios, Quality
|
from models.enums import AspectRatios, Quality
|
||||||
|
from config import settings
|
||||||
|
|
||||||
# Mock config
|
# Mock config
|
||||||
# Use the same host as aiws.py but different DB
|
# Use the same host as aiws.py but different DB
|
||||||
MONGO_HOST = os.getenv("MONGO_HOST", "mongodb://admin:super_secure_password@31.59.58.220:27017")
|
MONGO_HOST = settings.MONGO_HOST
|
||||||
DB_NAME = "bot_db_test_albums"
|
DB_NAME = "bot_db_test_albums"
|
||||||
|
|
||||||
async def test_albums():
|
async def test_albums():
|
||||||
@@ -83,8 +84,6 @@ async def test_albums():
|
|||||||
client.close()
|
client.close()
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
from dotenv import load_dotenv
|
|
||||||
load_dotenv()
|
|
||||||
try:
|
try:
|
||||||
asyncio.run(test_albums())
|
asyncio.run(test_albums())
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -1,29 +1,28 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from dotenv import load_dotenv
|
|
||||||
from motor.motor_asyncio import AsyncIOMotorClient
|
from motor.motor_asyncio import AsyncIOMotorClient
|
||||||
|
|
||||||
|
from config import settings
|
||||||
from models.Asset import Asset, AssetType
|
from models.Asset import Asset, AssetType
|
||||||
from repos.assets_repo import AssetsRepo
|
from repos.assets_repo import AssetsRepo
|
||||||
from adapters.s3_adapter import S3Adapter
|
from adapters.s3_adapter import S3Adapter
|
||||||
|
|
||||||
# Load env to get credentials
|
# Load env is not needed as settings handles it
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
async def test_integration():
|
async def test_integration():
|
||||||
print("🚀 Starting integration test...")
|
print("🚀 Starting integration test...")
|
||||||
|
|
||||||
# 1. Setup Dependencies
|
# 1. Setup Dependencies
|
||||||
mongo_uri = os.getenv("MONGO_HOST", "mongodb://localhost:27017")
|
mongo_uri = settings.MONGO_HOST
|
||||||
client = AsyncIOMotorClient(mongo_uri)
|
client = AsyncIOMotorClient(mongo_uri)
|
||||||
db_name = os.getenv("DB_NAME", "bot_db_test")
|
db_name = settings.DB_NAME + "_test"
|
||||||
|
|
||||||
s3_adapter = S3Adapter(
|
s3_adapter = S3Adapter(
|
||||||
endpoint_url=os.getenv("MINIO_ENDPOINT", "http://localhost:9000"),
|
endpoint_url=settings.MINIO_ENDPOINT,
|
||||||
aws_access_key_id=os.getenv("MINIO_ACCESS_KEY", "admin"),
|
aws_access_key_id=settings.MINIO_ACCESS_KEY,
|
||||||
aws_secret_access_key=os.getenv("MINIO_SECRET_KEY", "SuperSecretPassword123!"),
|
aws_secret_access_key=settings.MINIO_SECRET_KEY,
|
||||||
bucket_name=os.getenv("MINIO_BUCKET", "ai-char")
|
bucket_name=settings.MINIO_BUCKET
|
||||||
)
|
)
|
||||||
|
|
||||||
repo = AssetsRepo(client, s3_adapter, db_name=db_name)
|
repo = AssetsRepo(client, s3_adapter, db_name=db_name)
|
||||||
|
|||||||
Binary file not shown.
Binary file not shown.
@@ -3,12 +3,12 @@ from typing import Optional, Union, Any
|
|||||||
|
|
||||||
from jose import jwt
|
from jose import jwt
|
||||||
from passlib.context import CryptContext
|
from passlib.context import CryptContext
|
||||||
|
from config import settings
|
||||||
|
|
||||||
# Настройки безопасности (лучше вынести в config/env, но для старта здесь)
|
# Настройки безопасности берутся из config.py
|
||||||
# SECRET_KEY должен быть сложным и секретным в продакшене!
|
SECRET_KEY = settings.SECRET_KEY
|
||||||
SECRET_KEY = "CHANGE_ME_TO_A_SUPER_SECRET_KEY"
|
ALGORITHM = settings.ALGORITHM
|
||||||
ALGORITHM = "HS256"
|
ACCESS_TOKEN_EXPIRE_MINUTES = settings.ACCESS_TOKEN_EXPIRE_MINUTES
|
||||||
ACCESS_TOKEN_EXPIRE_MINUTES = 30 * 24 * 60 # 30 дней, например
|
|
||||||
|
|
||||||
pwd_context = CryptContext(schemes=["argon2"], deprecated="auto")
|
pwd_context = CryptContext(schemes=["argon2"], deprecated="auto")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user