Update src/streamlit_app.py

src/streamlit_app.py

@@ -11,7 +11,7 @@ from huggingface_hub import InferenceClient
 # CONFIG
 # ==============================
 st.set_page_config(page_title="Company ChatGPT", layout="wide")
-st.title("🏢 Company AI Assistant")
+st.title("🏢 Company AI Assistant (RAG Powered)")
 
 # ==============================
 # LOAD MODELS
@@ -19,9 +19,15 @@ st.title("🏢 Company AI Assistant")
 @st.cache_resource
 def load_models():
     embed_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    HF_TOKEN = os.environ.get("HF_TOKEN")
+    if not HF_TOKEN:
+        st.error("❌ Please add HF_TOKEN in Hugging Face Secrets")
+        st.stop()
+
     llm = InferenceClient(
         model="meta-llama/Meta-Llama-3-8B-Instruct",
-        token=os.environ.get("HF_TOKEN")
+        token=HF_TOKEN
     )
     return embed_model, llm
 
@@ -32,11 +38,20 @@ embed_model, llm = load_models()
 # ==============================
 @st.cache_data
 def load_data():
-    df = pd.read_csv("src/company_sample.csv")
+    path = "src/company_sample.csv"
+    if not os.path.exists(path):
+        st.error(f"❌ File not found: {path}")
+        st.stop()
+    df = pd.read_csv(path)
     return df
 
 df = load_data()
-documents = df["text"].tolist()
+
+if "text" not in df.columns:
+    st.error("❌ CSV must contain 'text' column")
+    st.stop()
+
+documents = df["text"].fillna("").tolist()
 
 # ==============================
 # CREATE VECTOR DB
@@ -46,9 +61,9 @@ def create_faiss(docs):
     embeddings = embed_model.encode(docs)
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(np.array(embeddings))
-    return index, embeddings
+    return index
 
-index, doc_embeddings = create_faiss(documents)
+index = create_faiss(documents)
 
 # ==============================
 # RETRIEVAL FUNCTION
@@ -56,7 +71,7 @@ index, doc_embeddings = create_faiss(documents)
 def retrieve(query, top_k=3):
     q_emb = embed_model.encode([query])
     D, I = index.search(np.array(q_emb), top_k)
-    return [documents[i] for i in I[0]]
+    return [documents[i] for i in I[0] if i < len(documents)]
 
 # ==============================
 # CHAT HISTORY
@@ -64,7 +79,6 @@ def retrieve(query, top_k=3):
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display history
 for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
 
@@ -77,29 +91,42 @@ if query:
     st.session_state.messages.append({"role": "user", "content": query})
     st.chat_message("user").write(query)
 
-    # 🔍 Retrieve relevant docs
+    # 🔍 Retrieve context
     context_docs = retrieve(query)
-    context = "\n".join(context_docs)
-
-    # 🧠 Build prompt
-    prompt = f"""
-You are a company assistant. Answer ONLY based on the context below.
-
+    context = "\n\n".join(context_docs)
+
+    # ==============================
+    # 🤖 LLM CALL (FIXED)
+    # ==============================
+    try:
+        response = llm.chat_completion(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a company assistant. Answer ONLY from given context. If not found, say 'Not available in company data.'"
+                },
+                {
+                    "role": "user",
+                    "content": f"""
 Context:
 {context}
 
 Question:
 {query}
-
-Answer:
 """
-
-    # 🤖 LLM Call
-    response = llm.text_generation(
-        prompt,
-        max_new_tokens=200,
-        temperature=0.5
-    )
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
-    st.chat_message("assistant").write(response)
+                }
+            ],
+            max_tokens=200,
+            temperature=0.5
+        )
+
+        answer = response.choices[0].message.content
+
+    except Exception as e:
+        answer = f"❌ Error: {str(e)}"
+
+    # ==============================
+    # DISPLAY RESPONSE
+    # ==============================
+    st.session_state.messages.append({"role": "assistant", "content": answer})
+    st.chat_message("assistant").write(answer)