Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,8 +1,93 @@
+# 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", layout="centered")
+
+# 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 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.markdown(msg["content"])
+
+# user_input = st.chat_input("Ask something...")
+
+# # ================= 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)
+
+#         full_response = ""
+
+#         while True:
+#             raw = await ws.recv()
+#             payload = json.loads(raw)
+
+#             if payload["type"] == "security_warning":
+#                 st.warning(payload["message"])
+#                 return None
+
+#             if payload["type"] == "response":
+#                 full_response += payload["content"]
+#                 placeholder.markdown(full_response)
+
+#             if payload["type"] == "complete":
+#                 return full_response
+
+# # ================= CHAT FLOW =================
+# if user_input:
+#     # Store user message
+#     ss.messages.append({"role": "user", "content": user_input})
+#     with st.chat_message("user"):
+#         st.markdown(user_input)
+
+#     with st.chat_message("assistant"):
+#         placeholder = st.empty()
+
+#         # IMPORTANT: Streamlit-safe asyncio execution
+#         response = asyncio.new_event_loop().run_until_complete(
+#             ws_stream(user_input, placeholder)
+#         )
+
+#         if response:
+#             ss.messages.append(
+#                 {"role": "assistant", "content": response}
+#             )
+
 import streamlit as st
 import uuid
 import json
 import asyncio
 import websockets
+import nest_asyncio
+
+nest_asyncio.apply()
 
 # ================= CONFIG =================
 MIDDLEWARE_WS = "wss://partha181098-middleware.hf.space/chat"
@@ -40,38 +125,54 @@ user_input = st.chat_input("Ask something...")
 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)
+    full_response = ""
 
-        full_response = ""
+    try:
+        async with websockets.connect(uri) as ws:
+            await ws.send(user_text)
 
-        while True:
-            raw = await ws.recv()
-            payload = json.loads(raw)
+            while True:
+                try:
+                    raw = await ws.recv()
+                except websockets.exceptions.ConnectionClosed:
+                    break
 
-            if payload["type"] == "security_warning":
-                st.warning(payload["message"])
-                return None
+                payload = json.loads(raw)
 
-            if payload["type"] == "response":
-                full_response += payload["content"]
-                placeholder.markdown(full_response)
+                if payload["type"] == "ping":
+                    continue
+
+                if payload["type"] == "security_warning":
+                    st.warning(payload["message"])
+                    return None
+
+                if payload["type"] == "response":
+                    full_response += payload["content"]
+                    placeholder.markdown(full_response)
+
+                if payload["type"] == "complete":
+                    return full_response
+
+                if payload["type"] == "error":
+                    st.error(payload["message"])
+                    return None
+
+    except Exception as e:
+        st.error(f"WebSocket error: {e}")
+        return None
 
-            if payload["type"] == "complete":
-                return full_response
 
 # ================= CHAT FLOW =================
 if user_input:
-    # Store user message
     ss.messages.append({"role": "user", "content": user_input})
+
     with st.chat_message("user"):
         st.markdown(user_input)
 
     with st.chat_message("assistant"):
         placeholder = st.empty()
 
-        # IMPORTANT: Streamlit-safe asyncio execution
-        response = asyncio.new_event_loop().run_until_complete(
+        response = asyncio.get_event_loop().run_until_complete(
             ws_stream(user_input, placeholder)
         )
 
@@ -79,4 +180,3 @@ if user_input:
             ss.messages.append(
                 {"role": "assistant", "content": response}
             )
-