feat: Add created_by and cost fields to generation models, populate created_by from the authenticated user, and implement cost calculation.

api/endpoints/generation_router.py

@@ -56,10 +56,10 @@ async def get_generations(character_id: Optional[str] = None, limit: int = 10, o
 
 @router.post("/_run", response_model=GenerationResponse)
 async def post_generation(generation: GenerationRequest, request: Request,
-                          generation_service: GenerationService = Depends(
+                          generation_service: GenerationService = Depends(get_generation_service),
-                              get_generation_service)) -> GenerationResponse:
+                          current_user: dict = Depends(get_current_user)) -> GenerationResponse:
     logger.info(f"post_generation (run) called. LinkedCharId: {generation.linked_character_id}, PromptLength: {len(generation.prompt)}")
-    return await generation_service.create_generation_task(generation)
+    return await generation_service.create_generation_task(generation, user_id=current_user.get("username"))
 
 
 @router.get("/{generation_id}", response_model=GenerationResponse)

api/models/GenerationRequest.py

@@ -42,6 +42,8 @@ class GenerationResponse(BaseModel):
     input_token_usage: Optional[int] = None
     output_token_usage: Optional[int] = None
     progress: int = 0
+    cost: Optional[float] = None
+    created_by: Optional[str] = None
     created_at: datetime = datetime.now(UTC)
     updated_at: datetime = datetime.now(UTC)
 

api/service/generation_service.py

@@ -109,12 +109,15 @@ class GenerationService:
     async def get_running_generations(self) -> List[Generation]:
         return await self.dao.generations.get_generations(status=GenerationStatus.RUNNING)
 
-    async def create_generation_task(self, generation_request: GenerationRequest) -> GenerationResponse:
+    async def create_generation_task(self, generation_request: GenerationRequest, user_id: Optional[str] = None) -> GenerationResponse:
         gen_id = None
         generation_model = None
 
         try:
             generation_model = Generation(**generation_request.model_dump())
+            if user_id:
+                generation_model.created_by = user_id
+            
             gen_id = await self.dao.generations.create_generation(generation_model)
             generation_model.id = gen_id
 

models/Generation.py

@@ -2,7 +2,7 @@ from datetime import datetime, UTC
 from enum import Enum
 from typing import List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, computed_field
 
 from models.Asset import Asset
 from models.enums import AspectRatios, Quality, GenType
@@ -34,5 +34,15 @@ class Generation(BaseModel):
     input_token_usage: Optional[int] = None
     output_token_usage: Optional[int] = None
     is_deleted: bool = False
+    album_id: Optional[str] = None
+    created_by: Optional[str] = None
     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:
+            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