Spaces:
Running
Running
File size: 3,294 Bytes
83eb16b a6f36cf 83eb16b a6f36cf 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 |
# ----- ----- IMPORTS
import streamlit as st
import requests, time, regex, re
from datetime import datetime
from dotenv import load_dotenv
import certifi
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()
|