Changes in Interface

app.py

@@ -21,7 +21,7 @@ class ChatInput(BaseModel):
     session_id: str
     input: str
 
-def chat_logic(session_id: str, user_input: str) -> str:
+def chat_logic(session_id: str, user_input: str):
     if session_id not in session_memory_dict:
         session_memory_dict[session_id] = []
 
@@ -48,14 +48,49 @@ def chat_logic(session_id: str, user_input: str) -> str:
     print(f"[Session: {session_id}] AI Response: {ai_response}")
     return ai_response
 
-# Gradio Interface
-def gradio_chat(user_input, session_id="gradio_default"):
-    return chat_logic(session_id=session_id, user_input=user_input)
-
-gr.Interface(
-    fn=gradio_chat,
-    inputs=[gr.Textbox(label="Your message"), gr.Textbox(label="Session ID", value="gradio_default")],
-    outputs=gr.Textbox(label="Response"),
-    title="Chatbot with memory",
-    description="Chat with Groq's LLaMA3 model. Handles sessions separately using IDs."
-).launch()
+# UI logic for message display
+def chat_interface(user_input, chat_history, session_id):
+    ai_reply = chat_logic(session_id, user_input)
+    chat_history.append((user_input, ai_reply))
+    return chat_history, ""
+
+# Clear session memory
+def clear_session(session_id):
+    session_memory_dict.pop(session_id, None)
+    return [], "gradio_default"
+
+with gr.Blocks(css="""
+.chatbot .message.user { background-color: #dcf8c6 !important; color: #000; }
+.chatbot .message.bot { background-color: #ececec !important; color: #000; }
+""") as demo:
+    gr.Markdown("### 🤖 WhatsApp-style Chatbot powered by Groq LLaMA3")
+    
+    with gr.Row():
+        session_id = gr.Textbox(label="Session ID", value="gradio_default", interactive=True)
+        clear_btn = gr.Button("🔄 Clear Session")
+
+    chatbot = gr.Chatbot(label="Chat with AI", elem_classes="chatbot")
+    user_input = gr.Textbox(label="Type your message...", placeholder="Message...", lines=1)
+    send_btn = gr.Button("Send", variant="primary")
+
+    state = gr.State([])
+
+    send_btn.click(
+        fn=chat_interface,
+        inputs=[user_input, state, session_id],
+        outputs=[chatbot, user_input]
+    )
+
+    user_input.submit(
+        fn=chat_interface,
+        inputs=[user_input, state, session_id],
+        outputs=[chatbot, user_input]
+    )
+
+    clear_btn.click(
+        fn=clear_session,
+        inputs=[session_id],
+        outputs=[chatbot, session_id]
+    )
+
+demo.launch()