55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
from datetime import datetime, UTC
|
|
from enum import Enum
|
|
|
|
from pydantic import BaseModel, Field, computed_field
|
|
|
|
from models.enums import AspectRatios, Quality
|
|
|
|
|
|
class GenerationStatus(str, Enum):
|
|
RUNNING = "running"
|
|
DONE = "done"
|
|
FAILED = "failed"
|
|
|
|
class Generation(BaseModel):
|
|
id: str | None = None
|
|
status: GenerationStatus = GenerationStatus.RUNNING
|
|
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
|
|
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: 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: 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))
|
|
|
|
@computed_field
|
|
def cost(self) -> float:
|
|
if self.status == GenerationStatus.DONE and self.input_token_usage and self.output_token_usage:
|
|
cost_input = self.input_token_usage * 0.000002
|
|
cost_output = self.output_token_usage * 0.00012
|
|
return round(cost_input + cost_output, 3)
|
|
return 0.0
|