File size: 4,361 Bytes
83eb16b
a6f36cf
83eb16b
 
 
 
28d010b
a6f36cf
83eb16b
 
 
 
 
 
 
 
 
 
 
 
 
88b5a5d
83eb16b
 
 
1125c1e
83eb16b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
28d010b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1125c1e
28d010b
 
 
 
 
 
 
 
 
 
1125c1e
28d010b
 
 
 
 
 
83eb16b
 
 
 
 
 
aeb27b9
 
 
 
 
83eb16b
 
 
 
 
 
 
84e66d1
83eb16b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
be52413
83eb16b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
# ----- ----- IMPORTS
import streamlit as st
import requests, time, regex, re
from datetime import datetime
from dotenv import load_dotenv
import certifi
import os

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,
    app_password,
    wx_api_key,
    wx_project_id,
    wx_url,
    info_tag
)
from src.helper_functions import (
    setup_watsonxai_client,
    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",
)


def check_password():
    def password_entered():
        if st.session_state["password"] == app_password:
            st.session_state["password_correct"] = True
            del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        st.markdown("\n\n")
        st.text_input(
            "Enter the password",
            type="password",
            on_change=password_entered,
            key="password",
        )
        st.divider()
        st.info(info_tag)
        return False
    elif not st.session_state["password_correct"]:
        st.markdown("\n\n")
        st.text_input(
            "Enter the password",
            type="password",
            on_change=password_entered,
            key="password",
        )
        st.divider()
        st.info(info_tag)
        st.error("😕 Password incorrect")
        return False
    else:
        return True


if not check_password():
    st.stop()

if "current_page" not in st.session_state:
    st.session_state.current_page = 0


def initialize_session_state():
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

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

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