update model serving api with extra checks

Author: divakaivan

src/make_api/app/main.py

@@ -1,15 +1,23 @@
 import os
+from functools import lru_cache
 
 import mlflow
 import pandas as pd
 import uvicorn
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
 mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
 
 model_uri = "models:/taxi_fare_prediction.taxi_fare_model@latest-model"
-pipe = mlflow.sklearn.load_model(model_uri)
+
+
+@lru_cache
+def get_model():
+    return mlflow.sklearn.load_model(model_uri)
+
+
+pipe = get_model()
 
 app = FastAPI()
 
@@ -26,16 +34,17 @@ class InferenceInput(BaseModel):
 
 @app.get("/")
 async def read_root():
-    return {"Hello": "Visit /docs for API documentation"}
+    return {"message": "Welcome to the Taxi Fare Prediction API!", "docs": "/docs"}
 
 
 @app.post("/predict")
 async def predict_one(data: InferenceInput) -> dict[str, float]:
-    df_input = pd.DataFrame([data.dict()])
-
-    prediction = pipe.predict(df_input)
-
-    return {"prediction": prediction[0]}
+    try:
+        df_input = pd.DataFrame([data.dict()])  # pd df might be overkill
+        prediction = pipe.predict(df_input)
+        return {"prediction": prediction[0]}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Prediction failed: {str(e)}") from e
 
 
 if __name__ == "__main__":