create UserIndex

src/endpoints/generate.py

@@ -7,7 +7,7 @@ from src.libs.image import save_img, generate_img_index
 from src.libs.model import CganCols, get_model
 from src.libs.s3 import s3client
 from src.models.generate import GenerateResult, ImageResult
-from src.models.main import User, Method
+from src.models.main import User, Method, UserIndex
 
 IMAGE_STORE_PATH = os.path.abspath("./src/store")
 BUCKET_NAME = "pimthaigans-image-container"
@@ -36,12 +36,12 @@ async def status():
 
 
 @router.post("/")
-async def generate(user: User, index: int | None = None) -> GenerateResult:
-    if index or index == 0:
-        result: GenerateResult = await generate_index(user, index)
+async def generate(user: UserIndex) -> GenerateResult:
+    if user.method == Method.index:
+        result: GenerateResult = await generate_index(user.user, user.index)
         return result
 
-    result: GenerateResult = await generate_all(user)
+    result: GenerateResult = await generate_all(user.user)
     return result
 
 

src/models/main.py

@@ -11,3 +11,9 @@ class User(BaseModel):
 class Method(str, Enum):
     all = "all"
     index = "index"
+
+
+class UserIndex(BaseModel):
+    user: User
+    method: Method
+    index: int