gradio_chat.py

import json 
import os 
import openai
import tiktoken
import gradio as gr

openai.api_key = os.getenv("OPENAI_API_KEY")

USE_GPT4 = False
OPENAI_MODEL = "gpt-4" if USE_GPT4 else "gpt-3.5-turbo"
CONTEXT_SIZE = 8191 if USE_GPT4 else 4096 # note: -1 because API (sometimes?) errors if we use max tokens

def count_tokens(text):
    encoding = tiktoken.encoding_for_model(OPENAI_MODEL)
    return len(encoding.encode(text))

def calc_max_tokens(messages):
    return CONTEXT_SIZE - len(json.dumps(messages)) # todo improve efficiency

messages = [ {"role": "system", "content": "You are a helpful assistant."} ]

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        history[-1][1] = ""
        user_msg = history[-1][0]
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": ""})

        for resp in openai.ChatCompletion.create(
            model=OPENAI_MODEL,
            messages=messages,
            temperature=0.7,
            max_tokens=calc_max_tokens(messages),
            stream=True
        ):
            delta = resp.choices[0].delta
            if 'role' in delta: 
                continue
            if 'content' in delta:
                token = delta.content
                history[-1][1] += token
                messages[-1]["content"] += token
                yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch()