### Importing Libraries

import os;
import sys;
import time;
import openai;
import tiktoken;

from PIL import Image;

import gradio as gr; # type: ignore
from functools import partial;

### OpenAI Setup

azure_api_key = os.environ['AZURE_OPENAI_KEY']

### OpenAI LLM Get Completion Function Example - Summarization of User Message

# client = openai.OpenAI();
client = openai.AzureOpenAI(
    azure_endpoint='https://wadiopenai.openai.azure.com/',
    api_key=azure_api_key,
    api_version="2024-02-01"
)
delimiter = '####';

sample_user_message = f"""
The Empire State Building is a steel-framed skyscraper rising 102 stories that was completed in New York City in 1931 and was the tallest building in the world until 1971. 
The Empire State Building is located in Midtown Manhattan, on Fifth Avenue at 34th Street. 
It remains one of the most distinctive and famous buildings in the United States and is one of the best examples of Modernist Art Deco design. 
At the time of its construction, there was fierce competition to win the title of the tallest building in the world. 
The Chrysler Building claimed the title in 1929, and the Empire State Building seized it in 1931, its height being 1,250 feet (381 metres) courtesy of its iconic spire, which was originally intended to serve as a mooring station for airships. 
A 222-foot (68-metre) antenna was added in 1950, increasing the building’s total height to 1,472 feet (449 metres), but the height was reduced to 1,454 feet (443 metres) in 1985 when the antenna was replaced. 
By that time the One World Trade Center, officially opened in 1972, had become the tallest building in the world.
""";

sample_riddle = f"""
You are a bus driver. The bus is initially empty. 
At the 1st stop of the day, 8 people get on board. 
At the 2nd stop, 4 people get off and 11 people get on. 
At the 3rd stop, 2 people get off and 6 people get on. 
At the 4th stop 13 people get off and 1 person gets on. 
How many people are now on the bus?
""";

def summarize(user_message, model='gpt-35', temperature = 0.1, max_tokens = 1028):

    response = client.chat.completions.create(
        model = model,
        messages = [
            {
                'role': 'system',
                'content': f"""
                You are a helpful assistant.
                Summarize the following text delimited by {delimiter}
                """
            }, 
            {
                'role': 'user',
                'content': f"""{delimiter} {user_message} {delimiter}"""
            }
        ],
        temperature = temperature,
        max_tokens = max_tokens
    );
    
    return response.choices[0].message.content;

### OpenAI LLM Get Completion Function Example for Multiple LLMs

system_prompt_2 = "You are a helpful assistant.";

def gpt_3p5t_completion(user_message, model = 'gpt-35', temperature = 0.1, max_tokens = 1028):

    response = client.chat.completions.create(
        model = model,
        messages = [
            {
                'role': 'system',
                'content': system_prompt_2
            }, 
            {
                'role': 'user',
                'content': f"""{delimiter} {user_message} {delimiter}"""
            }
        ],
        temperature = temperature,
        max_tokens = max_tokens
    );
    
    return response.choices[0].message.content;

def gpt_4t_completion(user_message, model = 'gpt-4o-mini', temperature = 0.0, max_tokens = 1028):

    response = client.chat.completions.create(
        model = model,
        messages = [
            {
                'role': 'system',
                'content': system_prompt_2
            }, 
            {
                'role': 'user',
                'content': f"""{delimiter} {user_message} {delimiter}"""
            }
        ],
        temperature = temperature,
        max_tokens = max_tokens
    );

    return response.choices[0].message.content;

def gpt_4om_completion(user_message, model = 'gpt-4o-mini', temperature = 0.1, max_tokens = 1028):

    response = client.chat.completions.create(
        model = model,
        messages = [
            {
                'role': 'system',
                'content': system_prompt_2
            }, 
            {
                'role': 'user',
                'content': f"""{delimiter} {user_message} {delimiter}"""
            }
        ],
        temperature = temperature,
        max_tokens = max_tokens
    );

    return response.choices[0].message.content;

### Gradio Interface First Demo - Only Structure & No Events

"""
gr.close_all();
demo = gr.Interface(
    fn = summarize,
    inputs = [gr.Textbox(label='Text to summarize', lines = 6)],
    outputs = [gr.Textbox(label = 'Result', lines = 3)],
    title = 'Text Summarization with GPT-3.5 Turbo',
    description = 'Summarize any text with OpenAI\'s GPT-3.5 Turbo'
);

demo.launch(share = True);
"""

### Gradio Blocks Demo - Application

gr.close_all();

def hide_textbox():
    return gr.Textbox(visible=False);

def inference(input):
        return gpt_3p5t_completion(input), gpt_4t_completion(input), gpt_4om_completion(input);

with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
    
    gr.Markdown('<center><h1> OpenAI LLM Explorer </h1><center>');
    gr.Markdown("Type your prompt below and compare results from 3 important OpenAI models.");
    prompt = gr.Textbox(label='Prompt', lines=3, max_lines=5);
    token = gr.Textbox(label='OpenAI Token', type='password', placeholder='Your OpenAI Token');

    generate_btn = gr.Button(variant='primary', size='sm'); 

    with gr.Row() as label_row:
        gpt_35_label = gr.Markdown('<b><h4> GPT-3.5-Turbo </h4></b>');
        gpt_4t_label = gr.Markdown('<b><h4> GPT-4-Turbo </h4></b>');
        gpt_4om_label = gr.Markdown('<b><h4> GPT-4o-mini </h4></b>');
    
    with gr.Row() as output_row:
        gpt_35_output = gr.Markdown('Response to show here');
        gpt_4t_output = gr.Markdown('Response to show here');
        gpt_4om_output = gr.Markdown('Response to show here');

    gr.on(
        triggers = [prompt.submit, generate_btn.click],
        fn = hide_textbox,
        inputs = None,
        outputs = [token]
    ).then(
        fn = inference,
        inputs = [prompt],
        outputs = [gpt_35_output, gpt_4t_output, gpt_4om_output]
    );

demo.launch(share = True);