Update app.py

app.py

@@ -1,60 +1,80 @@
-import os
-from langchain_groq import ChatGroq
-from langchain.chains import ConversationChain
-from langchain.memory import ConversationBufferMemory
-from dotenv import load_dotenv
-import gradio as gr
-
-# Load environment variables
-load_dotenv()
-groq_api_key = os.getenv('GROQ_API_KEY')
-
-if not groq_api_key:
-    raise ValueError("GROQ_API_KEY not found in environment variables. Please add it to your .env file.")
-
-# Initialize chatbot components
-llm = ChatGroq(groq_api_key=groq_api_key, model_name="mixtral-8x7b-32768")
-memory = ConversationBufferMemory()
-conversation = ConversationChain(llm=llm, memory=memory)
-
-def add_message(history, message):
-    history.append((message, None))
-    return history, gr.Textbox(value="", interactive=True)
-
-def bot(history):
-    if not history:
-        initial_message = "Hello! How can I assist you today?"
-        history.append((None, initial_message))
-        return history
-
-    message = history[-1][0]
-    
-    try:
-        response = conversation.predict(input=message)
-    except Exception as e:
-        response = f"I'm sorry, but I encountered an error: {str(e)}"
-    
-    history[-1] = (history[-1][0], response)
-    return history
-
-def clear_chat():
-    global memory
-    memory.clear()
-    return []
-
-with gr.Blocks() as demo:
-    with gr.Column():
-        chatbot = gr.Chatbot([], elem_id="chatbot")
-        msg = gr.Textbox(label="Message")
-        clear = gr.Button("Clear")
-
-        chat_msg = msg.submit(add_message, [chatbot, msg], [chatbot, msg], queue=False).then(
-            bot, chatbot, chatbot
-        )
-
-        clear.click(clear_chat, None, chatbot, queue=False)
-
-    demo.queue()
-
-if __name__ == "__main__":
+import os
+import time
+from langchain_groq import ChatGroq
+from langchain.chains import ConversationChain
+from langchain.memory import ConversationBufferMemory
+from dotenv import load_dotenv
+import gradio as gr
+
+# Load environment variables
+load_dotenv()
+groq_api_key = os.getenv('GROQ_API_KEY')
+
+if not groq_api_key:
+    raise ValueError("GROQ_API_KEY not found in environment variables. Please add it to your .env file.")
+
+# Global variables
+stop_generation = False
+
+# Initialize chatbot components
+llm = ChatGroq(groq_api_key=groq_api_key, model_name="mixtral-8x7b-32768")
+memory = ConversationBufferMemory()
+conversation = ConversationChain(llm=llm, memory=memory)
+
+def add_message(history, message):
+    history.append((message, None))
+    return history, gr.Textbox(value="", interactive=True)
+
+def bot(history):
+    global stop_generation
+    stop_generation = False
+    
+    if not history:
+        initial_message = "Hello! How can I assist you today?"
+        history.append((None, initial_message))
+        yield history
+        return
+
+    message = history[-1][0]
+    
+    try:
+        response = conversation.predict(input=message)
+    except Exception as e:
+        response = f"I'm sorry, but I encountered an error: {str(e)}"
+    
+    history[-1] = (history[-1][0], "")
+    for character in response:
+        if stop_generation:
+            break
+        history[-1] = (history[-1][0], history[-1][1] + character)
+        time.sleep(0.01)
+        yield history
+
+def stop_response():
+    global stop_generation
+    stop_generation = True
+
+def clear_chat():
+    global memory
+    memory.clear()
+    return []
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        chatbot = gr.Chatbot([], elem_id="chatbot")
+        msg = gr.Textbox(label="Message")
+        clear = gr.Button("Clear")
+
+        chat_msg = msg.submit(add_message, [chatbot, msg], [chatbot, msg], queue=False).then(
+            bot, chatbot, chatbot
+        )
+
+        clear.click(clear_chat, None, chatbot, queue=False)
+
+        stop_btn = gr.Button("Stop Generation")
+        stop_btn.click(stop_response)
+
+    demo.queue()
+
+if __name__ == "__main__":
     demo.launch()