Update app.py

app.py

@@ -1,51 +1,73 @@
+import os
 import streamlit as st
 import faiss
+import numpy as np
 from sentence_transformers import SentenceTransformer
 from groq import Groq
 
-# Initialize Groq API
-client = Groq(api_key="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")  # Ensure your API key is valid
+# --- Load API key from environment (HF Repo Secrets) ---
+API_KEY = os.getenv("GROQ_API_KEY")
+if not API_KEY:
+    st.error(
+        "GROQ_API_KEY not found. In your Space go to: "
+        "Settings → Repository secrets → Add new secret (Name: GROQ_API_KEY, Value: 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)."
+    )
+    st.stop()
+
+# --- Init Groq client (no key printed/logged) ---
+client = Groq(api_key=API_KEY)
+
+# --- Cache the embedding model to speed up reloads ---
+@st.cache_resource
+def load_embedder():
+    return SentenceTransformer("all-MiniLM-L6-v2")
 
-# Initialize Sentence Transformer
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+embedding_model = load_embedder()
 
-# FAISS Index
-dimension = 384  # Embedding dimension of the model
-index = faiss.IndexFlatL2(dimension)
+# --- FAISS index (384 dims for MiniLM) ---
+dimension = 384
+if "faiss_index" not in st.session_state:
+    st.session_state.faiss_index = faiss.IndexFlatL2(dimension)
+if "chunks_store" not in st.session_state:
+    st.session_state.chunks_store = []
 
-# Function to chunk text
+index = st.session_state.faiss_index
+chunks_store = st.session_state.chunks_store
+
+# ---- Utilities ----
 def chunk_text(text, max_length=500):
-    words = text.split()
-    chunks = []
-    chunk = []
-    for word in words:
-        if len(" ".join(chunk)) + len(word) <= max_length:
-            chunk.append(word)
+    words, chunks, cur = text.split(), [], []
+    for w in words:
+        if len(" ".join(cur)) + len(w) + 1 <= max_length:
+            cur.append(w)
         else:
-            chunks.append(" ".join(chunk))
-            chunk = [word]
-    if chunk:
-        chunks.append(" ".join(chunk))
+            chunks.append(" ".join(cur))
+            cur = [w]
+    if cur:
+        chunks.append(" ".join(cur))
     return chunks
 
-# Function to embed text and add to FAISS index
 def embed_and_store(chunks):
-    embeddings = embedding_model.encode(chunks)
-    index.add(embeddings)
-
-# Query handling using Groq's streaming completions
-def query_llm(prompt):
-    # Create a completion request using the Groq model
-    completion = client.chat.completions.create(
-        model="deepseek-r1-distill-llama-70b",  # Use the provided Groq model
+    if not chunks:
+        return
+    embs = embedding_model.encode(chunks, convert_to_numpy=True, normalize_embeddings=False)
+    # Ensure float32 for FAISS
+    embs = np.asarray(embs, dtype="float32")
+    index.add(embs)
+    chunks_store.extend(chunks)
+
+def query_llm(prompt: str) -> str:
+    # Streaming chat completion
+    stream = client.chat.completions.create(
+        model="deepseek-r1-distill-llama-70b",
         messages=[
             {
                 "role": "system",
                 "content": (
-                    "You are a relationship counselor. Analyze the given WhatsApp conversation "
-                    "and provide insights on potential red flags, toxicity, and room for improvement in behavior. "
-                    "Every response must start by rating the overall chat toxicity out of 10."
-                )
+                    "You are a relationship counselor. Analyze the WhatsApp conversation and "
+                    "provide insights on red flags, toxicity, and improvements. "
+                    "Start every answer with: 'Toxicity score: X/10'."
+                ),
             },
             {"role": "user", "content": prompt},
         ],
@@ -53,41 +75,43 @@ def query_llm(prompt):
         max_completion_tokens=1024,
         top_p=0.95,
         stream=True,
-        reasoning_format="raw"
+        reasoning_format="raw",
     )
-    
-    # Stream and collect the response
-    full_response = ""
-    for chunk in completion:
-        full_response += chunk.choices[0].delta.content or ""
-    return full_response
-
-# Streamlit App
+    out = []
+    for chunk in stream:
+        delta = chunk.choices[0].delta.content or ""
+        out.append(delta)
+    return "".join(out)
+
+# ---- UI ----
 st.title("AI Relationship Counsellor")
 
-uploaded_file = st.file_uploader("Upload a text file of your WhatsApp chat", type=["txt"])
+uploaded_file = st.file_uploader("Upload a .txt export of your WhatsApp chat", type=["txt"])
 
 if uploaded_file:
-    text = uploaded_file.read().decode("utf-8")
-    st.write("Chat Extracted Successfully!")
+    text = uploaded_file.read().decode("utf-8", errors="ignore")
+    st.success("Chat extracted successfully!")
 
-    # Chunk and embed text
     chunks = chunk_text(text)
     embed_and_store(chunks)
 
-    # Query Interface
     user_query = st.text_input("Ask a question about your relationship:")
     if user_query:
-        # Embed query and search FAISS for the top 5 relevant chunks
-        query_embedding = embedding_model.encode([user_query])
-        distances, indices = index.search(query_embedding, k=5)
-        relevant_chunks = [chunks[i] for i in indices[0]]
-
-        # Combine chunks to form context
-        context = " ".join(relevant_chunks)
-        final_prompt = f"Context: {context}\n\nQuestion: {user_query}"
-
-        # Get response from the Groq model
-        response = query_llm(final_prompt)
-        st.write("### AI Analysis")
-        st.write(response)
+        # Search top-k relevant chunks
+        k = min(5, index.ntotal) if index.ntotal > 0 else 0
+        if k == 0:
+            st.warning("No text indexed yet. Please upload a chat file.")
+        else:
+            q_emb = embedding_model.encode([user_query], convert_to_numpy=True)
+            q_emb = np.asarray(q_emb, dtype="float32")
+            distances, idxs = index.search(q_emb, k)
+            relevant = [chunks_store[i] for i in idxs[0] if 0 <= i < len(chunks_store)]
+
+            context = " ".join(relevant)
+            final_prompt = f"Context:\n{context}\n\nQuestion:\n{user_query}"
+
+            with st.spinner("Analyzing…"):
+                answer = query_llm(final_prompt)
+
+            st.markdown("### AI Analysis")
+            st.write(answer)