import joblib
from typing import Union
from fastapi import FastAPI
from pydantic import BaseModel

class Item(BaseModel):
    value: str

app = FastAPI()

@app.get("/")
async def read_main():
    return {"msg": "Welcome in simple web application that provides prediction whether the sent data is an attack attempt or not. Use /detect/ endpoints with GET or POST methods - POST endpoint awaits for simple json with value assigned to 'value' key :) "}

@app.get("/detect/{value}")
async def detect_input_get(value: Union[str]):
    nb_saved = joblib.load("nb.joblib")
    vec_saved = joblib.load("vec.joblib")

    return {"prediction": str(nb_saved.predict(vec_saved.transform([value])))}

@app.post("/detect/")
async def detect_input_post(item: Item):
    nb_saved = joblib.load("nb.joblib")
    vec_saved = joblib.load("vec.joblib")

    return {"prediction": str(nb_saved.predict(vec_saved.transform([item.value])))}