import streamlit as st
import google.generativeai as genai
import os

def main():
    GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
    genai.configure(api_key=GOOGLE_API_KEY)
    model = genai.GenerativeModel('gemini-1.5-flash')
    
    if prompt := st.chat_input("哈囉！請問有什麼是我可以幫忙的嗎？"):
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            def reply():
                for chunk in model.generate_content(contents=prompt, stream=True):
                    yield chunk.text
            st.write_stream(reply)

if __name__=="__main__":
    main()