main.py

import os
import textwrap
# import markdown_it
# from IPython.display import display
# from IPython.display import Markdown
import streamlit as st
import google.generativeai as genai

# Initialize Gemini-Pro 
genai.configure(api_key=st.secrets["GOOGLE_GEMINI_KEY"])

# Set up the model
generation_config = {
  "temperature": 0.1,
  "top_p": 1,
  "top_k": 1,
  "max_output_tokens": 2048,
}

safety_settings = [
  {
    "category": "HARM_CATEGORY_HARASSMENT",
    "threshold": "BLOCK_ONLY_HIGH"
  },
  {
    "category": "HARM_CATEGORY_HATE_SPEECH",
    "threshold": "BLOCK_ONLY_HIGH"
  },
  {
    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
    "threshold": "BLOCK_ONLY_HIGH"
  },
  {
    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
    "threshold": "BLOCK_ONLY_HIGH"
  },
]

model = genai.GenerativeModel(model_name="gemini-pro",
                              generation_config=generation_config,
                              safety_settings=safety_settings)

# def to_markdown(text):
#   text = text.replace('•', '  *')
#   return markdown_it(textwrap.indent(text, '> ', predicate=lambda _: True) )

# Gemini uses 'model' for assistant; Streamlit uses 'assistant'
def role_to_streamlit(role):
  if role == "model":
    return "Assistant"
  else:
    return role

# Add a Gemini Chat history object to Streamlit session state
if "chat" not in st.session_state:
    st.session_state.chat = model.start_chat(history = [])
# Display Form Title
st.title("Heang Project Alpha")

# Display chat messages from history above current input box
for message in st.session_state.chat.history:
    with st.chat_message(role_to_streamlit(message.role)):
        st.markdown(message.parts[0].text)

# Accept user's next message, add to context, resubmit context to Gemini
if prompt := st.chat_input("I possess a well of knowledge. What would you like to know?"):
    # Display user's last message
    st.chat_message("user").markdown(prompt)
    
    with st.spinner('💡Thinking'):
      # Send user entry to Gemini and read the response
      response = st.session_state.chat.send_message(prompt) 
      # response_markdown = to_markdown(response.text).data
      # Display last 
      with st.chat_message("Assistant"):
          st.markdown(response.text)