Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,113 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+# ----- ----- IMPORTS
 import streamlit as st
+import requests, time, regex, re
+from datetime import datetime
+from dotenv import load_dotenv
+import certifi
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+load_dotenv()
+
+from ibm_watsonx_ai import Credentials, APIClient
+from ibm_watsonx_ai.foundation_models import ModelInference
+from jinja2 import Template
+
+from src.bot_specs import bot_name, bot_icon, user_icon
+from src.parameters import (
+    model_id,
+    system_prompt,
+    params,
+    display_chat_history,
+    stream_outputs,
+    wx_api_key,
+    wx_project_id,
+    wx_url,
+)
+from src.helper_functions import (
+    check_password,
+    setup_watsonxai_client,
+    initialize_session_state,
+    create_pdf_from_chat,
+    watsonx_chat_prompt,
+    generate_response,
+)
+
+# ----- ----- PAGE CONFIG
+st.set_page_config(
+    page_title=bot_name,
+    page_icon=bot_icon,
+    initial_sidebar_state="collapsed",
+    layout="centered",
+)
+
+if not check_password():
+    st.stop()
+
+if "current_page" not in st.session_state:
+    st.session_state.current_page = 0
+
+wx_client = setup_watsonxai_client(
+    api_key=wx_api_key, project_id=wx_project_id, url=wx_url
+)
+
+
+# ----- ----- MAIN APP
+def main():
+    initialize_session_state()
+    st.subheader(f"{bot_name} {bot_icon}")
+
+    if display_chat_history:
+        for message in st.session_state.chat_history:
+            with st.chat_message(
+                message["role"],
+                avatar=user_icon if message["role"] == "user" else bot_icon,
+            ):
+                st.markdown(message["content"])
+    user_input = st.chat_input("You:", key="user_input")
+
+    if user_input:
+        # Add user message to chat history
+        st.session_state.chat_history.append({"role": "user", "content": user_input})
+        with st.chat_message("user", avatar=user_icon):
+            st.markdown(user_input)
+
+        with st.chat_message(bot_name, avatar=bot_icon):
+            # Build messages with baseline + chat history
+            messages = [{"role": "system", "content": system_prompt}]
+
+            messages.extend(st.session_state.chat_history)
+
+            stream_generator = watsonx_chat_prompt(
+                messages=messages,
+                stream=stream_outputs,
+                client=wx_client,
+                params=params,
+                model_id=model_id,
+            )
+            print(messages)
+            text_output = generate_response(stream_generator, stream=stream_outputs)
+
+            # Stream the response with typewriter effect
+            assistant_response = st.write_stream(text_output)
+
+            # Add assistant response to chat history
+            st.session_state.chat_history.append(
+                {"role": "assistant", "content": assistant_response}
+            )
+
+    if st.session_state.chat_history:
+        now = datetime.now()
+        date_str = now.strftime("%Y-%m-%d")
+        try:
+            pdf_buffer = create_pdf_from_chat(st.session_state.chat_history)
+            st.download_button(
+                label="Download Chat History as PDF",
+                data=pdf_buffer,
+                file_name=f"chat_history_{date_str}.pdf",
+                mime="application/pdf",
+            )
+        except Exception as e:
+            st.error(f"An error occurred while generating the PDF: {str(e)}")
+
+
+if __name__ == "__main__":
+    main()