feat: Add created_by and cost fields to generation models, populate created_by from the authenticated user, and implement cost calculation.
This commit is contained in:
Binary file not shown.
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user