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

# Setup page config
st.set_page_config(
    page_title="Gemini Streaming Demo",
    page_icon="🤖",
    layout="wide"
)

# Create a title
st.title("🤖 Gemini API Streaming Demo")

# Configure the Gemini API with environment variable
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])

# Initialize the model
model = genai.GenerativeModel('gemini-1.5-flash')

# Create input text area
user_input = st.text_area("Enter your prompt:", height=100)

# Create a button to generate content
if st.button("Generate", type="primary"):
    if user_input:
        # Create a placeholder for the streaming text
        response_placeholder = st.empty()
        
        # Initialize full_response
        full_response = ""
        
        # Generate streaming response
        try:
            response = model.generate_content(user_input, stream=True)
            
            # Stream the response
            for chunk in response:
                if chunk.text:
                    full_response += chunk.text
                    # Update the placeholder with the full response so far
                    response_placeholder.markdown(full_response + "▌")
                    time.sleep(0.05)  # Add a small delay for better visualization
            
            # Final update without the cursor
            response_placeholder.markdown(full_response)
            
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
    else:
        st.warning("Please enter a prompt first.")

# Add instructions in the sidebar
with st.sidebar:
    st.markdown("""
    ### Instructions
    1. Type your prompt in the text area
    2. Click 'Generate' to see the streaming response
    
    ### About
    This demo shows how to use Gemini's streaming capabilities to generate text responses in real-time.
    """)