Update app.py

app.py

@@ -22,10 +22,7 @@ import sys
 st.set_page_config(page_title="📥 Email Chat Application", layout="wide")
 st.title("💬 Turn Emails into Conversations—Effortless Chat with Your Inbox! 📩")
 
-# ===============================
-# 2. Gmail Authentication Configuration
-# ===============================
-SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']
+# Initialize session state variables
 if "authenticated" not in st.session_state:
     st.session_state.authenticated = False
 if "creds" not in st.session_state:
@@ -49,11 +46,15 @@ if "candidate_context" not in st.session_state:
 if "raw_candidates" not in st.session_state:
     st.session_state.raw_candidates = None
 
+# Initialize chat messages
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
 def count_tokens(text):
     return len(text.split())
 
 # ===============================
-# 3. Gmail Authentication Functions
+# 2. Gmail Authentication Functions
 # ===============================
 def reset_session_state():
     st.session_state.authenticated = False
@@ -66,6 +67,9 @@ def reset_session_state():
     st.session_state.vector_store = None
     st.session_state.candidate_context = None
     st.session_state.raw_candidates = None
+    st.session_state.messages = []
+    st.session_state.candidates_message_shown = False
+    st.session_state.vector_db_message_shown = False
     for filename in ["token.json", "data_chunks.pkl", "embeddings.pkl", "vector_store.index", "vector_database.pkl"]:
         if os.path.exists(filename):
             os.remove(filename)
@@ -79,7 +83,9 @@ def authenticate_gmail(credentials_file):
             if creds and creds.valid:
                 st.session_state.creds = creds
                 st.session_state.authenticated = True
-                st.success("✅ Authentication successful!")
+                if not st.session_state.get('candidates_message_shown', False):
+                    st.success("✅ Authentication successful!")
+                    st.session_state.candidates_message_shown = True
                 return creds
         except Exception as e:
             st.error(f"❌ Invalid token.json file: {e}")
@@ -89,7 +95,9 @@ def authenticate_gmail(credentials_file):
             creds.refresh(Request())
             st.session_state.creds = creds
             st.session_state.authenticated = True
-            st.success("✅ Authentication successful!")
+            if not st.session_state.get('candidates_message_shown', False):
+                st.success("✅ Authentication successful!")
+                st.session_state.candidates_message_shown = True
             with open('token.json', 'w') as token_file:
                 token_file.write(creds.to_json())
             return creds
@@ -107,14 +115,16 @@ def submit_auth_code():
         st.session_state.flow.fetch_token(code=st.session_state.auth_code)
         st.session_state.creds = st.session_state.flow.credentials
         st.session_state.authenticated = True
+        if not st.session_state.get('candidates_message_shown', False):
+            st.success("✅ Authentication successful!")
+            st.session_state.candidates_message_shown = True
         with open('token.json', 'w') as token_file:
             token_file.write(st.session_state.creds.to_json())
-        st.success("✅ Authentication successful!")
     except Exception as e:
         st.error(f"❌ Error during authentication: {e}")
 
 # ===============================
-# 4. Email Data Extraction, Embedding and Vector Store Functions
+# 3. Email Data Extraction, Embedding and Vector Store Functions
 # ===============================
 def extract_email_body(payload):
     if 'body' in payload and 'data' in payload['body'] and payload['body']['data']:
@@ -187,7 +197,9 @@ def create_chunks_from_gmail(service, label):
             data_chunks.append(email_dict)
             progress_bar.progress(min((idx + 1) / total, 1.0))
         st.session_state.data_chunks.extend(data_chunks)
-        st.success(f"✅ Data chunks created successfully from {label}! Total emails processed for this label: {len(data_chunks)}")
+        if not st.session_state.get('vector_db_message_shown', False):
+            st.success(f"✅ Data chunks created successfully from {label}! Total emails processed for this label: {len(data_chunks)}")
+            st.session_state.vector_db_message_shown = True
     except Exception as e:
         st.error(f"❌ Error creating chunks from Gmail for label {label}: {e}")
 
@@ -251,7 +263,7 @@ def save_vector_database():
         st.error(f"❌ Error saving vector database: {e}")
 
 # ===============================
-# 5. Handling User Queries (User-Controlled Threshold)
+# 4. Handling User Queries (User-Controlled Threshold)
 # ===============================
 def preprocess_query(query):
     return query.lower().strip()
@@ -283,9 +295,6 @@ def process_candidate_emails(query, similarity_threshold):
             )
             faiss.normalize_L2(query_embedding)
             
-            # Debug: Verify the type of vector_store
-            # st.write(f"Vector Store Type: {type(st.session_state.vector_store)}")
-            
             # Perform search
             distances, indices = st.session_state.vector_store.search(query_embedding, TOP_K)
             candidates = []
@@ -310,7 +319,8 @@ def process_candidate_emails(query, similarity_threshold):
 
             st.session_state.candidate_context = context_str
             st.session_state.raw_candidates = candidates
-            st.success("✅ Candidates retrieved and context built!")
+            # Append success message to chat messages
+            st.session_state.messages.append({"role": "system", "content": "✅ Candidates retrieved and context built!"})
         except Exception as e:
             st.error(f"❌ An error occurred during processing: {e}")
 
@@ -349,8 +359,8 @@ def call_llm_api(query):
             response.raise_for_status()  # Raises stored HTTPError, if one occurred.
             response_json = response.json()
             generated_text = response_json["choices"][0]["message"]["content"]
-            st.subheader("📝 AI Response:")
-            st.write(generated_text)
+            # Append AI response to chat messages
+            st.session_state.messages.append({"role": "assistant", "content": generated_text})
         except requests.exceptions.HTTPError as http_err:
             try:
                 error_info = response.json().get("error", {})
@@ -376,37 +386,36 @@ def handle_user_query():
             key='similarity_threshold'
         )
 
-    # Text input with callback on change (when Enter is pressed)
-    def query_callback():
-        query = st.session_state.query_input
-        if not query.strip():
-            st.warning("⚠️ Please enter a valid query.")
-            return
-        process_candidate_emails(query, similarity_threshold)
-        if st.session_state.raw_candidates:
-            with st.expander("🔎 Matching Email Chunks:", expanded=False):
-                for candidate, sim in st.session_state.raw_candidates:
-                    # Get a snippet (first 150 characters) of the body instead of full body content.
-                    body = candidate.get('body', 'No Content')
-                    snippet = (body[:150] + "...") if len(body) > 150 else body
-                    st.markdown(
-                        f"**From:** {candidate.get('sender','Unknown')}  <br>"
-                        f"**To:** {candidate.get('to','Unknown')}  <br>"
-                        f"**Date:** {candidate.get('date','Unknown')}  <br>"
-                        f"**Subject:** {candidate.get('subject','No Subject')}  <br>"
-                        f"**Body Snippet:** {snippet}  <br>"
-                        f"**Similarity:** {sim:.4f}",
-                        unsafe_allow_html=True
-                    )
-            # Then send the query along with the context to the LLM API.
-            call_llm_api(query)
-        # Clear the input field after processing
-        st.session_state.query_input = ""
-
-    st.text_input("Enter your query:", key="query_input", on_change=query_callback)
+    # Chat input for user queries
+    user_input = st.chat_input("Enter your query:")
+
+    if user_input:
+        # Append user message to chat
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        # Process the query
+        process_candidate_emails(user_input, similarity_threshold)
+        
+        if st.session_state.candidate_context:
+            # Send the query to the LLM API
+            call_llm_api(user_input)
+        
+        # Clear the input field is handled automatically by st.chat_input
+
+    # Display chat messages
+    for msg in st.session_state.messages:
+        if msg["role"] == "user":
+            with st.chat_message("user"):
+                st.markdown(msg["content"])
+        elif msg["role"] == "assistant":
+            with st.chat_message("assistant"):
+                st.markdown(msg["content"])
+        elif msg["role"] == "system":
+            with st.chat_message("system"):
+                st.markdown(msg["content"])
 
 # ===============================
-# 6. Main Application Logic
+# 5. Main Application Logic
 # ===============================
 def main():
     st.sidebar.header("🔒 Gmail Authentication")
@@ -430,7 +439,9 @@ def main():
                 st.session_state.vector_store = vector_db.get("vector_store")
                 st.session_state.embeddings = vector_db.get("embeddings")
                 st.session_state.data_chunks = vector_db.get("data_chunks")
-                st.success("📁 Vector database loaded successfully from upload!")
+                if not st.session_state.get('vector_db_message_shown', False):
+                    st.success("📁 Vector database loaded successfully from upload!")
+                    st.session_state.vector_db_message_shown = True
             except Exception as e:
                 st.error(f"❌ Error loading vector database: {e}")
     elif data_management_option == "Authenticate and Create New Data":