app.py

# docker build -t ml-endpoint .
# docker run --env-file .env -p 8000:8000 ml-endpoint

from fastapi import FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import os
import pandas as pd
import numpy as np
import tensorflow as tf
import re
from openai import OpenAI
import time
from pymongo import MongoClient
import tempfile
import bson
from datetime import datetime, timezone
from bson import ObjectId
from mangum import Mangum

ML_API_SECRET = os.environ.get("ML_API_SECRET")
# Connect to MongoDB and instantiate models collection
client = MongoClient(os.environ.get("MONGO_CONNECTION_STRING"))
db = client["spending-tracker-db"]
GLOBAL_MODEL_RETRIEVAL_VERBOSE=True
USER_MODEL_RETRIEVAL_VERBOSE=True
global_model_db = db["global_model"]
users_db = db["users"]
user_models_db = db["user_models"]

SAVE_NAME = "model_"
# SAVE_FOLDER_PATH = "./saved_models/mlp/text_embedding_3_small"
MODEL_EXTENSION = ".h5"

RULES_PATH = './data/rules.csv'
CASE_SENSITIVE_RULES_PATH = './data/case-sensitive-rules.csv'
CATEGORIES = ["Groceries", "Housing & Bills", "Finance & Fees", "Transport", "Income", "Shopping",
               "Eating Out", "Entertainment", "Health & Fitness", "Transfer", "Other / Misc"]
NUM_LABELS = len(CATEGORIES)
idx_to_labels = { k : v for k, v in enumerate(CATEGORIES)}
labels_to_idx = { k : v for v, k in enumerate(CATEGORIES)}

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIMS = 1536
EMBEDDING_VERBOSE=True


def save_global_model(model, collection, metadata={}):

    with tempfile.NamedTemporaryFile(suffix=MODEL_EXTENSION, delete=False) as tmp:
        tmp_filepath = tmp.name

    try:
        model.save(tmp_filepath)
        # Read back as bytes
        with open(tmp_filepath, "rb") as f:
            model_bytes = f.read()

        # Insert into MongoDB
        collection.insert_one({
            "modelBlob": bson.Binary(model_bytes),
            "createdAt": datetime.now(timezone.utc),
            "metadata": metadata,
            "stable": False,
        })

    finally:
        os.remove(tmp_filepath)

def get_global_model(collection, verbose=GLOBAL_MODEL_RETRIEVAL_VERBOSE):
    start = time.time()
    # Find latest model version
    doc = collection.find_one(
        { "stable": True }, 
        sort=[("createdAt", -1)]
    )
    model_bytes = doc["modelBlob"]

    # Write bytes to a temporary file
    with tempfile.NamedTemporaryFile(suffix=MODEL_EXTENSION, delete=False) as tmp:
        tmp.write(model_bytes)
        tmp_filepath = tmp.name

    try:
        loaded_model = tf.keras.models.load_model(tmp_filepath)
        return loaded_model
    
    finally:
        elapsed = time.time() - start
        if verbose:
            print(f"Time taken to retrieve model from global_model: {elapsed:2f} seconds")
        os.remove(tmp_filepath)

def save_user_model(model, collection, uid, is_auth_0_user):

    with tempfile.NamedTemporaryFile(suffix=MODEL_EXTENSION, delete=False) as tmp:
        tmp_filepath = tmp.name

    try:
        model.save(tmp_filepath)
        # Read back as bytes
        with open(tmp_filepath, "rb") as f:
            model_bytes = f.read()

        if not uid:
            raise ValueError("uid must be provided")
        
        if is_auth_0_user:
            collection.update_one(
                { "auth0Uid": uid },
                { 
                    "$set": { "modelBlob": bson.Binary(model_bytes) },
                    "$setOnInsert": { "auth0Uid": uid, "uid": None },
                },
                upsert=True,
            )

        else:
            collection.update_one(
                { "uid": ObjectId(uid) },
                { 
                    "$set": {"modelBlob": bson.Binary(model_bytes)},
                    "$setOnInsert": { "uid": ObjectId(uid), "auth0Uid": None },
                },
                upsert=True,
            )

    finally:
        os.remove(tmp_filepath)

def get_user_model(collection, uid, is_auth0_user, verbose=USER_MODEL_RETRIEVAL_VERBOSE):
    start = time.time()

    if not uid:
        raise ValueError(f"uid must be provided")
    
    doc = None
    if is_auth0_user:
        doc = collection.find_one({ "auth0Uid": uid })
    else:
        doc = collection.find_one({ "uid": ObjectId(uid) })
    
    if not doc:
        return None
    
    model_bytes = doc["modelBlob"]

    if not model_bytes:
        print("Model not found for user")
        return None

    # Write bytes to a temporary file
    with tempfile.NamedTemporaryFile(suffix=MODEL_EXTENSION, delete=False) as tmp:
        tmp.write(model_bytes)
        tmp_filepath = tmp.name

    try:
        loaded_model = tf.keras.models.load_model(tmp_filepath)
        return loaded_model
    
    finally:
        elapsed = time.time() - start
        if verbose:
            print(f"Time taken to retrieve model from users: {elapsed:2f} seconds")
        os.remove(tmp_filepath)


def create_embeddings(sentences, open_ai_model=EMBEDDING_MODEL, verbose=EMBEDDING_VERBOSE):
    start_time = time.time()

    response = client.embeddings.create(
        model=open_ai_model,
        input=sentences
    )

    end_time = time.time()
    elapsed = end_time - start_time

    if verbose:
        print(f"Time to embed {len(sentences)} sentences: {elapsed:.2f} seconds")
    
    # Extract the embedding vectors into a numpy array
    embeddings = np.array([object.embedding for object in response.data], dtype=np.float32)
    return embeddings


def format_description(description, lower=True):
        formatted_description = str(description).split('\t')[0].strip()
        if lower:
              formatted_description.lower()
        return re.sub(r"\s+", " ", formatted_description).strip()


def save_model(model, folder_path, EXTENSION=MODEL_EXTENSION, SAVE_NAME=SAVE_NAME):
    files = os.listdir(folder_path)
    files = [os.path.splitext(file)[0] for file in files if not file.startswith('.')]
    filename = SAVE_NAME
    
    num = 0
    if len(files) > 0:
        for file in files:
            num = max(num, int(file.split('_')[-1]) + 1)

    filename += str(num)
    if EXTENSION: filename += EXTENSION

    folderpath = os.path.join(folder_path, filename)
    # Convert to forward slashes
    folderpath = folderpath.replace("\\", "/")

    model.save(folderpath)

    return folderpath


def get_latest_model(folder_path, EXTENSION=MODEL_EXTENSION, SAVE_NAME=SAVE_NAME):
    files = os.listdir(folder_path)
    files = [os.path.splitext(file)[0] for file in files if not file.startswith('.')]
    filename = SAVE_NAME

    num = 0
    if len(files) > 0:
        for file in files:
            num = max(num, int(file.split('_')[-1]))

    filename += str(num)
    if EXTENSION: filename += EXTENSION
    folderpath = os.path.join(folder_path, filename)
    
    # Convert to forward slashes
    folderpath = folderpath.replace("\\", "/")

    return folderpath


# Load model on start
MLP_MODEL = get_global_model(global_model_db)

def predict(model, embeddings):
    predictions = model.predict(embeddings)
    maxProbs = predictions.max(-1)

    predicted_indices = np.argmax(predictions, axis=-1)
    predicted_categories = [idx_to_labels[idx] for idx in predicted_indices]

    return predicted_categories, maxProbs


# Load rules
rules = pd.read_csv(RULES_PATH)
case_sensitive_rules = pd.read_csv(CASE_SENSITIVE_RULES_PATH)

# Predict request schema
class PredictItems(BaseModel):
    descriptions: List[str]
    modelType: str
    isAuth0User: bool
    uid: str

# Predict request schema
class PredictRequest(BaseModel):
    predict_data: PredictItems

# Response schema
class Prediction(BaseModel):
    predicted_category: str
    confidence: float

# Train request schema
class TrainItems(BaseModel):
    descriptions: List[str]
    categories: List[str]
    modelType: str
    isAuth0User: bool
    uid: str

# Train request schema
class TrainRequest(BaseModel):
    train_data: TrainItems

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173"],  # frontend URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/predict", response_model=List[Prediction])
def predict_route(request: PredictRequest, ml_api_secret: Optional[str] = Header(None, alias="ml_api_secret")):
    if ml_api_secret != ML_API_SECRET:
        raise HTTPException(status_code=401, detail="Unauthorized")
    
    descriptions = request.predict_data.descriptions
    model_type = request.predict_data.modelType
    is_auth0_user = request.predict_data.isAuth0User
    uid = request.predict_data.uid

    # Create embeddings on formatted lowercased descriptions
    embeddings = create_embeddings([format_description(description) for description in descriptions])

    model = MLP_MODEL
    if model_type == "clientModel":
        retrieved_model = get_user_model(user_models_db, uid, is_auth0_user)

        if not retrieved_model:
            print("failed to retrieve client model")
            save_user_model(MLP_MODEL, user_models_db, uid, is_auth0_user)

        else:
            print("client model successfully retrieved")
            model = retrieved_model

    # Predict
    preds, probabilities = predict(model, embeddings)

    predicted_categories = []
    for i, d in enumerate(descriptions):
        matches = case_sensitive_rules[
            case_sensitive_rules["company_name"].apply(lambda rule: rule in d)
        ]
        if not matches.empty:
            predicted_categories.append(matches.iloc[0]["category"])
            probabilities[i] = 1.0
            continue

        matches = rules[rules["company_name"].str.lower().apply(lambda rule: rule in d.lower())]
        
        if not matches.empty:
            predicted_categories.append(matches.iloc[0]["category"])
            probabilities[i] = 1.0
            continue

        predicted_categories.append(preds[i])
        
    results = [
        Prediction(predicted_category=p, confidence=c)
        for p, c in zip(predicted_categories, probabilities)
    ]

    return results


@app.post("/train")
def train_route(request: TrainRequest, ml_api_secret: Optional[str] = Header(None, alias="ml_api_secret")):
    if ml_api_secret != ML_API_SECRET:
        raise HTTPException(status_code=401, detail="Unauthorized")

    uid = request.train_data.uid
    model_type = request.train_data.modelType
    is_auth0_user = request.train_data.isAuth0User

    global_model = MLP_MODEL
    client_model = None

    if model_type == "clientModel":
        retrieved_model = get_user_model(user_models_db, uid, is_auth0_user)

        if not retrieved_model:
            print("failed to retrieve client model")

        else:
            print("client model successfully retrieved")
            client_model = retrieved_model

    global_model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    if client_model:
        client_model.compile(
            optimizer='adam',
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy']
        )

    embeddings = create_embeddings(request.train_data.descriptions)
    categories = request.train_data.categories

    labels = np.array([labels_to_idx[label] for label in categories])

    global_model.fit(
        embeddings, 
        labels,
        epochs=1,
        batch_size=16,
    )
    
    if client_model:
        client_model.fit(
            embeddings, 
            labels,
            epochs=1,
            batch_size=16,
        )

    if client_model:
        save_user_model(client_model, user_models_db, uid, is_auth0_user)
        save_global_model(global_model, global_model_db)
    else:
        save_user_model(global_model, user_models_db, uid, is_auth0_user)
        save_global_model(global_model, global_model_db)

    return { "status": "model updated" }

# Lambda handler
handler = Mangum(app)