use gr chatbot

app.py

@@ -11,8 +11,8 @@ MODELS = [
     "Meta-Llama-3.1-8B-Instruct"
 ]
 
-def create_client(api_key = None):
-    if api_key is not None:
+def create_client(api_key=None):
+    if api_key:
         openai.api_key = api_key
         openai.api_base = "https://api.sambanova.ai/v1"  # Fixed Base URL
     else:
@@ -48,11 +48,11 @@ def respond(message, chat_history, model, system_prompt, thinking_budget, api_ke
         response = completion.choices[0].message['content']
         thinking_time = time.time() - start_time
         print("Response received from OpenAI API.")
-        yield response, thinking_time
+        return response, thinking_time
     except Exception as e:
         error_message = f"Error: {str(e)}"
         print(error_message)
-        yield error_message, time.time() - start_time
+        return error_message, time.time() - start_time
 
 def parse_response(response):
     answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
@@ -67,29 +67,23 @@ def parse_response(response):
 
 def process_chat(message, history, model, system_prompt, thinking_budget, api_key):
     print(f"Received message: {message}")
-    #if not api_key:
-    #    print("API key missing")
-    #    return "Please provide your API Key before starting the chat."
+    if not api_key:
+        print("API key missing")
+        return history + [("System", "Please provide your API Key before starting the chat.")]
 
     try:
         formatted_system_prompt = system_prompt.format(budget=thinking_budget)
     except KeyError as e:
         error_msg = f"System prompt missing placeholder: {str(e)}"
         print(error_msg)
-        return error_msg
+        return history + [("System", error_msg)]
 
-    full_response = ""
-    thinking_time = 0
+    response, thinking_time = respond(message, history, model, formatted_system_prompt, thinking_budget, api_key)
 
-    for response, elapsed_time in respond(message, history, model, formatted_system_prompt, thinking_budget, api_key):
-        print(f"Received response: {response}")
-        full_response = response
-        thinking_time = elapsed_time
+    if response.startswith("Error:"):
+        return history + [("System", response)]
 
-    if full_response.startswith("Error:"):
-        return full_response
-
-    answer, reflection, steps = parse_response(full_response)
+    answer, reflection, steps = parse_response(response)
 
     formatted_response = f"**Answer:** {answer}\n\n**Reflection:** {reflection}\n\n**Thinking Steps:**\n"
     for i, step in enumerate(steps, 1):
@@ -98,8 +92,7 @@ def process_chat(message, history, model, system_prompt, thinking_budget, api_ke
     formatted_response += f"\n**Thinking time:** {thinking_time:.2f} s"
 
     print(f"Appended response: {formatted_response}")
-    history.append((message, formatted_response))
-    return formatted_response
+    return history + [(message, formatted_response)]
 
 # Define the default system prompt
 default_system_prompt = """
@@ -169,42 +162,42 @@ with gr.Blocks() as demo:
     system_prompt = gr.Textbox(
         label="System Prompt",
         value=default_system_prompt,
-        lines=10
+        lines=15,
+        interactive=True
     )
 
-    msg = gr.Textbox(
-        label="Type your message here...",
-        placeholder="Enter your message..."
-    )
-    submit = gr.Button("Submit")
-    clear = gr.Button("Clear Chat")
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Type your message here...",
+            placeholder="Enter your message..."
+        )
+        submit = gr.Button("Submit")
+        clear = gr.Button("Clear Chat")
 
-    output = gr.Textbox(
-        label="Response",
-        lines=20,
-        interactive=False
+    chatbot = gr.Chatbot(
+        label="Chat History"
     )
 
-    # Initialize chat history
-    chat_history = []
+    # Initialize chat history as a Gradio state
+    chat_history = gr.State([])
 
     def handle_submit(message, history, model, system_prompt, thinking_budget, api_key):
-        response = process_chat(message, history, model, system_prompt, thinking_budget, api_key)
-        return response
+        updated_history = process_chat(message, history, model, system_prompt, thinking_budget, api_key)
+        return updated_history, ""
 
     def handle_clear():
-        return ""
+        return [], ""
 
     submit.click(
         handle_submit,
-        inputs=[msg, gr.State(chat_history), model, system_prompt, thinking_budget, api_key],
-        outputs=output
+        inputs=[msg, chat_history, model, system_prompt, thinking_budget, api_key],
+        outputs=[chatbot, msg]
     )
 
     clear.click(
-        lambda: "",
+        handle_clear,
         inputs=None,
-        outputs=output
+        outputs=[chatbot, msg]
     )
 
 demo.launch()