Spaces:
Sleeping
Sleeping
Update src/streamlit_app.py
Browse files- src/streamlit_app.py +111 -38
src/streamlit_app.py
CHANGED
@@ -1,40 +1,113 @@
|
|
1 |
-
|
2 |
-
import numpy as np
|
3 |
-
import pandas as pd
|
4 |
import streamlit as st
|
|
|
|
|
|
|
|
|
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 |
-
st.
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# ----- ----- IMPORTS
|
|
|
|
|
2 |
import streamlit as st
|
3 |
+
import requests, time, regex, re
|
4 |
+
from datetime import datetime
|
5 |
+
from dotenv import load_dotenv
|
6 |
+
import certifi
|
7 |
|
8 |
+
load_dotenv()
|
9 |
+
|
10 |
+
from ibm_watsonx_ai import Credentials, APIClient
|
11 |
+
from ibm_watsonx_ai.foundation_models import ModelInference
|
12 |
+
from jinja2 import Template
|
13 |
+
|
14 |
+
from src.bot_specs import bot_name, bot_icon, user_icon
|
15 |
+
from src.parameters import (
|
16 |
+
model_id,
|
17 |
+
system_prompt,
|
18 |
+
params,
|
19 |
+
display_chat_history,
|
20 |
+
stream_outputs,
|
21 |
+
wx_api_key,
|
22 |
+
wx_project_id,
|
23 |
+
wx_url,
|
24 |
+
)
|
25 |
+
from src.helper_functions import (
|
26 |
+
check_password,
|
27 |
+
setup_watsonxai_client,
|
28 |
+
initialize_session_state,
|
29 |
+
create_pdf_from_chat,
|
30 |
+
watsonx_chat_prompt,
|
31 |
+
generate_response,
|
32 |
+
)
|
33 |
+
|
34 |
+
# ----- ----- PAGE CONFIG
|
35 |
+
st.set_page_config(
|
36 |
+
page_title=bot_name,
|
37 |
+
page_icon=bot_icon,
|
38 |
+
initial_sidebar_state="collapsed",
|
39 |
+
layout="centered",
|
40 |
+
)
|
41 |
+
|
42 |
+
if not check_password():
|
43 |
+
st.stop()
|
44 |
+
|
45 |
+
if "current_page" not in st.session_state:
|
46 |
+
st.session_state.current_page = 0
|
47 |
+
|
48 |
+
wx_client = setup_watsonxai_client(
|
49 |
+
api_key=wx_api_key, project_id=wx_project_id, url=wx_url
|
50 |
+
)
|
51 |
+
|
52 |
+
|
53 |
+
# ----- ----- MAIN APP
|
54 |
+
def main():
|
55 |
+
initialize_session_state()
|
56 |
+
st.subheader(f"{bot_name} {bot_icon}")
|
57 |
+
|
58 |
+
if display_chat_history:
|
59 |
+
for message in st.session_state.chat_history:
|
60 |
+
with st.chat_message(
|
61 |
+
message["role"],
|
62 |
+
avatar=user_icon if message["role"] == "user" else bot_icon,
|
63 |
+
):
|
64 |
+
st.markdown(message["content"])
|
65 |
+
user_input = st.chat_input("You:", key="user_input")
|
66 |
+
|
67 |
+
if user_input:
|
68 |
+
# Add user message to chat history
|
69 |
+
st.session_state.chat_history.append({"role": "user", "content": user_input})
|
70 |
+
with st.chat_message("user", avatar=user_icon):
|
71 |
+
st.markdown(user_input)
|
72 |
+
|
73 |
+
with st.chat_message(bot_name, avatar=bot_icon):
|
74 |
+
# Build messages with baseline + chat history
|
75 |
+
messages = [{"role": "system", "content": system_prompt}]
|
76 |
+
|
77 |
+
messages.extend(st.session_state.chat_history)
|
78 |
+
|
79 |
+
stream_generator = watsonx_chat_prompt(
|
80 |
+
messages=messages,
|
81 |
+
stream=stream_outputs,
|
82 |
+
client=wx_client,
|
83 |
+
params=params,
|
84 |
+
model_id=model_id,
|
85 |
+
)
|
86 |
+
print(messages)
|
87 |
+
text_output = generate_response(stream_generator, stream=stream_outputs)
|
88 |
+
|
89 |
+
# Stream the response with typewriter effect
|
90 |
+
assistant_response = st.write_stream(text_output)
|
91 |
+
|
92 |
+
# Add assistant response to chat history
|
93 |
+
st.session_state.chat_history.append(
|
94 |
+
{"role": "assistant", "content": assistant_response}
|
95 |
+
)
|
96 |
+
|
97 |
+
if st.session_state.chat_history:
|
98 |
+
now = datetime.now()
|
99 |
+
date_str = now.strftime("%Y-%m-%d")
|
100 |
+
try:
|
101 |
+
pdf_buffer = create_pdf_from_chat(st.session_state.chat_history)
|
102 |
+
st.download_button(
|
103 |
+
label="Download Chat History as PDF",
|
104 |
+
data=pdf_buffer,
|
105 |
+
file_name=f"chat_history_{date_str}.pdf",
|
106 |
+
mime="application/pdf",
|
107 |
+
)
|
108 |
+
except Exception as e:
|
109 |
+
st.error(f"An error occurred while generating the PDF: {str(e)}")
|
110 |
+
|
111 |
+
|
112 |
+
if __name__ == "__main__":
|
113 |
+
main()
|