Spaces:
Sleeping
Sleeping
| 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} | |
| ) | |