Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,16 +1,98 @@
+# import streamlit as st
+# import uuid
+# import json
+# import asyncio
+# import websockets
+
+# MIDDLEWARE_WS = "wss://partha181098-middleware.hf.space/chat"
+
+# st.set_page_config(page_title="Secure AI Chat")
+
+# ss = st.session_state
+
+# # ------------------ SESSION ------------------
+# def new_session():
+#     ss.session_id = str(uuid.uuid4())
+#     ss.messages = []
+
+# if "session_id" not in ss:
+#     new_session()
+
+# if "messages" not in ss:
+#     ss.messages = []
+
+# # ------------------ UI ------------------
+# st.title("🔐 Secure Database Chatbot")
+
+# if st.button("New Chat"):
+#     new_session()
+#     st.rerun()
+
+# for msg in ss.messages:
+#     with st.chat_message(msg["role"]):
+#         st.write(msg["content"])
+
+# user_input = st.chat_input("Ask something...")
+
+# # ------------------ WS CLIENT ------------------
+# async def send_and_receive(user_text):
+#     async with websockets.connect(
+#         f"{MIDDLEWARE_WS}?session_id={ss.session_id}"
+#     ) as ws:
+
+#         await ws.send(user_text)
+
+#         assistant_msg = ""
+
+#         while True:
+#             data = await ws.recv()
+#             payload = json.loads(data)
+
+#             if payload["type"] == "security_warning":
+#                 st.warning(payload["message"])
+#                 return
+
+#             if payload["type"] == "response":
+#                 assistant_msg += payload["content"]
+#                 yield assistant_msg
+
+#             if payload["type"] == "complete":
+#                 break
+
+# # ------------------ CHAT FLOW ------------------
+# if user_input:
+#     ss.messages.append({"role": "user", "content": user_input})
+#     with st.chat_message("user"):
+#         st.write(user_input)
+
+#     with st.chat_message("assistant"):
+#         placeholder = st.empty()
+#         full_response = {"text": ""}
+
+#         async def runner():
+#             async for chunk in send_and_receive(user_input):
+#                 full_response["text"] = chunk
+#                 placeholder.markdown(full_response["text"])
+
+#         asyncio.run(runner())
+
+#     ss.messages.append({"role": "assistant", "content": full_response})
+
+
 import streamlit as st
 import uuid
 import json
 import asyncio
 import websockets
 
+# ================= CONFIG =================
 MIDDLEWARE_WS = "wss://partha181098-middleware.hf.space/chat"
 
-st.set_page_config(page_title="Secure AI Chat")
+st.set_page_config(page_title="Secure AI Chat", layout="centered")
 
 ss = st.session_state
 
-# ------------------ SESSION ------------------
+# ================= SESSION =================
 def new_session():
     ss.session_id = str(uuid.uuid4())
     ss.messages = []
@@ -21,59 +103,61 @@ if "session_id" not in ss:
 if "messages" not in ss:
     ss.messages = []
 
-# ------------------ UI ------------------
-st.title("🔐 Secure Database Chatbot")
+# ================= UI =================
+st.title("🔐 Secure AI Chatbot")
 
 if st.button("New Chat"):
     new_session()
     st.rerun()
 
+# Show chat history
 for msg in ss.messages:
     with st.chat_message(msg["role"]):
-        st.write(msg["content"])
+        st.markdown(msg["content"])
 
 user_input = st.chat_input("Ask something...")
 
-# ------------------ WS CLIENT ------------------
-async def send_and_receive(user_text):
-    async with websockets.connect(
-        f"{MIDDLEWARE_WS}?session_id={ss.session_id}"
-    ) as ws:
+# ================= STREAM HANDLER =================
+async def ws_stream(user_text, placeholder):
+    uri = f"{MIDDLEWARE_WS}?session_id={ss.session_id}"
 
+    async with websockets.connect(uri) as ws:
         await ws.send(user_text)
 
-        assistant_msg = ""
+        full_response = ""
 
         while True:
-            data = await ws.recv()
-            payload = json.loads(data)
+            raw = await ws.recv()
+            payload = json.loads(raw)
 
             if payload["type"] == "security_warning":
                 st.warning(payload["message"])
-                return
+                return None
 
             if payload["type"] == "response":
-                assistant_msg += payload["content"]
-                yield assistant_msg
+                full_response += payload["content"]
+                placeholder.markdown(full_response)
 
             if payload["type"] == "complete":
-                break
+                return full_response
 
-# ------------------ CHAT FLOW ------------------
+# ================= CHAT FLOW =================
 if user_input:
+    # Store user message
     ss.messages.append({"role": "user", "content": user_input})
     with st.chat_message("user"):
-        st.write(user_input)
+        st.markdown(user_input)
 
     with st.chat_message("assistant"):
         placeholder = st.empty()
-        full_response = {"text": ""}
 
-        async def runner():
-            async for chunk in send_and_receive(user_input):
-                full_response["text"] = chunk
-                placeholder.markdown(full_response["text"])
+        # IMPORTANT: Streamlit-safe asyncio execution
+        response = asyncio.new_event_loop().run_until_complete(
+            ws_stream(user_input, placeholder)
+        )
 
-        asyncio.run(runner())
+        if response:
+            ss.messages.append(
+                {"role": "assistant", "content": response}
+            )
 
-    ss.messages.append({"role": "assistant", "content": full_response})