File size: 4,312 Bytes
cf36c54
 
ce4ef7a
cf36c54
 
 
 
 
 
 
ce4ef7a
 
 
 
 
 
 
 
 
cf36c54
 
 
 
 
 
 
 
 
79d0a75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce4ef7a
 
 
 
 
de9d96d
ce4ef7a
de9d96d
ce4ef7a
 
 
 
 
 
 
 
 
 
 
cf36c54
ce4ef7a
 
 
 
cf36c54
 
 
 
 
 
 
 
ce4ef7a
cf36c54
ce4ef7a
 
 
d35c922
cf36c54
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce4ef7a
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import openai 
import streamlit as st
import extra_streamlit_components as stx
import uuid
import time

openai.api_key = st.secrets["OPENAI_API_KEY"]
assistant_id = st.secrets["OPENAI_ASSISTANT_ID"]
client = openai

st.set_page_config(page_title="Assistant API Chat", page_icon=":speech_balloon:")

@st.cache_resource(experimental_allow_widgets=True)
def get_manager():
    return stx.CookieManager()
cookie_manager = get_manager()

st.session_state.thread_id = cookie_manager.get('thread_id')

if "start_chat" not in st.session_state:
    st.session_state.start_chat = False
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
if "messages" not in st.session_state:
    st.session_state.messages = []

st.title(":speech_balloon: Assistant API Chat")

def get_all_cookies():
    '''
    WARNING: This uses unsupported feature of Streamlit
    Returns the cookies as a dictionary of kv pairs
    '''
    from streamlit.web.server.websocket_headers import _get_websocket_headers 
    from urllib.parse import unquote

    headers = _get_websocket_headers()
    if headers is None:
        return {}
    
    if 'Cookie' not in headers:
        return {}
    
    cookie_string = headers['Cookie']
    # A sample cookie string: "K1=V1; K2=V2; K3=V3"
    cookie_kv_pairs = cookie_string.split(';')

    cookie_dict = {}
    for kv in cookie_kv_pairs:
        k_and_v = kv.split('=')
        k = k_and_v[0].strip()
        v = k_and_v[1].strip()
        cookie_dict[k] = unquote(v) #e.g. Convert name%40company.com to name@company.com
    return cookie_dict


with st.sidebar:
    st.caption(f"**Session ID**:  \n  {st.session_state.session_id}")
    st.header("Configuration")
    thread_id = st.sidebar.text_input("Enter your Thread ID:", placeholder="Leave empty to start a new thread")
    c1, c2 = st.columns(2)
    st.write(get_all_cookies())
    if c1.button("Start Chat", use_container_width=True):
        
        st.session_state.start_chat = True
        if thread_id:
            st.session_state.thread_id = thread_id
        else:
            thread = client.beta.threads.create(
                metadata={
                    'session_id': st.session_state.session_id,
                }
            )
            st.session_state.thread_id = thread.id
        cookie_manager.set('thread_id', st.session_state.thread_id)

    if c2.button("Clear Chat", use_container_width=True):
        st.session_state.start_chat = False
        st.session_state.thread_id = None
        cookie_manager.delete('thread_id')

#def process_message(message):
#    message_content = message.content[0].text
#    annotations = message_content.annotations if hasattr(message_content, 'annotations') else []
#    citations = []
#    full_response = message_content.value + '\n\n' + '\n'.join(citations)
#    return full_response

if st.session_state.thread_id:

    st.session_state.messages = client.beta.threads.messages.list(
        thread_id=st.session_state.thread_id
    )
    st.caption(f"**Thread ID**: {st.session_state.thread_id}")
    for message in reversed(st.session_state.messages.data):
        with st.chat_message(message.role):
            st.markdown(message.content[0].text.value)

    if prompt := st.chat_input():

        with st.chat_message("user"):
            st.markdown(prompt)

        client.beta.threads.messages.create(
            thread_id=st.session_state.thread_id,
            role="user",
            content=prompt
        )

        run = client.beta.threads.runs.create(
            thread_id=st.session_state.thread_id,
            assistant_id=assistant_id,
        )

        while run.status != 'completed':
            time.sleep(1)
            run = client.beta.threads.runs.retrieve(
                thread_id=st.session_state.thread_id,
                run_id=run.id
            )

        st.session_state.messages = client.beta.threads.messages.list(
            thread_id=st.session_state.thread_id
        )

        for message in reversed(st.session_state.messages.data):
            if message.run_id == run.id and message.role == "assistant":
                with st.chat_message("assistant"):
                    st.markdown(message.content[0].text.value)
else:
    st.write("Click on 'Start Chat' to start a new thread.")