Add response schema

app.py

@@ -1,19 +1,11 @@
 import os
 from typing import Union
 
-from fastapi import (
-    Depends,
-    FastAPI,
-    File,
-    HTTPException,
-    Query,
-    Response,
-    UploadFile,
-    status,
-)
+from fastapi import Depends, FastAPI, File, HTTPException, Response, UploadFile, status
 from fastapi.security import APIKeyQuery
 
 from model import Model
+from schema import EmbeddingResponse, SimilarityResponse
 
 app = FastAPI(
     title="Facial Expression Embedding Service",
@@ -42,6 +34,7 @@ async def validate_token(
     "/embed",
     status_code=status.HTTP_200_OK,
     dependencies=[Depends(validate_token)],
+    response_model=EmbeddingResponse,
 )
 async def calculate_embedding(
     image: UploadFile = File(...),
@@ -49,7 +42,7 @@ async def calculate_embedding(
     try:
         image_content = await image.read()
         pred = model.predict(model.preprocess(image_content))
-        return {"embedding": pred.tolist()}
+        return EmbeddingResponse(embedding=pred.tolist())
     except Exception as e:
         return Response(
             content=str(e), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
@@ -60,6 +53,7 @@ async def calculate_embedding(
     "/similarity",
     status_code=status.HTTP_200_OK,
     dependencies=[Depends(validate_token)],
+    response_model=SimilarityResponse,
 )
 async def calculate_similarity_score(
     image1: UploadFile = File(...),
@@ -70,7 +64,7 @@ async def calculate_similarity_score(
         image2_content = await image2.read()
         pred1 = model.predict(model.preprocess(image1_content))
         pred2 = model.predict(model.preprocess(image2_content))
-        return {"score": float(model.distance(pred1, pred2))}
+        return SimilarityResponse(score=float(model.distance(pred1, pred2)))
     except Exception as e:
         return Response(
             content=str(e), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR

schema.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class EmbeddingResponse:
+    embedding: List[float]
+
+
+@dataclass
+class SimilarityResponse:
+    score: float