Update app.py

app.py

@@ -1,40 +1,58 @@
-import gradio as gr
+import os
 import google.generativeai as genai
+import gradio as gr
 
-# Set your API key
-genai.configure(api_key='YOUR_API_KEY')
+# Configure the API key
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 
-# Initialize the model
-model = genai.GenerativeModel('gemini-1.5-flash')
+# Create the model
+generation_config = {
+    "temperature": 1,
+    "top_p": 0.95,
+    "top_k": 64,
+    "max_output_tokens": 8192,
+    "response_mime_type": "text/plain",
+}
 
-def respond(user_message, history, chat_state):
-    if chat_state is None:
-        # Start a new chat with an initial greeting
-        chat_state = model.start_chat(
-            messages=[
-                {"author": "user", "content": "Hello"},
-                {"author": "assistant", "content": "Great to meet you. What would you like to know?"},
-            ]
-        )
-        # Initialize history if it's empty
-        if not history:
-            history = [["Hello", "Great to meet you. What would you like to know?"]]
-    else:
-        # Continue the conversation
-        pass  # No need to manually append messages; send_message handles it
+model = genai.GenerativeModel(
+    model_name="gemini-1.5-flash",
+    generation_config=generation_config,
+    # safety_settings can be adjusted
+    # See https://ai.google.dev/gemini-api/docs/safety-settings
+)
 
-    # Send the user's message to the model
-    response = chat_state.send_message(user_message)
-    # Append the user's message and model's response to the history
-    history.append([user_message, response.text])
-    return history, chat_state, ''
+# Function to handle chat
+def chat_with_model(user_input, history):
+    # Update the history with the user's input
+    history.append({"role": "user", "content": user_input})
+    # Start or continue the chat session
+    chat_session = model.start_chat(history=history)
+    response = chat_session.send_message(user_input)
+    # Add the assistant's response to the history
+    history.append({"role": "assistant", "content": response.text})
+    # Format the history for display in Gradio Chatbot
+    messages = []
+    for i in range(0, len(history), 2):
+        user_msg = history[i]["content"]
+        assistant_msg = history[i+1]["content"] if i+1 < len(history) else ""
+        messages.append((user_msg, assistant_msg))
+    return messages, history
 
+# Create Gradio app
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 align='center'>Gemini 1.5 Flash Chatbot Demo</h1>")
-    chatbot = gr.Chatbot([["Hello", "Great to meet you. What would you like to know?"]])
-    msg = gr.Textbox(placeholder="Type your message here...", show_label=False)
-    state = gr.State()  # To store the chat_state object
-
-    msg.submit(respond, [msg, chatbot, state], [chatbot, state, msg])
+    gr.Markdown("# Chat with Gemini Model")
+    chatbot = gr.Chatbot()
+    state = gr.State([])
+    with gr.Row():
+        user_input = gr.Textbox(
+            show_label=False,
+            placeholder="Type your message and press Enter"
+        ).style(container=False)
+        send_btn = gr.Button("Send")
+    send_btn.click(chat_with_model, [user_input, state], [chatbot, state])
+    user_input.submit(chat_with_model, [user_input, state], [chatbot, state])
+    # Clear input after sending
+    send_btn.click(lambda: "", None, user_input)
+    user_input.submit(lambda: "", None, user_input)
 
 demo.launch()