from transformers import pipeline
from fastapi import FastAPI, Request, HTTPException
from payload import SomeText
from api_contract import components_dict as api_components
from error_handling import ErrorCodes

app = FastAPI()

app=FastAPI(title="Huggingface Gen LLM gest",
            version="1.0",
            debug=True,
            components=api_components,
            swagger_ui_bundle_js= "//unpkg.com/swagger-ui-dist@3/swagger-ui-bundle.js",
            swagger_ui_standalone_preset_js= "//unpkg.com/swagger-ui-dist@3/swagger-ui-standalone-preset.js",
            summary="API to perform generative prompt completion using small LLM (without GPU).",
)

pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")

@app.get("/")
def read_root():
    return {'detail': 'API running. Try out the endpoints in swagger'}

@app.get("/api/check-heartbeat")
def get_heartbeat():
    return {"detail": "seems to be working"}

@app.post("/api/generatel-language", summary="Generate text from prompt", tags=["Generate"])
def inference(request: Request, input_prompt: SomeText):

    if len(input_prompt.text) >0:

        output = pipe_flan(input_prompt.text)
        return {"output": output[0]["generated_text"]}

    else:
        raise HTTPException(status_code=400, detail = ErrorCodes.REQUEST_VALIDATION_ERROR.value[1])