eaglelandsonce/VertexAIAgentExample

1import vertexai2import streamlit as st3from vertexai.generative_models import GenerativeModel4import vertexai.preview.generative_models as generative_models5 6PROJECT_ID = "koverholt-devrel-355716"7LOCATION = "us-central1"8 9vertexai.init(project=PROJECT_ID, location=LOCATION)10 11model = GenerativeModel(12    "gemini-1.5-pro-001",13    system_instruction="""You are a pirate that gives life advice"""14)15 16generation_config = {17    "max_output_tokens": 8192,18    "temperature": 1,19    "top_p": 0.95,20}21 22safety_settings = {23    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,24    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,25    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,26    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,27}28 29st.title("Chat with Gemini")30 31if "chat" not in st.session_state:32    st.session_state["chat"] = model.start_chat()33 34if "messages" not in st.session_state:35    st.session_state.messages = []36 37for message in st.session_state.messages:38    with st.chat_message(message["role"]):39        st.markdown(message["content"])40 41if prompt := st.chat_input("Ask me anything!"):42    st.session_state.messages.append({"role": "user", "content": prompt})43    with st.chat_message("assistant"):44        response = st.session_state["chat"].send_message(45            prompt,46            generation_config=generation_config,47            safety_settings=safety_settings48        )49        st.markdown(response.text)50        st.session_state.messages.append({"role": "assistant", "content": response.text})51