diff --git a/__pycache__/main.cpython-313.pyc b/__pycache__/main.cpython-313.pyc index 3c3528c..7e90820 100644 Binary files a/__pycache__/main.cpython-313.pyc and b/__pycache__/main.cpython-313.pyc differ diff --git a/api/endpoints/__pycache__/generation_router.cpython-313.pyc b/api/endpoints/__pycache__/generation_router.cpython-313.pyc index c9920bb..e785c8f 100644 Binary files a/api/endpoints/__pycache__/generation_router.cpython-313.pyc and b/api/endpoints/__pycache__/generation_router.cpython-313.pyc differ diff --git a/api/endpoints/generation_router.py b/api/endpoints/generation_router.py index 0c3152b..cfdf4e8 100644 --- a/api/endpoints/generation_router.py +++ b/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( - get_generation_service)) -> GenerationResponse: + generation_service: GenerationService = Depends(get_generation_service), + 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) diff --git a/api/models/GenerationRequest.py b/api/models/GenerationRequest.py index 45e9d06..2ad06cd 100644 --- a/api/models/GenerationRequest.py +++ b/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) diff --git a/api/models/__pycache__/GenerationRequest.cpython-313.pyc b/api/models/__pycache__/GenerationRequest.cpython-313.pyc index daf9354..8dbd694 100644 Binary files a/api/models/__pycache__/GenerationRequest.cpython-313.pyc and b/api/models/__pycache__/GenerationRequest.cpython-313.pyc differ diff --git a/api/service/__pycache__/generation_service.cpython-313.pyc b/api/service/__pycache__/generation_service.cpython-313.pyc index 4d22de6..7bfe1a3 100644 Binary files a/api/service/__pycache__/generation_service.cpython-313.pyc and b/api/service/__pycache__/generation_service.cpython-313.pyc differ diff --git a/api/service/generation_service.py b/api/service/generation_service.py index 24d6421..b157c64 100644 --- a/api/service/generation_service.py +++ b/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 diff --git a/models/Generation.py b/models/Generation.py index 88e87f7..ea4500b 100644 --- a/models/Generation.py +++ b/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 \ No newline at end of file diff --git a/models/__pycache__/Generation.cpython-313.pyc b/models/__pycache__/Generation.cpython-313.pyc index 7353cf5..0457de5 100644 Binary files a/models/__pycache__/Generation.cpython-313.pyc and b/models/__pycache__/Generation.cpython-313.pyc differ diff --git a/repos/__pycache__/dao.cpython-313.pyc b/repos/__pycache__/dao.cpython-313.pyc index 8e0d2f5..f0b61bb 100644 Binary files a/repos/__pycache__/dao.cpython-313.pyc and b/repos/__pycache__/dao.cpython-313.pyc differ diff --git a/repos/__pycache__/generation_repo.cpython-313.pyc b/repos/__pycache__/generation_repo.cpython-313.pyc index 94910e4..0c9a3be 100644 Binary files a/repos/__pycache__/generation_repo.cpython-313.pyc and b/repos/__pycache__/generation_repo.cpython-313.pyc differ