Update app.py

app.py

@@ -1,26 +1,25 @@
 import os
 import time
 import uuid
+import sqlite3
 from typing import List, Tuple, Optional, Dict, Union
+from PIL import Image
+from io import BytesIO
 
-import sqlite3
 import google.generativeai as genai
 import streamlit as st
-from PIL import Image
-import requests
-from io import BytesIO
 
 # Database setup
 conn = sqlite3.connect('chat_history.db') 
 c = conn.cursor()
 
 c.execute('''
-         CREATE TABLE IF NOT EXISTS history 
-         (role TEXT, message TEXT)
-         ''')
+       CREATE TABLE IF NOT EXISTS history 
+       (role TEXT, message TEXT)
+       ''')
 
 # Generative AI setup
-api_key = "AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM"  
+api_key = "AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM" 
 genai.configure(api_key=api_key)
 
 generation_config = {
@@ -33,70 +32,67 @@ safety_settings = []
 # Streamlit UI
 st.set_page_config(page_title="Chatbot", page_icon="🤖")
 
-# Header
+# Header with logo
 st.markdown("""
 <style>
 .container {
-   display: flex;
+ display: flex;
 }
 .logo-text {
-   font-weight:700 !important;
-   font-size:50px !important;
-   color: #f9a01b !important;
-   padding-top: 75px !important;
+ font-weight:700 !important;
+ font-size:50px !important;
+ color: #f9a01b !important;
+ padding-top: 75px !important;
 }
 .logo-img {
-   float:right;
+ float:right;
 }
 </style>
 <div class="container">
-   <p class="logo-text">Chatbot</p>
-   <img class="logo-img" src="https://media.roboflow.com/spaces/gemini-icon.png" width=120 height=120>
+ <p class="logo-text">Chatbot</p>
+ <img class="logo-img" src="https://media.roboflow.com/spaces/gemini-icon.png" width=120 height=120>
 </div>
 """, unsafe_allow_html=True)
 
-# Sidebar
+# Sidebar for parameters and model selection
 st.sidebar.title("Parameters")
 temperature = st.sidebar.slider(
-   "Temperature",
-   min_value=0.0,
-   max_value=1.0,
-   value=0.9,
-   step=0.01,
-   help="Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that expect a true or correct response, while higher temperatures can lead to more diverse or unexpected results."
+ "Temperature",
+ min_value=0.0,
+ max_value=1.0,
+ value=0.9,
+ step=0.01,
+ help="Temperature controls the degree of randomness in token selection. Lower temperatures are good for prompts that expect a true or correct response, while higher temperatures can lead to more diverse or unexpected results."
 )
 max_output_tokens = st.sidebar.slider(
-   "Token limit",
-   min_value=1,
-   max_value=2048,
-   value=3000,
-   step=1,
-   help="Token limit determines the maximum amount of text output from one prompt. A token is approximately four characters. The default value is 2048."
+ "Token limit",
+ min_value=1,
+ max_value=2048,
+ value=3000,
+ step=1,
+ help="Token limit determines the maximum amount of text output from one prompt. A token is approximately four characters. The default value is 2048."
 )
 st.sidebar.title("Model")
 model_name = st.sidebar.selectbox(
-   "Select a model",
-   options=["gemini-pro", "gemini-pro-vision"],
-   index=0,
-   help="Gemini Pro is a text-only model that can generate natural language responses based on the chat history. Gemini Pro Vision is a multimodal model that can generate natural language responses based on the chat history and the uploaded images."
+ "Select a model",
+ options=["gemini-pro", "gemini-pro-vision"],
+ index=0,
+ help="Gemini Pro is a text-only model that can generate natural language responses based on the chat history. Gemini Pro Vision is a multimodal model that can generate natural language responses based on the chat history and the uploaded images."
 )
 
 # Initialize user_input in session state
-st.session_state["user_input"] = ""
+if "user_input" not in st.session_state:
+    st.session_state["user_input"] = ""
 
 # Chat history
 st.title("Chatbot")
-chat_history = st.session_state.get("chat_history", []) 
+if "chat_history" not in st.session_state:
+    st.session_state["chat_history"] = [] 
 
-if len(chat_history) % 2 == 0:
-   role = "user"
-else:
-   role = "model"
+for message in st.session_state["chat_history"]:
+    r, t = message["role"], message["parts"][0]["text"]
+    st.markdown(f"**{r.title()}:** {t}")
 
-for message in chat_history:
-   r, t = message["role"], message["parts"][0]["text"]
-   st.markdown(f"**{r.title()}:** {t}")
-   
 # User input
 user_input = st.text_area("", height=5, key="user_input") 
 
@@ -105,9 +101,9 @@ uploaded_files = st.file_uploader("Upload images here or paste screenshots", typ
 
 # If files are uploaded, open and display them
 if uploaded_files:
-   for uploaded_file in uploaded_files:
-       image = Image.open(uploaded_file)
-       st.image(image)
+    for uploaded_file in uploaded_files:
+        image = Image.open(uploaded_file)
+        st.image(image)
 
 # Run button
 run_button = st.button("Run", key="run_button")
@@ -125,62 +121,64 @@ progress_bar = st.progress(0)
 st.markdown("""
 <style>
 .footer {
-   position: fixed;
-   left: 0;
-   bottom: 0;
-   width: 100%;
-   background-color: #f9a01b;
-   color: white;
-   text-align: center;
+ position: fixed;
+ left: 0;
+ bottom: 0;
+ width: 100%;
+ background-color: #f9a01b;
+ color: white;
+ text-align: center;
 }
 </style>
 <div class="footer">
-   <p>Made with Streamlit and Google Generative AI</p>
+ <p>Made with Streamlit and Google Generative AI</p>
 </div>
 """, unsafe_allow_html=True)
 
 # Clear chat history and image uploader
 if clear_button:
-   chat_history.clear()
-   st.session_state["chat_history"] = chat_history
-   st.session_state["user_input"] = ""
-   st.session_state["uploaded_files"] = None
-   st.experimental_rerun()
+    st.session_state["chat_history"] = []
+# Update progress bar
+progress_bar.progress(1)
+
+# Generate model response
+if run_button:
+    if model_name == "gemini-pro":
+        response = genai.generate(
+            prompt=st.session_state["user_input"],
+            max_tokens=max_output_tokens,
+            temperature=temperature,
+            safety_settings=safety_settings
+        )
+    elif model_name == "gemini-pro-vision":
+        images = [Image.open(file).convert('RGB') for file in uploaded_files]
+        response = genai.generate(
+            prompt=st.session_state["user_input"],
+            max_tokens=max_output_tokens,
+            temperature=temperature,
+            safety_settings=safety_settings,
+            images=images
+        )
+
+    # Add model response to chat history
+    st.session_state["chat_history"].append({"role": "model", "parts": [{"text": response}]})
+
+    # Save chat history to database
+    c.execute("INSERT INTO history VALUES (?, ?)", (role, message))
+    conn.commit()
+
+    # Clear user input
+    st.session_state["user_input"] = ""
+
+    # Rerun the app
+    st.experimental_rerun()
 
 # Save chat history to a text file
 if download_button:
-   chat_text = "\n".join([f"{r.title()}: {t}" for r, t in chat_history])
-   st.download_button(
-       label="Download chat history",
-       data=chat_text,
-       file_name="chat_history.txt",
-       mime="text/plain"
-   )
-
- # Generate model response
-  if model_name == "gemini-pro":
-      response = genai.generate(
-          prompt=user_input,
-          temperature=temperature,
-          max_output_tokens=max_output_tokens
-      )
-  elif model_name == "gemini-pro-vision":
-      response = genai.generate(
-          prompt=user_input,
-          temperature=temperature,
-          max_output_tokens=max_output_tokens,
-          images=uploaded_files
-      )
-
-  # Update chat history
-  chat_history.append({"role": "model", "parts": [{"text": response}]})
-  st.session_state["chat_history"] = chat_history
-
-  # Update progress bar
-  progress_bar.progress(1)
-
-  # Clear user input
-  st.session_state["user_input"] = ""
-
-  # Rerun the app
-  st.experimental_rerun()
+    chat_text = "\n".join([f"{r.title()}: {t}" for r, t in st.session_state["chat_history"]])
+    st.download_button(
+        label="Download chat history",
+        data=chat_text,
+        file_name="chat_history.txt",
+        mime="text/plain"
+    )