Spaces:
Sleeping
Sleeping
import streamlit as st | |
import google.generativeai as genai | |
from dotenv import load_dotenv | |
import os | |
# Load environment variables | |
load_dotenv() | |
# Configure Google Generative AI with API key | |
api_key = os.getenv("GENERATIVEAI_API_KEY") | |
genai.configure(api_key=api_key) | |
# Initialize session state for message history if it doesn't exist | |
if 'messages' not in st.session_state: | |
st.session_state.messages = [] | |
# Initialize session state for chat | |
if 'chat' not in st.session_state: | |
st.session_state.chat = None | |
# Generation configuration and safety settings | |
generation_config = { | |
"temperature": 0.9, | |
"top_p": 0.5, | |
"top_k": 5, | |
"max_output_tokens": 1000, | |
} | |
safety_settings = [ | |
{ | |
"category": "HARM_CATEGORY_HARASSMENT", | |
"threshold": "BLOCK_MEDIUM_AND_ABOVE" | |
}, | |
{ | |
"category": "HARM_CATEGORY_HATE_SPEECH", | |
"threshold": "BLOCK_MEDIUM_AND_ABOVE" | |
}, | |
{ | |
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", | |
"threshold": "BLOCK_MEDIUM_AND_ABOVE" | |
}, | |
{ | |
"category": "HARM_CATEGORY_DANGEROUS_CONTENT", | |
"threshold": "BLOCK_MEDIUM_AND_ABOVE" | |
}, | |
] | |
def text_summary(text, isNew=False): | |
if isNew or st.session_state.chat is None: | |
model = genai.GenerativeModel( | |
model_name="gemini-pro", | |
generation_config=generation_config, | |
safety_settings=safety_settings | |
) | |
st.session_state.chat = model.start_chat() | |
st.session_state.chat.send_message(""" | |
Act as a financial advisor and generate financial summaries in a structured and tabular format... | |
""") # Your existing prompt here | |
response = st.session_state.chat.send_message(text) | |
return response.text | |
# Custom CSS for chat interface | |
st.markdown(""" | |
<style> | |
.chat-container { | |
height: 600px; | |
overflow-y: auto; | |
padding: 20px; | |
margin-bottom: 100px; | |
} | |
.user-message { | |
background-color: #e6f3ff; | |
padding: 10px; | |
border-radius: 10px; | |
margin: 10px 0; | |
max-width: 80%; | |
margin-left: auto; | |
} | |
.bot-message { | |
background-color: #f0f0f0; | |
padding: 10px; | |
border-radius: 10px; | |
margin: 10px 0; | |
max-width: 80%; | |
} | |
.input-container { | |
position: fixed; | |
bottom: 0; | |
left: 0; | |
right: 0; | |
padding: 20px; | |
background-color: white; | |
border-top: 1px solid #ddd; | |
} | |
.stTextInput input { | |
border-radius: 20px; | |
} | |
</style> | |
""", unsafe_allow_html=True) | |
# Main title | |
st.title("Financial Summary Chatbot") | |
# Chat message container | |
chat_container = st.container() | |
# Function to process the message | |
def process_message(message: str): | |
if message.strip(): | |
# Add user message to history | |
st.session_state.messages.append({"role": "user", "content": message}) | |
# Get bot response | |
response = text_summary(message) | |
# Add bot response to history | |
st.session_state.messages.append({"role": "assistant", "content": response}) | |
# Create a container for the input area at the bottom | |
with st.container(): | |
st.markdown('<div class="input-container">', unsafe_allow_html=True) | |
# Create two columns for input and button | |
col1, col2 = st.columns([5,1]) | |
with col1: | |
user_input = st.text_input("Message", key="user_input", label_visibility="collapsed") | |
with col2: | |
if st.button("Send"): | |
process_message(user_input) | |
# Instead of trying to clear the input directly, we'll use a rerun | |
st.rerun() | |
st.markdown('</div>', unsafe_allow_html=True) | |
# Handle Enter key press | |
if user_input and len(user_input.strip()) > 0: | |
if '\n' in user_input or st.session_state.get('enter_pressed', False): | |
process_message(user_input) | |
st.session_state.enter_pressed = False | |
st.rerun() | |
# Display chat messages in the container | |
with chat_container: | |
st.markdown('<div class="chat-container">', unsafe_allow_html=True) | |
for message in st.session_state.messages: | |
if message["role"] == "user": | |
st.markdown(f'<div class="user-message">{message["content"]}</div>', unsafe_allow_html=True) | |
else: | |
st.markdown(f'<div class="bot-message">{message["content"]}</div>', unsafe_allow_html=True) | |
st.markdown('</div>', unsafe_allow_html=True) |