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

adapters/google_adapter.py

@@ -131,9 +131,17 @@ class GoogleAdapter:
             else:
                 logger.warning("No images text generated from parts")
 
+            input_tokens = 0
+            output_tokens = 0
+            if response.usage_metadata:
+                input_tokens = response.usage_metadata.prompt_token_count
+                output_tokens = response.usage_metadata.candidates_token_count
+
             metrics = {
                 "api_execution_time_seconds": api_duration,
-                "token_usage": token_usage
+                "token_usage": token_usage,
+                "input_token_usage": input_tokens,
+                "output_token_usage": output_tokens
             }
             return generated_images, metrics
 

api/models/GenerationRequest.py

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

api/service/generation_service.py

@@ -154,6 +154,7 @@ 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")
+            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}"""
 
@@ -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")

models/Generation.py

@@ -19,6 +19,7 @@ class Generation(BaseModel):
     failed_reason: Optional[str] = None
     linked_character_id: Optional[str] = None
     telegram_id: Optional[int] = None
+    use_profile_image: bool = True
     aspect_ratio: AspectRatios
     quality: Quality
     prompt: str
@@ -29,6 +30,8 @@ class Generation(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
     
     created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
     updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))