from pydantic import BaseModel,Json
from fastapi import FastAPI,Request,Response
from fastapi.middleware.cors import CORSMiddleware
import requests


app = FastAPI()
origins = ["*"]

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

class Item(BaseModel):
    transcript: Json

@app.get("/")
def base():
    return """PROXY ONLINE on /llm route $curl -X POST https://iakashpaul-cors-proxy-baseten.hf.space/llm --data '{"prompt":"hello"}'"""

prefix_prompt="""<s>[INST]You will be given a user query for interacting  with record/lines/rows of a db.\nYou need to guess whether it is only among 'create'(eg. generate, write row, create record number 45), 'update'(eg. revise, update, modify), 'read'(eg. see, show me, view, read) or 'delete'(eg. remove record, delete row number 5) logically & whether the user has mentioned numbers or digits of any rows/record/lines between 1 to 100. \nYour response should strictly be as follows & no extra words/explanations or respond with just a single output. \nIn case no action matches, just respond with {"NA":"NA"}, don't assume anything just reply to the user in-case it is not clear what the action it is they want to perform. Do not add comments after the JSON \nFor given example input: Show me record number 3.\noutput:{"action":"read","number":"3"}[/INST]\n"""
suffix_prompt="""\n"""
import os
baseten_url = os.getenv("BASETEN_URL")
baseten_key = os.getenv("BASETEN_KEY")
@app.post("/llm")
async def main(request: Request):
    input_json =  await request.json()
    print(input_json)
    final_prompt = prefix_prompt + str(input_json["prompt"]) + suffix_prompt
    resp = requests.post(
    baseten_url,
    headers={"Authorization": f"Api-Key {baseten_key}"},
    json={'prompt': final_prompt ,'temperature': 0.001, 'max_new_tokens': 100, 'repetition_penalty':1.2},
)   
    llm_response = resp.json()
    llm_response = llm_response.rsplit("[/INST]")[-1].split("</s>")[0];
    print(llm_response)
    return {"text":str(llm_response)}