MilanM commited on
Commit
83eb16b
·
verified ·
1 Parent(s): 5707c74

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +111 -38
src/streamlit_app.py CHANGED
@@ -1,40 +1,113 @@
1
- import altair as alt
2
- import numpy as np
3
- import pandas as pd
4
  import streamlit as st
 
 
 
 
5
 
6
- """
7
- # Welcome to Streamlit!
8
-
9
- Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
10
- If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
11
- forums](https://discuss.streamlit.io).
12
-
13
- In the meantime, below is an example of what you can do with just a few lines of code:
14
- """
15
-
16
- num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
17
- num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
18
-
19
- indices = np.linspace(0, 1, num_points)
20
- theta = 2 * np.pi * num_turns * indices
21
- radius = indices
22
-
23
- x = radius * np.cos(theta)
24
- y = radius * np.sin(theta)
25
-
26
- df = pd.DataFrame({
27
- "x": x,
28
- "y": y,
29
- "idx": indices,
30
- "rand": np.random.randn(num_points),
31
- })
32
-
33
- st.altair_chart(alt.Chart(df, height=700, width=700)
34
- .mark_point(filled=True)
35
- .encode(
36
- x=alt.X("x", axis=None),
37
- y=alt.Y("y", axis=None),
38
- color=alt.Color("idx", legend=None, scale=alt.Scale()),
39
- size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
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()