from fastai import FastApi
from fastapi.middleware.cors import CORSMiddleware
from generator import generation_function
from pydantic import BaseModel


app = FastApi()
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class Prompt(BaseModel):
    prompts: str
    lang: str


@app.post('/generate')
def generate(prompt: Prompt):
    request_data = prompt.dict()

    text = request_data["prompts"]

    if request_data.get("lang"):
        generated_recepie = generation_function(
            text, lang=request_data["lang"])
    else:
        generated_recepie = generation_function(text)

    return generated_recepie