Update app.py

app.py

@@ -1,96 +1,178 @@
 import streamlit as st
-from pypdf import PdfReader
-from sentence_transformers import SentenceTransformer
+import os
+import pypdf
 import chromadb
-from chromadb.utils import embedding_functions
+from sentence_transformers import SentenceTransformer
 from groq import Groq
-import os
+from typing import List, Dict, Any, Optional
+
+# CONFIG
+SIMILARITY_THRESHOLD = 0.2
+TOP_K_CHUNKS = 3
+CHUNK_SIZE = 300
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+
+# PDF extraction
+def extract_text_from_pdf(pdf_file) -> Dict[str, Any]:
+    try:
+        pdf_reader = pypdf.PdfReader(pdf_file)
+        pages_text = []
+        for page_num, page in enumerate(pdf_reader.pages):
+            page_text = page.extract_text()
+            if page_text and page_text.strip():
+                pages_text.append({
+                    'page_number': page_num + 1,
+                    'text': page_text.strip()
+                })
+        return {"success": True, "pages": pages_text, "total_pages": len(pages_text)}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# Chunking
+def create_chunks(pages_text: List[Dict]) -> List[Dict]:
+    chunks = []
+    chunk_id = 0
+    for page_data in pages_text:
+        words = page_data['text'].split()
+        for i in range(0, len(words), CHUNK_SIZE):
+            chunk_words = words[i:i + CHUNK_SIZE]
+            if len(chunk_words) > 20:
+                chunks.append({
+                    "id": chunk_id,
+                    "text": " ".join(chunk_words),
+                    "page_number": page_data['page_number'],
+                    "word_count": len(chunk_words)
+                })
+                chunk_id += 1
+    return chunks
+
+# Embedding model
+@st.cache_resource
+def load_embedding_model():
+    return SentenceTransformer(EMBEDDING_MODEL)
+
+# Vector database
+def create_vector_database(chunks: List[Dict], embedding_model) -> Optional[Any]:
+    try:
+        client = chromadb.Client()
+        # use get_or_create instead of create
+        collection = client.get_or_create_collection("pdf_chunks")
+
+        texts = [c['text'] for c in chunks]
+        embeddings = embedding_model.encode(texts).tolist()
 
-# -------------------------------
-# 1. Setup
-# -------------------------------
-st.set_page_config(page_title="📖 RAG Tutor", layout="wide")
-st.title("📚 RAG Tutor – Learn from Your Book")
-
-# Load API key from Hugging Face secrets
-api_key = os.environ.get("GROQ_API_KEY")
-if not api_key:
-    st.error("❌ Missing GROQ_API_KEY. Please add it in Hugging Face Secrets.")
-    st.stop()
-
-client = Groq(api_key=api_key)
-
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Create ChromaDB in-memory instance
-chroma_client = chromadb.Client()
-collection = chroma_client.create_collection(
-    name="book_chunks",
-    embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
-        model_name="all-MiniLM-L6-v2"
-    )
-)
-
-# -------------------------------
-# 2. PDF Upload + Processing
-# -------------------------------
-uploaded_file = st.file_uploader("📂 Upload a PDF book", type=["pdf"])
-
-if uploaded_file:
-    reader = PdfReader(uploaded_file)
-    text = ""
-    for i, page in enumerate(reader.pages):
-        page_text = page.extract_text()
-        if page_text:
-            text += f"[Page {i+1}]\n" + page_text + "\n"
-
-    # Split into ~300 word chunks
-    words = text.split()
-    chunks = [" ".join(words[i:i+300]) for i in range(0, len(words), 300)]
-
-    # Store chunks in ChromaDB
-    for idx, chunk in enumerate(chunks):
         collection.add(
-            documents=[chunk],
-            metadatas=[{"source": f"Page {idx//1+1}"}],
-            ids=[str(idx)]
+            embeddings=embeddings,
+            documents=texts,
+            metadatas=[{
+                "page_number": c["page_number"],
+                "chunk_id": c["id"],
+                "word_count": c["word_count"]
+            } for c in chunks],
+            ids=[str(c["id"]) for c in chunks]
         )
-    st.success("✅ PDF processed and stored in memory!")
-
-# -------------------------------
-# 3. Ask Questions
-# -------------------------------
-question = st.text_input("❓ Ask a question about the book")
-
-if st.button("Get Answer") and question:
-    q_embedding = embedder.encode(question).tolist()
-
-    results = collection.query(
-        query_embeddings=[q_embedding],
-        n_results=3
-    )
-
-    if results["documents"][0]:
-        context = "\n\n".join(results["documents"][0])
-        sources = [m["source"] for m in results["metadatas"][0]]
-
-        # Prompt LLM with context
+        return collection
+    except Exception as e:
+        st.error(f"Vector DB error: {e}")
+        return None
+
+def query_vector_database(collection, query: str, embedding_model, k: int = TOP_K_CHUNKS) -> List[Dict]:
+    try:
+        query_emb = embedding_model.encode([query]).tolist()
+        results = collection.query(query_embeddings=query_emb, n_results=k)
+        relevant_chunks = []
+        for i in range(len(results['documents'][0])):
+            distance = results['distances'][0][i]
+            similarity = max(0, 1 - distance)
+            if similarity >= SIMILARITY_THRESHOLD:
+                relevant_chunks.append({
+                    "text": results['documents'][0][i],
+                    "page_number": results['metadatas'][0][i]["page_number"],
+                    "similarity": similarity,
+                    "chunk_id": results['metadatas'][0][i]["chunk_id"]
+                })
+        return relevant_chunks
+    except Exception as e:
+        st.error(f"Query error: {e}")
+        return []
+
+# Groq setup
+def setup_groq():
+    api_key = st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")
+    if not api_key:
+        st.error("❌ No GROQ_API_KEY found. Please add it to secrets or env.")
+        return None
+    return Groq(api_key=api_key)
+
+def generate_answer_with_groq(client, query: str, relevant_chunks: List[Dict]) -> str:
+    try:
+        context = "\n\n".join([f"[Page {c['page_number']}]: {c['text']}" for c in relevant_chunks])
         prompt = f"""
-You are a tutor limited to the given book excerpts.
-Answer ONLY from the book. Always provide [Page].
-If not enough info, say: ❌ Insufficient evidence.
+        Based ONLY on the following context from a PDF document, answer the user's question.
+
+        Context:
+        {context}
 
-Context:
-{context}
+        Question: {query}
 
-Question: {question}
-Answer:
-"""
-        response = client.chat.completions.create(
+        Instructions:
+        - Answer ONLY using info from the context above
+        - If not enough info, reply: ❌ Insufficient evidence
+        - Always include page citations like [Page X]
+        """
+        chat = client.chat.completions.create(
             model="llama3-8b-8192",
-            messages=[{"role": "user", "content": prompt}]
+            messages=[
+                {"role": "system", "content": "You are a helpful tutor AI."},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.1,
+            max_tokens=500
         )
-        st.write(response.choices[0].message.content)
-        st.caption(f"📑 Sources: {', '.join(sources)}")
-    else:
-        st.error("❌ Insufficient evidence.")
+        return chat.choices[0].message.content
+    except Exception as e:
+        return f"Error generating answer: {e}"
+
+# Main answer pipeline
+def generate_answer(query: str, relevant_chunks: List[Dict]) -> str:
+    if not relevant_chunks:
+        return "❌ Insufficient evidence"
+    client = setup_groq()
+    if client:
+        return generate_answer_with_groq(client, query, relevant_chunks)
+    return "❌ No LLM configured."
+
+# -----------------------------
+# STREAMLIT MAIN
+# -----------------------------
+def main():
+    st.set_page_config(page_title="PageMentor", layout="wide")
+
+    st.title("📚 PageMentor")
+
+    if "vector_db" not in st.session_state:
+        st.session_state.vector_db = None
+        st.session_state.embedding_model = load_embedding_model()
+
+    uploaded_file = st.file_uploader("Upload PDF", type="pdf")
+
+    if uploaded_file and st.button("🚀 Process PDF"):
+        pdf_result = extract_text_from_pdf(uploaded_file)
+        if pdf_result["success"]:
+            chunks = create_chunks(pdf_result["pages"])
+            st.session_state.vector_db = create_vector_database(chunks, st.session_state.embedding_model)
+            if st.session_state.vector_db:
+                st.success(f"✅ Processed {pdf_result['total_pages']} pages, {len(chunks)} chunks ready!")
+        else:
+            st.error(pdf_result["error"])
+
+    if st.session_state.vector_db:
+        query = st.text_input("Ask a question:")
+        if query and st.button("🔍 Get Answer"):
+            relevant_chunks = query_vector_database(st.session_state.vector_db, query, st.session_state.embedding_model)
+            answer = generate_answer(query, relevant_chunks)
+            st.markdown("### 🎯 Answer")
+            st.write(answer)
+
+if __name__ == "__main__":
+    main()