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}
            )