from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from langchain.llms import OpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import openai
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()

# Configure OpenAI for Azure
openai.api_type = "azure"
openai.api_base = "https://amplifai-openai.openai.azure.com/"
openai.api_version = "2023-07-01-preview"
openai.api_key = os.getenv("OPENAI_API_KEY")

app = FastAPI()

class Prompt(BaseModel):
    text: str

class GenerateResponse(BaseModel):
    text: str

def generate_text(prompt: str):
    if prompt == "":
        return {"detail": "Please provide a prompt."}
    else:
        prompt_template = PromptTemplate(template=prompt, input_variables=['Prompt'])
        llm = OpenAI(api_key=openai.api_key)
        llmchain = LLMChain(
            prompt=prompt_template,
            llm=llm
        )
        llm_response = llmchain.run(
            {"Prompt": prompt},
            temperature=0.7,
            max_tokens=750,
            top_p=1.0,
            frequency_penalty=0.0,
            presence_penalty=0.0
        )
        return GenerateResponse(text=llm_response)

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

@app.get("/", tags=["Home"])
def api_home():
    return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}

@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=GenerateResponse)
def inference(input_prompt: Prompt):
    return generate_text(prompt=input_prompt.text)