Update app.py

app.py

@@ -1,171 +1,134 @@
 import os
 from pathlib import Path
-from dotenv import load_dotenv
 import streamlit as st
+from dotenv import load_dotenv
 
-# LangChain modules
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
-from langchain_community.vectorstores import Chroma
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_community.vectorstores import FAISS
 
-# Google Generative AI
 import google.generativeai as genai
 
+# -----------------------------
+# Load API key
+# -----------------------------
 load_dotenv()
-
-# -------------------------
-# Config
-# -------------------------
-GOOGLE_API_KEY = os.getenv("GOOGLE_API")
-if not GOOGLE_API_KEY:
-    st.error("❌ GOOGLE_API key missing. Add it under: Space → Settings → Secrets.")
+GOOGLE_API = os.getenv("GOOGLE_API")
+if not GOOGLE_API:
+    st.error("❌ GOOGLE_API key missing. Add it in Space → Settings → Secrets")
     st.stop()
 
-genai.configure(api_key=GOOGLE_API_KEY)
-
-WORKDIR = Path(".")
-DATA_FILE = WORKDIR / "350_QA_dataset.pdf"
-CHROMA_DIR = WORKDIR / "chroma_db"
+genai.configure(api_key=GOOGLE_API)
 
+# -----------------------------
+# File paths
+# -----------------------------
+DATA_FILE = Path("350_QA_dataset.pdf")  # upload this PDF to Space root
+DB_DIR = Path("vectorstore")
 
-# -------------------------
-# Detect Google Models
-# -------------------------
-def pick_models():
-    models = genai.list_models()
-
-    embed = None
-    chat = None
-
-    for m in models:
-        caps = getattr(m, "supported_generation_methods", [])
-        if "embedText" in caps and embed is None:
-            embed = m.name
-        if ("generateContent" in caps or "generateText" in caps) and chat is None:
-            chat = m.name
-
-    if embed is None:
-        embed = "models/text-embedding-004"
+# -----------------------------
+# System Prompt
+# -----------------------------
+SYSTEM_PROMPT = """
+You are an EV Service Expert Assistant for a customer support team of an electric vehicle manufacturer.
+Your primary knowledge source is an internal 350-entry complaint and resolution knowledge base extracted from "350_QA_dataset.pdf".
 
-    if chat is None:
-        chat = "models/gemini-1.5-flash"  # common universal model
+You will receive retrieved chunks and the user's query. Use ONLY the retrieved context.
 
-    return embed, chat
+Respond using this structure:
+1. Issue summary
+2. Likely cause / explanation
+3. Recommended solution / actions
+4. When to visit the service center
 
+If no matching context exists, say:
+"This specific issue is not covered in my internal EV complaint database. Based on general patterns, here are some safe next steps..."
+"""
 
-# -------------------------
-# Build vectorstore
-# -------------------------
-def build_vectorstore(embed_model):
+# -----------------------------
+# Build vector store
+# -----------------------------
+def build_store():
     if not DATA_FILE.exists():
-        st.error("❌ PDF file missing. Upload 350_QA_dataset.pdf to Space root.")
+        st.error("❌ PDF file missing. Upload '350_QA_dataset.pdf' in the Space root.")
         return
 
-    st.info("📄 Loading PDF...")
     loader = PyPDFLoader(str(DATA_FILE))
     docs = loader.load()
 
-    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=800,
+        chunk_overlap=150,
+        separators=["\n\n", "\n", " ", ""]
+    )
     chunks = splitter.split_documents(docs)
 
-    st.info(f"🔢 Creating embeddings using: {embed_model}")
     embeddings = GoogleGenerativeAIEmbeddings(
-        model=embed_model,
-        google_api_key=GOOGLE_API_KEY
+        model="models/text-embedding-004",
+        google_api_key=GOOGLE_API
     )
 
-    st.info("📦 Building Chroma vector DB...")
-    if CHROMA_DIR.exists():
-        import shutil
-        shutil.rmtree(CHROMA_DIR)
-
-    db = Chroma.from_documents(chunks, embedding=embeddings, persist_directory=str(CHROMA_DIR))
-    db.persist()
+    vectorstore = FAISS.from_documents(chunks, embeddings)
+    DB_DIR.mkdir(exist_ok=True)
+    vectorstore.save_local(str(DB_DIR))
 
-    st.success("✅ Vector store created successfully!")
+    st.success("✅ Vector store built successfully!")
 
 
-# -------------------------
-# System Prompt
-# -------------------------
-SYSTEM_PROMPT = """
-You are an EV Service Expert Assistant for customer support.
-
-Use ONLY retrieved context.
-
-Respond using this structure:
-1. Issue summary
-2. Likely cause
-3. Recommended solution
-4. When to visit service center
-"""
-
-
-# -------------------------
-# RAG Query
-# -------------------------
-def rag_query(chat_model, query, history, k=4):
-    if not CHROMA_DIR.exists():
-        return "❌ Vector DB missing. Build it first.", []
-
+# -----------------------------
+# Load vector store
+# -----------------------------
+def load_store():
     embeddings = GoogleGenerativeAIEmbeddings(
-        model=chat_model,
-        google_api_key=GOOGLE_API_KEY
+        model="models/text-embedding-004",
+        google_api_key=GOOGLE_API
     )
-    db = Chroma(persist_directory=str(CHROMA_DIR), embedding_function=embeddings)
+    return FAISS.load_local(str(DB_DIR), embeddings, allow_dangerous_deserialization=True)
 
-    docs = db.similarity_search(query, k=k)
-    ctx = "\n\n---\n".join([d.page_content for d in docs]) if docs else "[No context found]"
 
-    model = genai.GenerativeModel(chat_model)
+# -----------------------------
+# Query function
+# -----------------------------
+def answer_query(query):
+    vectorstore = load_store()
+    docs = vectorstore.similarity_search(query, k=5)
 
-    prompt = [
-        SYSTEM_PROMPT,
-        "\n\nPrevious conversation:\n",
-        str(history),
-        "\n\nContext:\n",
-        ctx,
-        "\n\nUser Query:\n",
-        query
-    ]
+    context = "\n\n---\n\n".join([d.page_content for d in docs]) if docs else "[No matching context]"
 
-    response = model.generate_content(prompt)
-    return response.text, docs
+    model = genai.GenerativeModel("gemini-2.5-flash")
 
+    prompt = f"""
+{SYSTEM_PROMPT}
 
-# -------------------------
-# Streamlit UI
-# -------------------------
-st.title("⚡ EV RAG Assistant – Hugging Face Space")
+Retrieved context:
+{context}
 
-embed_model, chat_model = pick_models()
+User question:
+{query}
+"""
 
-with st.expander("Detected Models"):
-    st.write("Embedding model:", embed_model)
-    st.write("Chat model:", chat_model)
+    response = model.generate_content(prompt)
+    return response.text
 
-if st.button("📘 Build Vector Store"):
-    build_vectorstore(embed_model)
 
-st.divider()
+# -----------------------------
+# Streamlit UI
+# -----------------------------
+st.title("🔋 EV Service Expert — RAG Chatbot")
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+# Build vector store button
+if not DB_DIR.exists():
+    st.warning("Vector store missing. Click the button below to build it.")
+    if st.button("Build Vector Store"):
+        with st.spinner("Building vector store..."):
+            build_store()
 
-query = st.text_input("Enter your EV issue:")
+# Query input
+query = st.text_input("Ask a question about EV issues:")
 
 if st.button("Submit") and query.strip():
-    answer, docs = rag_query(chat_model, query, st.session_state.messages)
-    st.session_state.messages.append({"role": "user", "content": query})
-    st.session_state.messages.append({"role": "assistant", "content": answer, "sources": docs})
-
-for m in st.session_state.messages:
-    if m["role"] == "user":
-        st.markdown(f"**User:** {m['content']}")
-    else:
-        st.markdown(f"**Assistant:** {m['content']}")
-        if m.get("sources"):
-            with st.expander("Retrieved Context"):
-                for d in m["sources"]:
-                    st.write(d.page_content)
+    with st.spinner("Searching knowledge base..."):
+        answer = answer_query(query)
+    st.markdown("### 🧠 Assistant Response")
+    st.write(answer)