add run_id, refactor output response

Author: TheAIWizard

src/api/main.py

@@ -24,6 +24,7 @@ async def lifespan(app: FastAPI):
     logger.info("🚀 Starting API lifespan")
 
     app.state.model = load_model()
+    app.state.run_id = app.state.model.metadata.run_id
 
     yield
     logger.info("🛑 Shutting down API lifespan")

src/api/models/responses.py

@@ -1,6 +1,6 @@
-from typing import Dict, Union
+from typing import Any, Dict, Mapping, Union
 
-from pydantic import BaseModel, RootModel
+from pydantic import BaseModel, RootModel, model_validator
 
 
 class Prediction(BaseModel):
@@ -9,5 +9,72 @@ class Prediction(BaseModel):
     libelle: str
 
 
-class PredictionResponse(RootModel[Dict[str, Union[Prediction, float, str]]]):
-    pass
+class OutputResponse(RootModel[Dict[str, Union[Prediction, float, str]]]):
+    """
+    Contract for the output response of the API including:
+    - KV of PredictionResponse: normalized prediction responses generated by the model artifact
+    - MLversion: run_id as version of the ML model
+
+    Expected flat structure after normalization:
+
+    {
+    "1": Prediction,
+    "2": Prediction,
+    ...,
+    "IC": float,      # required confidence score
+    "MLversion": str  # required run_id as model version
+    }
+
+    Notes:
+    - The output reflects what the model artifact produces, but the API applies
+    `model_dump()` in `predict()` before returning to ensure schema consistency.
+    - Any changes to the output schema (e.g., new fields, renaming) must be documented
+    both here and in training repo (codif-ape-training) to maintain API contract clarity.
+    """
+
+    @model_validator(mode="after")
+    @classmethod
+    def _normalize(cls, data: Any) -> "OutputResponse":
+        # unwrap root if called with an instance
+        raw = data.root if isinstance(data, cls) else data
+
+        if not isinstance(raw, Mapping):
+            raise TypeError("OutputResponse: expected a dict/mapping")
+
+        # IC (required) - accept numbers or numeric strings
+        try:
+            ic = float(raw["IC"])
+        except KeyError:
+            raise ValueError("OutputResponse: missing required key 'IC'")
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"OutputResponse: 'IC' not convertible to float: {e}") from e
+
+        # MLversion (required)
+        try:
+            ml_version = str(raw["MLversion"])
+        except KeyError:
+            raise ValueError("OutputResponse: missing required key 'MLversion'")
+        except Exception as e:
+            raise ValueError(f"OutputResponse: 'MLversion' not convertible to str: {e}") from e
+
+        # allow only digit keys + IC + MLversion
+        allowed = {k for k in raw.keys() if k.isdigit()} | {"IC", "MLversion"}
+        extra = set(raw.keys()) - allowed
+        if extra:
+            raise ValueError(f"OutputResponse: unexpected keys: {sorted(extra)}")
+
+        # ensure digit keys map to Prediction
+        for k in (k for k in raw.keys() if k.isdigit()):
+            val = raw[k]
+            if not isinstance(val, (Mapping, Prediction)):
+                raise ValueError(
+                    f"OutputResponse: value for key '{k}'must be a mapping or Prediction"
+                )
+
+        # normalize in place
+        normalized = dict(raw)
+        normalized["IC"] = ic
+        normalized["MLversion"] = ml_version
+
+        data.root = normalized
+        return data

src/api/routes/predict.py

@@ -4,13 +4,13 @@
 from fastapi.security import HTTPBasicCredentials
 
 from api.models.forms import BatchForms
-from api.models.responses import PredictionResponse
+from api.models.responses import OutputResponse
 from utils.security import get_credentials
 
 router = APIRouter(prefix="/predict", tags=["Predict NACE code for a list of activities"])
 
 
-@router.post("/", response_model=List[PredictionResponse])
+@router.post("/", response_model=List[OutputResponse])
 async def predict(
     credentials: Annotated[HTTPBasicCredentials, Depends(get_credentials)],
     request: Request,
@@ -27,12 +27,17 @@ async def predict(
         credentials (HTTPBasicCredentials): The credentials for authentication.
         forms (Forms): The input data in the form of Forms object.
         nb_echos_max (int, optional): The maximum number of predictions to return. Defaults to 5.
-        prob_min (float, optional): The minimum probability threshold for predictions. Defaults to 0.01.
-        num_workers (int, optional): Number of CPU for multiprocessing in Dataloader. Defaults to 1.
+        prob_min (float, optional): The minimum probability threshold for predictions.
+                                    Defaults to 0.01.
+        num_workers (int, optional): Number of CPU for multiprocessing in Dataloader.
+                                     Defaults to 1.
         batch_size (int, optional): Size of a batch for batch prediction.
 
-    For single predictions, we recommend keeping num_workers and batch_size to 1 for better performance.
-    For batched predictions, consider increasing these two parameters (num_workers can range from 4 to 12, batch size can be increased up to 256) to optimize performance.
+    For single predictions, we recommend keeping num_workers and batch_size to 1
+        for better performance.
+    For batched predictions, consider increasing these two parameters
+        (num_workers can range from 4 to 12, batch size can be increased up to 256)
+        to optimize performance.
 
     Returns:
         list: The list of predicted responses.
@@ -51,4 +56,7 @@ async def predict(
     }
 
     output = request.app.state.model.predict(input_data, params=params_dict)
-    return [out.model_dump() for out in output]
+    return [
+        OutputResponse({**out.model_dump(), "MLversion": request.app.state.run_id})
+        for out in output
+    ]

src/utils/logging.py

@@ -1,6 +1,6 @@
 import logging
 
-from api.models.responses import PredictionResponse
+from api.models.responses import OutputResponse
 
 
 def configure_logging():
@@ -13,6 +13,6 @@ def configure_logging():
     )
 
 
-def log_prediction(query: dict, response: PredictionResponse, index: int = 0):
+def log_prediction(query: dict, response: OutputResponse, index: int = 0):
     query_line = {key: value[index] for key, value in query.items()}
     logging.info(f"{{'Query': {query_line}, 'Response': {response.model_dump()}}}")