added illustration

app.py

@@ -1,15 +1,12 @@
 import gradio as gr
 
+from queue import Queue
 from threading import Thread
-from queue import Queue, Empty
-# from callbacks import StreamingGradioCallbackHandler, job_done
+from callbacks import StreamingGradioCallbackHandler, job_done
 
-from langchain.schema import SystemMessage
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationChain
-from langchain.prompts import ChatPromptTemplate
 from langchain.memory import ConversationBufferMemory
-from langchain.callbacks.base import BaseCallbackHandler
 
 # huggingface.co/spaces/huggingface-projects/llama-2-13b-chat
 DEFAULT_SYSTEM_PROMPT = """\
@@ -19,74 +16,66 @@ ensure that your responses are socially unbiased and positive in nature.\n\nIf a
 is not factually coherent, explain why instead of answering something not correct. If you don't know the answer \
 to a question, please don't share false information."""
 
-class QueueCallback(BaseCallbackHandler):
-    """Callback handler for streaming LLM responses to a queue."""
-
-    def __init__(self, q):
-        self.q = q
-
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        print(token)
-        self.q.put(token)
-
-    def on_llm_end(self, *args, **kwargs) -> None:
-        print("Done")
-        return self.q.empty()
-
 def respond(openai_api_key, openai_model, creativity, max_tokens, message, buffer_memory, chat_history):
-    # print(buffer_memory.buffer)
-    chat_history.append([message, None])
     q = Queue()
-    job_done = object()
-    callback = QueueCallback(q)
     conversation = ConversationChain(
         llm = ChatOpenAI(
+            streaming=True,
             model=openai_model,
             max_tokens=max_tokens,
             temperature=creativity,
             openai_api_key=openai_api_key,
-            streaming=True,
-            callbacks=[callback]
+            callbacks=[StreamingGradioCallbackHandler(q)]
         ), 
         memory = buffer_memory
     )
-    def task():
-        resp = conversation.predict(input=message)
-        q.put(job_done)
-    thread = Thread(target=task)
+    chat_history.append([message, ""])
+
+    thread = Thread(target=conversation.predict, kwargs={
+        "input": message,
+    })
     thread.start()
-    chat_history[-1] = (chat_history[-1][0], "")
+
     while True:
         next_token = q.get(block=True) # Blocks until an input is available
         if next_token is job_done:
             break
         chat_history[-1] = (chat_history[-1][0], chat_history[-1][1] + next_token)
-        yield "", buffer_memory, chat_history  # Yield the chatbot's response as a string
+        yield "", buffer_memory, chat_history  # Yield the chatbot's response
     thread.join()
 
 
+def init_buffer_memory():
+    memory = ConversationBufferMemory()
+    memory.save_context({"input": DEFAULT_SYSTEM_PROMPT}, {"output": "Of course!"})
+    return memory
+
+
 with gr.Blocks(css="#component-0 { max-width: 900px; margin: auto; padding-top: 1.5rem; }") as demo:
-    # with gr.Column():
-    with gr.Row(visible=True) as primary_settings:
-        openai_key = gr.Textbox(
-            label="OpenAI Key",
-            type="password",
-            placeholder="sk-a83jv6fn3x8ndm78b5W...",
-        )
-        model = gr.Dropdown(
-            ["gpt-4", 
-             "gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-instruct", 
-             "text-davinci-002", "text-davinci-003"], 
-            label="OpenAI Model",
-            value="gpt-3.5-turbo",
-            interactive=True
-        )
-    # with gr.Accordion("Advances Settings"):
-    #     gr.Dropdown(
-    #         [-1, 1, 5, 10, 25], label="Conversation Buffer (k)"
-    #     )
+    gr.Markdown(
+        """
+        ![](/home/abdulla/gradio/banner.png)
+        """
+    )
+    
+    with gr.Group():
+        with gr.Row(visible=True) as primary_settings:
+            openai_key = gr.Textbox(
+                container=False,
+                type="password",
+                placeholder="OpenAI Key: sk-a83jv6fn3x8ndm78b5W...",
+            )
+            model = gr.Dropdown(
+                ["gpt-4", 
+                "gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-instruct", 
+                "text-davinci-002", "text-davinci-003"],
+                container=False,
+                value="gpt-3.5-turbo",
+                interactive=True
+            )
+        
     with gr.Group() as chat:
-        memory = gr.State(ConversationBufferMemory())
+        memory = gr.State(init_buffer_memory())
         chatbot = gr.Chatbot(label='Chatbot')
         with gr.Row():
             query = gr.Textbox(
@@ -95,10 +84,16 @@ with gr.Blocks(css="#component-0 { max-width: 900px; margin: auto; padding-top:
                 placeholder='Type a message...',
                 scale=10,
             )
-            submit = gr.Button('Submit',
-                        variant='primary',
-                        scale=1,
-                        min_width=0)
+            submit = gr.Button(
+                'Submit',
+                variant='primary',
+                scale=1,
+                min_width=0
+            )
+
+    with gr.Row():
+        regenerate = gr.Button("Regenerate")
+        clear_history = gr.Button("Clear History")
 
     with gr.Accordion(label='Advanced options', open=False):
         system_prompt = gr.Textbox(label='System prompt', value=DEFAULT_SYSTEM_PROMPT, lines=6)
@@ -128,5 +123,7 @@ with gr.Blocks(css="#component-0 { max-width: 900px; margin: auto; padding-top:
     # Event Handling
     query.submit(respond, [openai_key, model, temperature, max_new_tokens, query, memory, chatbot], [query, memory, chatbot])
     submit.click(respond, [openai_key, model, temperature, max_new_tokens, query, memory, chatbot], [query, memory, chatbot])
+
+    regenerate.click()
     
 demo.queue().launch()