Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,58 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+
 import streamlit as st
+from google import genai
+import os
+# Set GEMINI_API_KEY from GOOGLE_API_KEY if not already set
+if not os.getenv('GEMINI_API_KEY'):
+    os.environ["GEMINI_API_KEY"] = os.getenv("GOOGLE_API_KEY")
+
+# initializes chat client
+if 'client' not in st.session_state:
+    st.session_state['client'] = genai.Client()
+
+st.set_page_config(page_title="AI Chat Demo", page_icon="💬", layout="centered")
+
+# Initialize chat history in session state
+if "messages" not in st.session_state:
+    st.session_state["messages"] = [
+        {"role": "ai", "content": "👋 Welcome! How can I help you today?"}
+    ]
+    st.session_state["dummy_idx"] = 0
+
+
+st.title("AI Chat Demo")
+
+# Chat history display
+
+chat_container = st.container()
+with chat_container:
+    for msg in st.session_state["messages"]:
+        if msg["role"] == "user":
+            st.markdown(f"<div style='text-align: right; color: #1a73e8;'><b>You:</b> {msg['content']}</div>", unsafe_allow_html=True)
+        else:
+            st.markdown(f"<div style='text-align: left; color: #444;'><b>AI:</b> {msg['content']}</div>", unsafe_allow_html=True)
+
+# User input
+with st.form(key="chat_form", clear_on_submit=True):
+    user_input = st.text_input("Type your message:", "", key="input")
+    submitted = st.form_submit_button("Send")
+
+_ = st.button("Clear Chat", on_click=lambda: st.session_state.clear(), key="clear_chat")
+
+
+if submitted and user_input.strip():
+    # Add user message
+    st.session_state["messages"].append({"role": "user", "content": user_input.strip()})
+
+    # Prepare the content for the AI model by joining all messages with role tags
+    content = " /n ".join(f'<{msg["role"]}> {msg["content"]}' for msg in st.session_state["messages"])
+
+    # Generate AI response using the prepared content
+    ai_message = st.session_state['client'].models.generate_content(
+        model="gemini-2.5-flash",
+        contents=content
+    ).text
+
+    st.session_state["messages"].append({"role": "ai", "content": ai_message})
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+    st.rerun()