Spaces:
Sleeping
Sleeping
File size: 2,118 Bytes
2ce7d31 5286ac5 2ce7d31 f748655 7700a5b 5286ac5 f748655 5286ac5 f748655 5286ac5 f748655 5286ac5 f748655 5286ac5 f748655 5286ac5 f748655 5286ac5 3f1bf29 5286ac5 3f1bf29 afc62f4 3f1bf29 afc62f4 3f1bf29 afc62f4 3f1bf29 5286ac5 3f1bf29 5286ac5 f748655 5286ac5 3f1bf29 5286ac5 f748655 5286ac5 3f1bf29 f748655 5286ac5 f748655 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
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}
)
|