ai-char/ai-char-bot
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

models + refactor

Browse Source
This commit is contained in:
xds
2026-02-27 20:37:24 +03:00
parent d9caececd7
commit e011805186
31 changed files with 234 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
adapters/google_adapter.py
View File

@@ -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'],

4
api/endpoints/admin.py
View File

@@ -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)]

17
api/endpoints/album_router.py
View File

@@ -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)

22
api/endpoints/assets_router.py
View File

@@ -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:

10
api/endpoints/character_router.py
View File

@@ -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:

3
api/endpoints/environment_router.py
View File

@@ -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),

28
api/endpoints/generation_router.py
View File

@@ -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)

7
api/endpoints/idea_router.py
View File

@@ -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),

5
api/endpoints/inspiration_router.py
View File

@@ -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),

11
api/endpoints/post_router.py
View File

@@ -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),
):

9
api/endpoints/project_router.py
View File

@@ -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)

7
api/models/AssetDTO.py
View File

@@ -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

21
api/models/CharacterDTO.py
View File

@@ -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

13
api/models/EnvironmentRequest.py
View File

@@ -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]

23
api/models/ExternalGenerationDTO.py
View File

@@ -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."""

7
api/models/FinancialUsageDTO.py
View File

@@ -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

55
api/models/GenerationRequest.py
View File

@@ -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):

15
api/models/IdeaRequest.py
View File

@@ -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

7
api/models/InspirationRequest.py
View File

@@ -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

11
api/models/PostRequest.py
View File

@@ -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]

13
api/service/generation_service.py
View File

@@ -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,

9
models/Album.py
View File

@@ -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))

26
models/Asset.py
View File

@@ -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))

17
models/Character.py
View File

@@ -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

7
models/Environment.py
View File

@@ -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)

46
models/Generation.py
View File

@@ -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))

9
models/Idea.py
View File

@@ -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)

1
models/Inspiration.py
View File

@@ -1,5 +1,4 @@
from datetime import datetime, UTC
from typing import Optional
from pydantic import BaseModel, Field

7
models/Post.py
View File

@@ -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))

7
models/Project.py
View File

@@ -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)

17
models/enums.py
View File

@@ -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