from fastapi import FastAPI

import os

from langchain.document_loaders import DirectoryLoader

import wandb

import huggingface_hub

from scripts.summarization import falcon_summ

from scripts.decision_clf import seq_clf

from scripts.path_gen import paths_gen

from scripts.text_gen import story_gen

app = FastAPI()

wandb.login(key = os.getenv('wandb_key'))

huggingface_hub.login(token = os.getenv('hf_key'))

os.environ['OPENAI_API_KEY'] = os.getenv('openapi_key')

summarizer = falcon_summ.prep_pipeline()

token_path= 'models/dec_clf/tokenizer.pkl'

model_path = 'models/dec_clf/nlp.h5'

@app.get("/hello")
def hello():
    return {"message": "Hello World"}

@app.post("/summ")
def summ(text: str):
    return { "summary": falcon_summ.gen_summary(summarizer, text)}

@app.post("/clf")
def clf(text: str):
    return {"decision": seq_clf.predict(text, model_path, token_path)}

@app.post("/gen_path")
def gen_path(text: str, name: str):
    db = paths_gen.get_vectordb([], name)

    return paths_gen.gen_sample(text, db)

@app.post("/gen_story")
def gen_story(text: str, name: str, decision: str):
    db = paths_gen.get_vectordb([], name)

    return story_gen.gen_sample(text, decision, db)