Update app.py

app.py

@@ -8,6 +8,7 @@ import numpy as np
 import faiss
 from sentence_transformers import SentenceTransformer
 from PIL import Image
+import io
 
 # ---------------------------
 # CONFIG
@@ -15,8 +16,14 @@ from PIL import Image
 st.set_page_config(page_title="Krish GPT Pro", page_icon="🤖", layout="wide")
 st.markdown("""
 <style>
+/* Chat bubbles */
 .chat-container { max-width: 800px; margin:auto; }
+.chat-bubble { padding: 12px; border-radius: 12px; margin-bottom: 8px; max-width: 80%; }
+.user { background-color:#646cff; color:white; margin-left:auto; text-align:right; }
+.assistant { background-color:#f1f0f0; color:black; margin-right:auto; text-align:left; }
 .stButton>button { background-color: #646cff; color:white; }
+.upload-btn { display:flex; align-items:center; gap:5px; cursor:pointer; padding:5px 10px; border:1px solid #ccc; border-radius:5px; }
+.upload-btn:hover { background-color:#f5f5f5; }
 </style>
 """, unsafe_allow_html=True)
 st.title("🤖 Krish GPT Pro")
@@ -58,7 +65,7 @@ def ocr_space_image(file, api_key):
     return text
 
 # ---------------------------
-# AUTH SYSTEM (Signup/Login/Reset)
+# AUTH SYSTEM
 # ---------------------------
 if "user" not in st.session_state:
     st.session_state.user = None
@@ -103,46 +110,21 @@ elif auth_choice == "Anonymous":
     st.sidebar.info("You are using the app anonymously. History will not be saved.")
 
 # ---------------------------
-# FILE UPLOAD & TEXT EXTRACTION
-# ---------------------------
-uploaded_file = st.file_uploader("📄 Upload PDF / 🖼️ Image", type=["pdf","png","jpg","jpeg"])
-file_text = ""
-
-if uploaded_file:
-    if uploaded_file.type == "application/pdf":
-        reader = PdfReader(uploaded_file)
-        for page in reader.pages:
-            t = page.extract_text()
-            if t:
-                file_text += t
-    elif "image" in uploaded_file.type:
-        file_text = ocr_space_image(uploaded_file, OCR_API_KEY)
-
-# ---------------------------
-# CHUNK + FAISS
+# SESSION STATE
 # ---------------------------
-def chunk_text(text, chunk_size=500):
-    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
-
-def build_index(chunks):
-    embeddings = embedder.encode(chunks)
-    dim = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dim)
-    index.add(np.array(embeddings))
-    return index, embeddings
-
-def search(query, chunks, index):
-    q_emb = embedder.encode([query])
-    D, I = index.search(np.array(q_emb), k=min(3,len(chunks)))
-    return "\n".join([chunks[i] for i in I[0]])
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-if uploaded_file and file_text:
-    chunks = chunk_text(file_text)
-    index, embeddings = build_index(chunks)
-    st.session_state.rag_data = (chunks, index)
+# Get current user safely
+current_user_resp = supabase.auth.get_user()
+if current_user_resp and current_user_resp.user:
+    user = current_user_resp.user
+else:
+    user = st.session_state.user  # fallback to session_state
+st.session_state.user = user
 
 # ---------------------------
-# CHAT FUNCTIONS
+# CHAT HISTORY FUNCTIONS
 # ---------------------------
 def save_chat(user_id, role, content):
     supabase.table("chats").insert({
@@ -155,32 +137,71 @@ def load_chat(user_id):
     data = supabase.table("chats").select("*").eq("user_id", user_id).order("created_at").execute()
     return data.data
 
-# ---------------------------
-# SESSION STATE
-# ---------------------------
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
 # Load chat history if registered
-if st.session_state.user:
-    chat_history = load_chat(st.session_state.user.id)
+if user:
+    chat_history = load_chat(user.id)
     for msg in chat_history:
-        st.chat_message(msg["role"]).write(msg["content"])
+        st.markdown(f'<div class="chat-bubble {msg["role"]}">{msg["content"]}</div>', unsafe_allow_html=True)
         st.session_state.messages.append({"role": msg["role"], "content": msg["content"]})
 
 # ---------------------------
-# USER INPUT
+# FILE UPLOAD + RAG
 # ---------------------------
-prompt = st.chat_input("💬 Ask anything...")
+if "rag_data" not in st.session_state:
+    st.session_state.rag_data = None
+
+uploaded_file = st.file_uploader("📎 Upload PDF/Image", type=["pdf","png","jpg","jpeg"], key="file_upload")
+file_text = ""
+
+if uploaded_file:
+    if uploaded_file.type == "application/pdf":
+        reader = PdfReader(uploaded_file)
+        for page in reader.pages:
+            t = page.extract_text()
+            if t:
+                file_text += t
+    elif "image" in uploaded_file.type:
+        file_text = ocr_space_image(uploaded_file, OCR_API_KEY)
+
+    if file_text:
+        # Chunk + FAISS
+        def chunk_text(text, chunk_size=500):
+            return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
+        def build_index(chunks):
+            embeddings = embedder.encode(chunks)
+            dim = embeddings.shape[1]
+            index = faiss.IndexFlatL2(dim)
+            index.add(np.array(embeddings))
+            return index, embeddings
+        def search(query, chunks, index):
+            q_emb = embedder.encode([query])
+            D, I = index.search(np.array(q_emb), k=min(3,len(chunks)))
+            return "\n".join([chunks[i] for i in I[0]])
+
+        chunks = chunk_text(file_text)
+        index, embeddings = build_index(chunks)
+        st.session_state.rag_data = (chunks, index)
+
+        # Show uploaded file inside chat
+        st.markdown(f'<div class="chat-bubble assistant">📎 Uploaded file processed successfully!</div>', unsafe_allow_html=True)
+
+# ---------------------------
+# CHAT INPUT
+# ---------------------------
+prompt = st.chat_input("💬 Type your message...")
 
 if prompt:
     st.session_state.messages.append({"role":"user","content":prompt})
-    st.chat_message("user").write(prompt)
+    st.markdown(f'<div class="chat-bubble user">{prompt}</div>', unsafe_allow_html=True)
+
     context = ""
-    if "rag_data" in st.session_state:
+    if st.session_state.rag_data:
         chunks, index = st.session_state.rag_data
-        context = search(prompt, chunks, index)
-    with st.chat_message("assistant"):
+        q_emb = embedder.encode([prompt])
+        D, I = index.search(np.array(q_emb), k=min(3,len(chunks)))
+        context = "\n".join([chunks[i] for i in I[0]])
+
+    with st.spinner("🤖 Krish GPT Pro is thinking..."):
         try:
             response = client.chat.completions.create(
                 model="llama-3.3-70b-versatile",
@@ -191,8 +212,10 @@ if prompt:
             reply = response.choices[0].message.content
         except Exception as e:
             reply = f"❌ Error: {str(e)}"
-        st.write(reply)
-        st.session_state.messages.append({"role":"assistant","content":reply})
-        if st.session_state.user:
-            save_chat(st.session_state.user.id,"user",prompt)
-            save_chat(st.session_state.user.id,"assistant",reply)
+
+    st.session_state.messages.append({"role":"assistant","content":reply})
+    st.markdown(f'<div class="chat-bubble assistant">{reply}</div>', unsafe_allow_html=True)
+
+    if user:
+        save_chat(user.id,"user",prompt)
+        save_chat(user.id,"assistant",reply)