Add mongo for api calls

Author: elliotdes

api/main.py

@@ -1,11 +1,12 @@
-import logging
+import datetime
 from typing import List
 import os
 
 import pandas as pd
 from fastapi import FastAPI
 import mlflow
 import numpy as np
+import pymongo
 
 from ml.data import get_data
 from ml.evaluate import eval_metrics
@@ -16,6 +17,13 @@
 app = FastAPI()
 
 
+def response_to_mongo(r: dict):
+    client = pymongo.MongoClient("mongodb://mongo:27017")
+    db = client["models"]
+    model_collection = db["example-model"]
+    model_collection.insert_one(r)
+
+
 @app.on_event("startup")
 def load_api_model():
     model_id = os.environ.get("MODEL_ID")
@@ -41,6 +49,9 @@ async def predict_model(features: List[float]):
     )
 
     response = {"predictions": prediction.tolist()}
+    response_to_mongo(
+        {"predictions": prediction.tolist(), "date": datetime.datetime.utcnow()},
+    )
     return response
 
 

docker-compose.yml

@@ -1,11 +1,12 @@
 version: "3.9"
 
 services:
-  postgres:
-    image: "postgres"
-    restart: always
-    environment:
-      POSTGRES_PASSWORD: password
+  mongo:
+    image: mongo:latest
+    ports:
+      - 27017:27017
+    volumes:
+      - mongodb_data_container:/data/db
 
   api:
     build: .
@@ -34,3 +35,6 @@ services:
       - "5001:5001"
     volumes:
       - ./mlruns:/code/mlruns
+
+volumes:
+  mongodb_data_container:

poetry.lock

@@ -12,6 +12,7 @@ typer = "^0.4.0"
 fastapi = "^0.75.0"
 uvicorn = "^0.17.5"
 streamlit = "^1.7.0"
+pymongo = "^4.0.2"
 
 [tool.poetry.dev-dependencies]
 black = "^22.1.0"