diff --git a/adapters/google_adapter.py b/adapters/google_adapter.py index a4cd452..4e14f35 100644 --- a/adapters/google_adapter.py +++ b/adapters/google_adapter.py @@ -8,7 +8,7 @@ from google import genai from google.genai import types from adapters.Exception import GoogleGenerationException -from models.enums import AspectRatios, Quality +from models.enums import AspectRatios, Quality, TextModel, ImageModel logger = logging.getLogger(__name__) @@ -19,10 +19,6 @@ class GoogleAdapter: raise ValueError("API Key for Gemini is missing") self.client = genai.Client(api_key=api_key) - # Константы моделей - self.TEXT_MODEL = "gemini-3.1-pro-preview" - self.IMAGE_MODEL = "gemini-3-pro-image-preview" - def _prepare_contents(self, prompt: str, images_list: List[bytes] | None = None) -> tuple: """Вспомогательный метод для подготовки контента (текст + картинки). Returns (contents, opened_images) — caller MUST close opened_images after use.""" @@ -41,16 +37,19 @@ class GoogleAdapter: logger.info("Preparing content with no images") return contents, opened_images - def generate_text(self, prompt: str, images_list: List[bytes] | None = None) -> str: + def generate_text(self, prompt: str, model: str = "gemini-3.1-pro-preview", images_list: List[bytes] | None = None) -> str: """ Генерация текста (Чат или Vision). Возвращает строку с ответом. """ + if model not in [m.value for m in TextModel]: + raise ValueError(f"Invalid model for text generation: {model}. Expected one of: {[m.value for m in TextModel]}") + contents, opened_images = self._prepare_contents(prompt, images_list) - logger.info(f"Generating text: {prompt}") + logger.info(f"Generating text: {prompt} with model: {model}") try: response = self.client.models.generate_content( - model=self.TEXT_MODEL, + model=model, contents=contents, config=types.GenerateContentConfig( response_modalities=['TEXT'], @@ -74,21 +73,23 @@ class GoogleAdapter: for img in opened_images: img.close() - def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] | None = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]: + def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, model: str = "gemini-3-pro-image-preview", images_list: List[bytes] | None = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]: """ Генерация изображений (Text-to-Image или Image-to-Image). Возвращает список байтовых потоков (готовых к отправке). """ + if model not in [m.value for m in ImageModel]: + raise ValueError(f"Invalid model for image generation: {model}. Expected one of: {[m.value for m in ImageModel]}") 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}, Model: {model}") start_time = datetime.now() token_usage = 0 try: response = self.client.models.generate_content( - model=self.IMAGE_MODEL, + model=model, contents=contents, config=types.GenerateContentConfig( response_modalities=['IMAGE'], diff --git a/api/endpoints/admin.py b/api/endpoints/admin.py index 59f89d5..1fa8ab2 100644 --- a/api/endpoints/admin.py +++ b/api/endpoints/admin.py @@ -1,4 +1,4 @@ -from typing import Annotated, List +from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, status from fastapi.security import OAuth2PasswordBearer @@ -54,7 +54,7 @@ class UserResponse(BaseModel): class Config: from_attributes = True -@router.get("/approvals", response_model=List[UserResponse]) +@router.get("/approvals", response_model=list[UserResponse]) async def list_pending_users( admin: Annotated[dict, Depends(get_current_admin)], repo: Annotated[UsersRepo, Depends(get_users_repo)] diff --git a/api/endpoints/album_router.py b/api/endpoints/album_router.py index 252a619..09b4a4c 100644 --- a/api/endpoints/album_router.py +++ b/api/endpoints/album_router.py @@ -1,4 +1,3 @@ -from typing import List, Optional from fastapi import APIRouter, HTTPException, status, Request from pydantic import BaseModel @@ -13,18 +12,18 @@ router = APIRouter(prefix="/api/albums", tags=["Albums"]) class AlbumCreateRequest(BaseModel): name: str - description: Optional[str] = None + description: str | None = None class AlbumUpdateRequest(BaseModel): - name: Optional[str] = None - description: Optional[str] = None + name: str | None = None + description: str | None = None class AlbumResponse(BaseModel): id: str name: str - description: Optional[str] = None - generation_ids: List[str] = [] - cover_asset_id: Optional[str] = None # Not implemented yet + description: str | None = None + generation_ids: list[str] = [] + cover_asset_id: str | None = None # Not implemented yet @router.post("", response_model=AlbumResponse) async def create_album(request: Request, album_in: AlbumCreateRequest): @@ -32,7 +31,7 @@ async def create_album(request: Request, album_in: AlbumCreateRequest): album = await service.create_album(name=album_in.name, description=album_in.description) return AlbumResponse(**album.model_dump()) -@router.get("", response_model=List[AlbumResponse]) +@router.get("", response_model=list[AlbumResponse]) async def get_albums(request: Request, limit: int = 10, offset: int = 0): service: AlbumService = request.app.state.album_service albums = await service.get_albums(limit=limit, offset=offset) @@ -77,7 +76,7 @@ async def remove_generation_from_album(request: Request, album_id: str, generati raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Album or Generation not found") return {"status": "success"} -@router.get("/{album_id}/generations", response_model=List[GenerationResponse]) +@router.get("/{album_id}/generations", response_model=list[GenerationResponse]) async def get_album_generations(request: Request, album_id: str, limit: int = 10, offset: int = 0): service: AlbumService = request.app.state.album_service generations = await service.get_generations_by_album(album_id, limit=limit, offset=offset) diff --git a/api/endpoints/assets_router.py b/api/endpoints/assets_router.py index a1073c9..87172d0 100644 --- a/api/endpoints/assets_router.py +++ b/api/endpoints/assets_router.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Dict, Any +from typing import Any from aiogram.types import BufferedInputFile from bson import ObjectId @@ -135,22 +135,22 @@ async def delete_orphan_assets_from_minio( *, assets_collection: str = "assets", generations_collection: str = "generations", - asset_type: Optional[str] = "generated", - project_id: Optional[str] = None, + asset_type: str | None = "generated", + project_id: str | None = None, dry_run: bool = True, mark_assets_deleted: bool = False, batch_size: int = 500, -) -> Dict[str, Any]: +) -> dict[str, Any]: db = mongo['bot_db'] # БД уже выбрана в get_mongo_client assets = db[assets_collection] - match_assets: Dict[str, Any] = {} + match_assets: dict[str, Any] = {} if asset_type is not None: match_assets["type"] = asset_type if project_id is not None: match_assets["project_id"] = project_id - pipeline: List[Dict[str, Any]] = [ + pipeline: list[dict[str, Any]] = [ {"$match": match_assets} if match_assets else {"$match": {}}, { "$lookup": { @@ -192,8 +192,8 @@ async def delete_orphan_assets_from_minio( deleted_objects = 0 deleted_assets = 0 - errors: List[Dict[str, Any]] = [] - orphan_asset_ids: List[ObjectId] = [] + errors: list[dict[str, Any]] = [] + orphan_asset_ids: list[ObjectId] = [] async for asset in cursor: aid = asset["_id"] @@ -259,7 +259,7 @@ async def delete_asset( @router.get("", dependencies=[Depends(get_current_user)]) -async def get_assets(request: Request, dao: DAO = Depends(get_dao), type: Optional[str] = None, limit: int = 10, offset: int = 0, current_user: dict = Depends(get_current_user), project_id: Optional[str] = Depends(get_project_id)) -> AssetsResponse: +async def get_assets(request: Request, dao: DAO = Depends(get_dao), type: str | None = None, limit: int = 10, offset: int = 0, current_user: dict = Depends(get_current_user), project_id: str | None = Depends(get_project_id)) -> AssetsResponse: logger.info(f"get_assets called. Limit: {limit}, Offset: {offset}") user_id_filter = current_user["id"] @@ -286,10 +286,10 @@ async def get_assets(request: Request, dao: DAO = Depends(get_dao), type: Option @router.post("/upload", response_model=AssetResponse, status_code=status.HTTP_201_CREATED) async def upload_asset( file: UploadFile = File(...), - linked_char_id: Optional[str] = Form(None), + linked_char_id: str | None = Form(None), dao: DAO = Depends(get_dao), current_user: dict = Depends(get_current_user), - project_id: Optional[str] = Depends(get_project_id) + project_id: str | None = Depends(get_project_id) ): logger.info(f"upload_asset called. Filename: {file.filename}, ContentType: {file.content_type}, LinkedCharId: {linked_char_id}") if not file.content_type: diff --git a/api/endpoints/character_router.py b/api/endpoints/character_router.py index 67d0c07..fe47b2a 100644 --- a/api/endpoints/character_router.py +++ b/api/endpoints/character_router.py @@ -1,4 +1,4 @@ -from typing import List, Any, Coroutine, Optional +from typing import Any, Coroutine from fastapi import APIRouter, Depends from pydantic import BaseModel @@ -23,15 +23,15 @@ from api.dependency import get_project_id router = APIRouter(prefix="/api/characters", tags=["Characters"], dependencies=[Depends(get_current_user)]) -@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), + project_id: str | None = Depends(get_project_id), limit: int = 100, offset: int = 0 -) -> List[Character]: +) -> list[Character]: logger.info(f"get_characters called. Limit: {limit}, Offset: {offset}") user_id_filter = str(current_user["_id"]) @@ -102,7 +102,7 @@ async def get_character_by_id(character_id: str, request: Request, dao: DAO = De @router.post("/", response_model=Character) async def create_character( char_req: CharacterCreateRequest, - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), dao: DAO = Depends(get_dao), current_user: dict = Depends(get_current_user) ) -> Character: diff --git a/api/endpoints/environment_router.py b/api/endpoints/environment_router.py index 3b1f6b0..56caac6 100644 --- a/api/endpoints/environment_router.py +++ b/api/endpoints/environment_router.py @@ -1,5 +1,4 @@ import logging -from typing import List, Optional from fastapi import APIRouter, Depends, HTTPException from starlette import status @@ -50,7 +49,7 @@ async def create_environment( return created_env -@router.get("/character/{character_id}", response_model=List[Environment]) +@router.get("/character/{character_id}", response_model=list[Environment]) async def get_character_environments( character_id: str, dao: DAO = Depends(get_dao), diff --git a/api/endpoints/generation_router.py b/api/endpoints/generation_router.py index 641ef70..b7725a9 100644 --- a/api/endpoints/generation_router.py +++ b/api/endpoints/generation_router.py @@ -1,6 +1,5 @@ import logging import json -from typing import List, Optional from fastapi import APIRouter, UploadFile, File, Form, Header, HTTPException from fastapi.params import Depends @@ -30,7 +29,7 @@ logger = logging.getLogger(__name__) router = APIRouter(prefix='/api/generations', tags=["Generation"]) -async def check_project_access(project_id: Optional[str], current_user: dict, dao: DAO): +async def check_project_access(project_id: str | None, current_user: dict, dao: DAO): """Helper to check if user has access to project.""" if not project_id: return @@ -46,31 +45,36 @@ async def ask_prompt_assistant( current_user: dict = Depends(get_current_user) ) -> PromptResponse: logger.info(f"ask_prompt_assistant: {len(prompt_request.prompt)} chars") - generated_prompt = await generation_service.ask_prompt_assistant(prompt_request.prompt, prompt_request.linked_assets) + generated_prompt = await generation_service.ask_prompt_assistant( + prompt_request.prompt, + prompt_request.linked_assets, + prompt_request.model + ) return PromptResponse(prompt=generated_prompt) @router.post("/prompt-from-image", response_model=PromptResponse) async def prompt_from_image( - prompt: Optional[str] = Form(None), - images: List[UploadFile] = File(...), + prompt: str | None = Form(None), + model: str = Form("gemini-3.1-pro-preview"), + images: list[UploadFile] = File(...), generation_service: GenerationService = Depends(get_generation_service), current_user: dict = Depends(get_current_user) ) -> PromptResponse: images_bytes = [await img.read() for img in images] - generated_prompt = await generation_service.generate_prompt_from_images(images_bytes, prompt) + generated_prompt = await generation_service.generate_prompt_from_images(images_bytes, prompt, model) return PromptResponse(prompt=generated_prompt) @router.get("", response_model=GenerationsResponse) async def get_generations( - character_id: Optional[str] = None, + character_id: str | None = None, limit: int = 10, offset: int = 0, only_liked: bool = False, generation_service: GenerationService = Depends(get_generation_service), current_user: dict = Depends(get_current_user), - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), dao: DAO = Depends(get_dao) ): await check_project_access(project_id, current_user, dao) @@ -92,10 +96,10 @@ async def get_generations( @router.get("/usage", response_model=FinancialReport) async def get_usage_report( - breakdown: Optional[str] = None, # "user" or "project" + breakdown: str | None = 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), + project_id: str | None = Depends(get_project_id), dao: DAO = Depends(get_dao) ) -> FinancialReport: await check_project_access(project_id, current_user, dao) @@ -120,7 +124,7 @@ async def post_generation( generation: GenerationRequest, generation_service: GenerationService = Depends(get_generation_service), current_user: dict = Depends(get_current_user), - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), dao: DAO = Depends(get_dao) ) -> GenerationGroupResponse: await check_project_access(project_id, current_user, dao) @@ -137,7 +141,7 @@ async def post_generation( async def get_running_generations( generation_service: GenerationService = Depends(get_generation_service), current_user: dict = Depends(get_current_user), - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), dao: DAO = Depends(get_dao) ): await check_project_access(project_id, current_user, dao) diff --git a/api/endpoints/idea_router.py b/api/endpoints/idea_router.py index 7003fa1..137f55b 100644 --- a/api/endpoints/idea_router.py +++ b/api/endpoints/idea_router.py @@ -1,4 +1,3 @@ -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 @@ -14,7 +13,7 @@ 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), + project_id: str | None = Depends(get_project_id), current_user: dict = Depends(get_current_user), idea_service: IdeaService = Depends(get_idea_service) ): @@ -28,9 +27,9 @@ async def create_idea( inspiration_id=request.inspiration_id ) -@router.get("", response_model=List[IdeaResponse]) +@router.get("", response_model=list[IdeaResponse]) async def get_ideas( - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), limit: int = 20, offset: int = 0, current_user: dict = Depends(get_current_user), diff --git a/api/endpoints/inspiration_router.py b/api/endpoints/inspiration_router.py index e942016..e32ecba 100644 --- a/api/endpoints/inspiration_router.py +++ b/api/endpoints/inspiration_router.py @@ -1,4 +1,3 @@ -from typing import List, Optional from fastapi import APIRouter, Depends, HTTPException, status from api.dependency import get_inspiration_service, get_project_id @@ -13,7 +12,7 @@ router = APIRouter(prefix="/api/inspirations", tags=["Inspirations"]) @router.post("", response_model=InspirationResponse, status_code=status.HTTP_201_CREATED) async def create_inspiration( request: InspirationCreateRequest, - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), current_user: dict = Depends(get_current_user), service: InspirationService = Depends(get_inspiration_service) ): @@ -30,7 +29,7 @@ async def create_inspiration( @router.get("", response_model=InspirationListResponse) async def get_inspirations( - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), limit: int = 20, offset: int = 0, current_user: dict = Depends(get_current_user), diff --git a/api/endpoints/post_router.py b/api/endpoints/post_router.py index a9f3353..407bcff 100644 --- a/api/endpoints/post_router.py +++ b/api/endpoints/post_router.py @@ -1,4 +1,3 @@ -from typing import List, Optional from datetime import datetime from fastapi import APIRouter, Depends, HTTPException @@ -14,7 +13,7 @@ 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), + project_id: str | None = Depends(get_project_id), current_user: dict = Depends(get_current_user), post_service: PostService = Depends(get_post_service), ): @@ -28,13 +27,13 @@ async def create_post( ) -@router.get("", response_model=List[Post]) +@router.get("", response_model=list[Post]) async def get_posts( - project_id: Optional[str] = Depends(get_project_id), + project_id: str | None = Depends(get_project_id), limit: int = 200, offset: int = 0, - date_from: Optional[datetime] = None, - date_to: Optional[datetime] = None, + date_from: datetime | None = None, + date_to: datetime | None = None, current_user: dict = Depends(get_current_user), post_service: PostService = Depends(get_post_service), ): diff --git a/api/endpoints/project_router.py b/api/endpoints/project_router.py index 96375f9..1aa7635 100644 --- a/api/endpoints/project_router.py +++ b/api/endpoints/project_router.py @@ -1,4 +1,3 @@ -from typing import List, Optional from bson import ObjectId from fastapi import APIRouter, Depends, HTTPException, status @@ -12,7 +11,7 @@ router = APIRouter(prefix="/api/projects", tags=["Projects"]) class ProjectCreate(BaseModel): name: str - description: Optional[str] = None + description: str | None = None class ProjectMemberResponse(BaseModel): id: str @@ -21,9 +20,9 @@ class ProjectMemberResponse(BaseModel): class ProjectResponse(BaseModel): id: str name: str - description: Optional[str] = None + description: str | None = None owner_id: str - members: List[ProjectMemberResponse] + members: list[ProjectMemberResponse] is_owner: bool = False async def _get_project_response(project: Project, current_user_id: str, dao: DAO) -> ProjectResponse: @@ -78,7 +77,7 @@ async def create_project( return await _get_project_response(new_project, user_id, dao) -@router.get("", response_model=List[ProjectResponse]) +@router.get("", response_model=list[ProjectResponse]) async def get_my_projects( dao: DAO = Depends(get_dao), current_user: dict = Depends(get_current_user) diff --git a/api/models/AssetDTO.py b/api/models/AssetDTO.py index 8d46977..fa544d5 100644 --- a/api/models/AssetDTO.py +++ b/api/models/AssetDTO.py @@ -1,5 +1,4 @@ from datetime import datetime -from typing import List, Optional from pydantic import BaseModel @@ -11,10 +10,10 @@ class AssetResponse(BaseModel): name: str type: str # uploaded / generated content_type: str # image / prompt - linked_char_id: Optional[str] = None + linked_char_id: str | None = None created_at: datetime - url: Optional[str] = None + url: str | None = None class AssetsResponse(BaseModel): - assets: List[AssetResponse] + assets: list[AssetResponse] total_count: int \ No newline at end of file diff --git a/api/models/CharacterDTO.py b/api/models/CharacterDTO.py index ca73053..4cdb4d4 100644 --- a/api/models/CharacterDTO.py +++ b/api/models/CharacterDTO.py @@ -1,18 +1,17 @@ -from typing import Optional from pydantic import BaseModel class CharacterCreateRequest(BaseModel): name: str character_bio: str - character_image_doc_tg_id: Optional[str] = None - avatar_image: Optional[str] = None - character_image_tg_id: Optional[str] = None - project_id: Optional[str] = None + character_image_doc_tg_id: str | None = None + avatar_image: str | None = None + character_image_tg_id: str | None = None + project_id: str | None = None class CharacterUpdateRequest(BaseModel): - name: Optional[str] = None - character_bio: Optional[str] = None - character_image_doc_tg_id: Optional[str] = None - avatar_image: Optional[str] = None - character_image_tg_id: Optional[str] = None - project_id: Optional[str] = None + name: str | None = None + character_bio: str | None = None + character_image_doc_tg_id: str | None = None + avatar_image: str | None = None + character_image_tg_id: str | None = None + project_id: str | None = None diff --git a/api/models/EnvironmentRequest.py b/api/models/EnvironmentRequest.py index 430bc0a..bb12418 100644 --- a/api/models/EnvironmentRequest.py +++ b/api/models/EnvironmentRequest.py @@ -1,18 +1,17 @@ -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]] = [] + description: str | None = None + asset_ids: list[str] | None = [] class EnvironmentUpdate(BaseModel): - name: Optional[str] = Field(None, min_length=1) - description: Optional[str] = None - asset_ids: Optional[List[str]] = None + name: str | None = Field(None, min_length=1) + description: str | None = None + asset_ids: list[str] | None = None class AssetToEnvironment(BaseModel): @@ -20,4 +19,4 @@ class AssetToEnvironment(BaseModel): class AssetsToEnvironment(BaseModel): - asset_ids: List[str] + asset_ids: list[str] diff --git a/api/models/ExternalGenerationDTO.py b/api/models/ExternalGenerationDTO.py index 4744015..5b3feb0 100644 --- a/api/models/ExternalGenerationDTO.py +++ b/api/models/ExternalGenerationDTO.py @@ -1,4 +1,3 @@ -from typing import Optional from pydantic import BaseModel, Field from models.enums import AspectRatios, Quality @@ -7,29 +6,31 @@ class ExternalGenerationRequest(BaseModel): """Request model for importing external generations.""" prompt: str - tech_prompt: Optional[str] = None + tech_prompt: str | None = None # Image can be provided as base64 string OR URL (one must be provided) - image_data: Optional[str] = Field(None, description="Base64-encoded image data") - image_url: Optional[str] = Field(None, description="URL to download image from") + image_data: str | None = Field(None, description="Base64-encoded image data") + image_url: str | None = Field(None, description="URL to download image from") nsfw: bool = False # Generation metadata aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN # "1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9","21:9" quality: Quality = Quality.ONEK + model: str | None = None + seed: int | None = None # Optional linking - linked_character_id: Optional[str] = None + linked_character_id: str | None = None created_by: str = Field(..., description="User ID from external system") - project_id: Optional[str] = None + project_id: str | None = None # Performance metrics - execution_time_seconds: Optional[float] = None - api_execution_time_seconds: Optional[float] = None - token_usage: Optional[int] = None - input_token_usage: Optional[int] = None - output_token_usage: Optional[int] = None + execution_time_seconds: float | None = None + api_execution_time_seconds: float | None = None + token_usage: int | None = None + input_token_usage: int | None = None + output_token_usage: int | None = None def validate_image_source(self): """Ensure at least one image source is provided.""" diff --git a/api/models/FinancialUsageDTO.py b/api/models/FinancialUsageDTO.py index 90f1afd..fd163a9 100644 --- a/api/models/FinancialUsageDTO.py +++ b/api/models/FinancialUsageDTO.py @@ -1,5 +1,4 @@ from pydantic import BaseModel -from typing import List, Optional class UsageStats(BaseModel): total_runs: int @@ -9,10 +8,10 @@ class UsageStats(BaseModel): total_cost: float class UsageByEntity(BaseModel): - entity_id: Optional[str] = None + entity_id: str | None = None stats: UsageStats class FinancialReport(BaseModel): summary: UsageStats - by_user: Optional[List[UsageByEntity]] = None - by_project: Optional[List[UsageByEntity]] = None + by_user: list[UsageByEntity] | None = None + by_project: list[UsageByEntity] | None = None diff --git a/api/models/GenerationRequest.py b/api/models/GenerationRequest.py index e24d9c0..6f2713f 100644 --- a/api/models/GenerationRequest.py +++ b/api/models/GenerationRequest.py @@ -1,24 +1,24 @@ from datetime import datetime, UTC -from typing import List, Optional from pydantic import BaseModel, Field from models.Asset import Asset from models.Generation import GenerationStatus -from models.enums import AspectRatios, Quality, GenType +from models.enums import AspectRatios, Quality, GenType, ImageModel, TextModel class GenerationRequest(BaseModel): - linked_character_id: Optional[str] = None + linked_character_id: str | None = None aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN # "1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9","21:9" quality: Quality = Quality.ONEK prompt: str - telegram_id: Optional[int] = None + model: ImageModel = Field(default=ImageModel.GEMINI_3_PRO_IMAGE_PREVIEW) + telegram_id: int | None = None use_profile_image: bool = True - assets_list: List[str] - environment_id: Optional[str] = None - project_id: Optional[str] = None - idea_id: Optional[str] = None + assets_list: list[str] + environment_id: str | None = None + project_id: str | None = None + idea_id: str | None = None nsfw: bool = False count: int = Field(default=1, ge=1, le=10) @@ -28,33 +28,35 @@ class NsfwRequest(BaseModel): class GenerationsResponse(BaseModel): - generations: List["GenerationResponse"] + generations: list["GenerationResponse"] total_count: int class GenerationResponse(BaseModel): id: str status: GenerationStatus - failed_reason: Optional[str] = None + failed_reason: str | None = None project_id: str | None = None - linked_character_id: Optional[str] = None + linked_character_id: str | None = None aspect_ratio: AspectRatios quality: Quality prompt: str - tech_prompt: Optional[str] = None - assets_list: List[str] - result_list: List[str] = [] - result: Optional[str] = None - execution_time_seconds: Optional[float] = None - api_execution_time_seconds: Optional[float] = None - token_usage: Optional[int] = None - input_token_usage: Optional[int] = None - output_token_usage: Optional[int] = None + model: ImageModel | None = None + seed: int | None = None + tech_prompt: str | None = None + assets_list: list[str] + result_list: list[str] = [] + result: str | None = None + execution_time_seconds: float | None = None + api_execution_time_seconds: float | None = None + token_usage: int | None = None + input_token_usage: int | None = None + output_token_usage: int | None = None progress: int = 0 - cost: Optional[float] = None - created_by: Optional[str] = None - generation_group_id: Optional[str] = None - idea_id: Optional[str] = None + cost: float | None = None + created_by: str | None = None + generation_group_id: str | None = None + idea_id: str | None = None likes_count: int = 0 is_liked: bool = False nsfw: bool = False @@ -64,12 +66,13 @@ class GenerationResponse(BaseModel): class GenerationGroupResponse(BaseModel): generation_group_id: str - generations: List[GenerationResponse] + generations: list[GenerationResponse] class PromptRequest(BaseModel): prompt: str - linked_assets: List[str] = [] + model: TextModel = Field(default=TextModel.GEMINI_3_1_PRO_PREVIEW) + linked_assets: list[str] = [] class PromptResponse(BaseModel): diff --git a/api/models/IdeaRequest.py b/api/models/IdeaRequest.py index 327adbb..436b679 100644 --- a/api/models/IdeaRequest.py +++ b/api/models/IdeaRequest.py @@ -1,18 +1,17 @@ -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 - inspiration_id: Optional[str] = None + description: str | None = None + project_id: str | None = None # Optional in body if passed via header/dependency + inspiration_id: str | None = None class IdeaUpdateRequest(BaseModel): - name: Optional[str] = None - description: Optional[str] = None - inspiration_id: Optional[str] = None + name: str | None = None + description: str | None = None + inspiration_id: str | None = None class IdeaResponse(Idea): - last_generation: Optional[GenerationResponse] = None + last_generation: GenerationResponse | None = None diff --git a/api/models/InspirationRequest.py b/api/models/InspirationRequest.py index 802dc3b..3e26417 100644 --- a/api/models/InspirationRequest.py +++ b/api/models/InspirationRequest.py @@ -1,5 +1,4 @@ from datetime import datetime -from typing import List, Optional from pydantic import BaseModel @@ -8,8 +7,8 @@ from models.Inspiration import Inspiration class InspirationCreateRequest(BaseModel): source_url: str - caption: Optional[str] = None - project_id: Optional[str] = None + caption: str | None = None + project_id: str | None = None class InspirationResponse(BaseModel): @@ -25,5 +24,5 @@ class InspirationResponse(BaseModel): class InspirationListResponse(BaseModel): - inspirations: List[InspirationResponse] + inspirations: list[InspirationResponse] total_count: int diff --git a/api/models/PostRequest.py b/api/models/PostRequest.py index 18be225..7eb5cc4 100644 --- a/api/models/PostRequest.py +++ b/api/models/PostRequest.py @@ -1,19 +1,18 @@ 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 + generation_ids: list[str] = [] + project_id: str | None = None class PostUpdateRequest(BaseModel): - date: Optional[datetime] = None - topic: Optional[str] = None + date: datetime | None = None + topic: str | None = None class AddGenerationsRequest(BaseModel): - generation_ids: List[str] + generation_ids: list[str] diff --git a/api/service/generation_service.py b/api/service/generation_service.py index 819bc53..c0c1841 100644 --- a/api/service/generation_service.py +++ b/api/service/generation_service.py @@ -34,6 +34,7 @@ async def generate_image_task( media_group_bytes: List[bytes], aspect_ratio: AspectRatios, quality: Quality, + model: str, gemini: GoogleAdapter, ) -> Tuple[List[bytes], Dict[str, Any]]: """ @@ -47,6 +48,7 @@ async def generate_image_task( images_list=media_group_bytes, aspect_ratio=aspect_ratio, quality=quality, + model=model, ) generated_images_io, metrics = result logger.info(f"generate_image_task completed, received {len(generated_images_io) if generated_images_io else 0} images") @@ -75,7 +77,7 @@ class GenerationService: # --- Public API --- - async def ask_prompt_assistant(self, prompt: str, assets: list[str] | None = None) -> str: + async def ask_prompt_assistant(self, prompt: str, assets: list[str] | None = None, model: str = "gemini-3.1-pro-preview") -> str: 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. " @@ -87,17 +89,17 @@ class GenerationService: assets_db = await self.dao.assets.get_assets_by_ids(assets) assets_data.extend(asset.data for asset in assets_db if asset.data) - generated_prompt = await asyncio.to_thread(self.gemini.generate_text, future_prompt, assets_data) + generated_prompt = await asyncio.to_thread(self.gemini.generate_text, future_prompt, model, assets_data) logger.info(f"Prompt Assistant: {generated_prompt}") return generated_prompt - async def generate_prompt_from_images(self, images: List[bytes], user_prompt: Optional[str] = None) -> str: + async def generate_prompt_from_images(self, images: List[bytes], user_prompt: Optional[str] = None, model: str = "gemini-3.1-pro-preview") -> str: technical_prompt = "You are a prompt engineer. Describe this image in detail to create a stable diffusion using this image as reference. " if user_prompt: technical_prompt += f"User also provided this context: {user_prompt}. " technical_prompt += "Provide ONLY the detailed prompt." - 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, model=model, images_list=images) async def get_generations(self, **kwargs) -> GenerationsResponse: current_user_id = kwargs.pop('current_user_id', None) @@ -162,6 +164,7 @@ class GenerationService: media_group_bytes=media_group_bytes, aspect_ratio=generation.aspect_ratio, quality=generation.quality, + model=generation.model or "gemini-3-pro-image-preview", gemini=self.gemini ) self._update_generation_metrics(generation, metrics) @@ -205,7 +208,9 @@ class GenerationService: aspect_ratio=external_gen.aspect_ratio, quality=external_gen.quality, prompt=external_gen.prompt, + model=external_gen.model, tech_prompt=external_gen.tech_prompt, + seed=external_gen.seed, result_list=[new_asset.id], result=new_asset.id, progress=100, diff --git a/models/Album.py b/models/Album.py index 26e6b18..b3bae59 100644 --- a/models/Album.py +++ b/models/Album.py @@ -1,12 +1,11 @@ from datetime import datetime, UTC -from typing import Optional, List from pydantic import BaseModel, Field class Album(BaseModel): - id: Optional[str] = None + id: str | None = None name: str - description: Optional[str] = None - cover_asset_id: Optional[str] = None - generation_ids: List[str] = [] + description: str | None = None + cover_asset_id: str | None = None + generation_ids: list[str] = [] created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/Asset.py b/models/Asset.py index 29ab56e..496e24b 100644 --- a/models/Asset.py +++ b/models/Asset.py @@ -1,6 +1,6 @@ from datetime import datetime, UTC from enum import Enum -from typing import Optional, Any, List +from typing import Any from pydantic import BaseModel, computed_field, Field, model_validator @@ -17,21 +17,21 @@ class AssetType(str, Enum): class Asset(BaseModel): - id: Optional[str] = None + id: str | None = None name: str type: AssetType = AssetType.GENERATED content_type: AssetContentType = AssetContentType.IMAGE - linked_char_id: Optional[str] = None - data: Optional[bytes] = None - tg_doc_file_id: Optional[str] = None - tg_photo_file_id: Optional[str] = None - minio_object_name: Optional[str] = None - minio_bucket: Optional[str] = None - minio_thumbnail_object_name: Optional[str] = None - thumbnail: Optional[bytes] = None - tags: List[str] = [] - created_by: Optional[str] = None - project_id: Optional[str] = None + linked_char_id: str | None = None + data: bytes | None = None + tg_doc_file_id: str | None = None + tg_photo_file_id: str | None = None + minio_object_name: str | None = None + minio_bucket: str | None = None + minio_thumbnail_object_name: str | None = None + thumbnail: bytes | None = None + tags: list[str] = [] + created_by: str | None = None + project_id: str | None = None is_deleted: bool = False created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/Character.py b/models/Character.py index f174659..d1fdc8c 100644 --- a/models/Character.py +++ b/models/Character.py @@ -1,16 +1,15 @@ -from typing import Optional from pydantic import BaseModel from pydantic_core.core_schema import computed_field class Character(BaseModel): - id: Optional[str] = None + id: str | None = None name: str - avatar_asset_id: Optional[str] = None - avatar_image: Optional[str] = None - character_image_doc_tg_id: Optional[str] = None - character_image_tg_id: Optional[str] = None - character_bio: Optional[str] = None - created_by: Optional[str] = None - project_id: Optional[str] = None + avatar_asset_id: str | None = None + avatar_image: str | None = None + character_image_doc_tg_id: str | None = None + character_image_tg_id: str | None = None + character_bio: str | None = None + created_by: str | None = None + project_id: str | None = None diff --git a/models/Environment.py b/models/Environment.py index 0b1e3f1..21ffb63 100644 --- a/models/Environment.py +++ b/models/Environment.py @@ -1,15 +1,14 @@ -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") + id: str | None = 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) + description: str | None = 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) diff --git a/models/Generation.py b/models/Generation.py index e3d2b70..51ef27a 100644 --- a/models/Generation.py +++ b/models/Generation.py @@ -1,11 +1,9 @@ from datetime import datetime, UTC from enum import Enum -from typing import List, Optional from pydantic import BaseModel, Field, computed_field -from models.Asset import Asset -from models.enums import AspectRatios, Quality, GenType +from models.enums import AspectRatios, Quality class GenerationStatus(str, Enum): @@ -14,33 +12,35 @@ class GenerationStatus(str, Enum): FAILED = "failed" class Generation(BaseModel): - id: Optional[str] = None + id: str | None = None status: GenerationStatus = GenerationStatus.RUNNING - failed_reason: Optional[str] = None - linked_character_id: Optional[str] = None - telegram_id: Optional[int] = None + failed_reason: str | None = None + linked_character_id: str | None = None + telegram_id: int | None = None use_profile_image: bool = True aspect_ratio: AspectRatios quality: Quality prompt: str - tech_prompt: Optional[str] = None - assets_list: List[str] = Field(default_factory=list) - result_list: List[str] = Field(default_factory=list) - result: Optional[str] = None + model: str | None = None + seed: int | None = None + tech_prompt: str | None = None + assets_list: list[str] = Field(default_factory=list) + result_list: list[str] = Field(default_factory=list) + result: str | None = None progress: int = 0 - execution_time_seconds: Optional[float] = None - api_execution_time_seconds: Optional[float] = None - token_usage: Optional[int] = None - input_token_usage: Optional[int] = None - output_token_usage: Optional[int] = None + execution_time_seconds: float | None = None + api_execution_time_seconds: float | None = None + token_usage: int | None = None + input_token_usage: int | None = None + output_token_usage: int | None = None is_deleted: bool = False - 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) - project_id: Optional[str] = None - idea_id: Optional[str] = None - liked_by: List[str] = Field(default_factory=list) + album_id: str | None = None + environment_id: str | None = None + generation_group_id: str | None = None + created_by: str | None = None # Stores User ID (Telegram ID or Web User ObjectId) + project_id: str | None = None + idea_id: str | None = None + liked_by: list[str] = Field(default_factory=list) nsfw: bool = False created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/Idea.py b/models/Idea.py index b65093b..56985ec 100644 --- a/models/Idea.py +++ b/models/Idea.py @@ -1,13 +1,12 @@ from datetime import datetime -from typing import Optional, List from pydantic import BaseModel, Field class Idea(BaseModel): - id: Optional[str] = None + id: str | None = None name: str = "New Idea" - description: Optional[str] = None - project_id: Optional[str] = None - inspiration_id: Optional[str] = None # Link to Inspiration + description: str | None = None + project_id: str | None = None + inspiration_id: str | None = None # Link to Inspiration created_by: str # User ID is_deleted: bool = False created_at: datetime = Field(default_factory=datetime.now) diff --git a/models/Inspiration.py b/models/Inspiration.py index 0649772..bb75b7d 100644 --- a/models/Inspiration.py +++ b/models/Inspiration.py @@ -1,5 +1,4 @@ from datetime import datetime, UTC -from typing import Optional from pydantic import BaseModel, Field diff --git a/models/Post.py b/models/Post.py index 4ab5837..2598098 100644 --- a/models/Post.py +++ b/models/Post.py @@ -1,14 +1,13 @@ 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 + id: str | None = None date: datetime topic: str - generation_ids: List[str] = Field(default_factory=list) - project_id: Optional[str] = None + generation_ids: list[str] = Field(default_factory=list) + project_id: str | None = None created_by: str is_deleted: bool = False created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/models/Project.py b/models/Project.py index 65bbb59..34815ac 100644 --- a/models/Project.py +++ b/models/Project.py @@ -1,12 +1,11 @@ from datetime import datetime -from typing import List, Optional from pydantic import BaseModel, Field class Project(BaseModel): - id: Optional[str] = None + id: str | None = None name: str - description: Optional[str] = None + description: str | None = None owner_id: str - members: List[str] = [] # List of User IDs + members: list[str] = [] # List of User IDs is_deleted: bool = False created_at: datetime = Field(default_factory=datetime.now) diff --git a/models/enums.py b/models/enums.py index cfc7401..0401bb1 100644 --- a/models/enums.py +++ b/models/enums.py @@ -52,3 +52,20 @@ class GenType(str, Enum): GenType.TEXT: 'Text', GenType.IMAGE: 'Image', }[self] + + +class TextModel(str, Enum): + GEMINI_3_1_PRO_PREVIEW = "gemini-3.1-pro-preview" + + @property + def value_model(self) -> str: + return self.value + + +class ImageModel(str, Enum): + GEMINI_3_PRO_IMAGE_PREVIEW = "gemini-3-pro-image-preview" + GEMINI_3_1_FLASH_IMAGE_PREVIEW = "gemini-3.1-flash-image-preview" + + @property + def value_model(self) -> str: + return self.value