feat: Add use_profile_image and detailed token usage fields to generation models.

This commit is contained in:
xds
2026-02-05 22:19:57 +03:00
parent 002c949f08
commit c0debab0cb
8 changed files with 21 additions and 3 deletions

View File

@@ -14,6 +14,7 @@ class GenerationRequest(BaseModel):
quality: Quality = Quality.ONEK
prompt: str
telegram_id: Optional[int] = None
use_profile_image: bool = True
assets_list: List[str]
@@ -32,6 +33,8 @@ class GenerationResponse(BaseModel):
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
progress: int = 0
created_at: datetime = datetime.now(UTC)
updated_at: datetime = datetime.now(UTC)

View File

@@ -154,7 +154,8 @@ class GenerationService:
char_info = await self.dao.chars.get_character(generation.linked_character_id, with_image_data=True)
if char_info is None:
raise Exception(f"Character ID {generation.linked_character_id} not found")
media_group_bytes.append(char_info.character_image_data)
if generation.use_profile_image:
media_group_bytes.append(char_info.character_image_data)
generation_prompt = f"""You are creating image for {char_info.character_bio}"""
reference_assets = await self.dao.assets.get_assets_by_ids(generation.assets_list)
@@ -182,9 +183,12 @@ class GenerationService:
gemini=self.gemini
)
# Update metrics from API (Common for both)
# Update metrics from API (Common for both)
generation.api_execution_time_seconds = metrics.get("api_execution_time_seconds")
generation.token_usage = metrics.get("token_usage")
generation.input_token_usage = metrics.get("input_token_usage")
generation.output_token_usage = metrics.get("output_token_usage")
except GoogleGenerationException as e:
generation.status = GenerationStatus.FAILED
@@ -247,7 +251,7 @@ class GenerationService:
end_time = datetime.now()
generation.execution_time_seconds = (end_time - start_time).total_seconds()
logger.info(f"DEBUG: Saving generation {generation.id}. Metrics: api_exec={generation.api_execution_time_seconds}, tokens={generation.token_usage}, exec={generation.execution_time_seconds}")
logger.info(f"DEBUG: Saving generation {generation.id}. Metrics: api_exec={generation.api_execution_time_seconds}, tokens={generation.token_usage}, in_tokens={generation.input_token_usage}, out_tokens={generation.output_token_usage}, exec={generation.execution_time_seconds}")
await self.dao.generations.update_generation(generation)
logger.info(f"Generation {generation.id} completed successfully. {len(created_assets)} assets created. Total Time: {generation.execution_time_seconds:.2f}s")