Spaces:
Configuration error
Configuration error
| import streamlit as st | |
| from rag.pdf_loader import load_pdf | |
| from rag.chunker import chunk_text | |
| from rag.embedder import embed_chunks, embed_query | |
| from rag.vector_store import build_index, search_index | |
| from rag.llm import ask_llm | |
| # ── Page configuration ──────────────────────────────────────── | |
| # This must be the very first Streamlit command in the script | |
| st.set_page_config( | |
| page_title = "RAG Assistant", | |
| page_icon = "📄", | |
| layout = "wide" | |
| ) | |
| # ── Initialize session state ────────────────────────────────── | |
| # session_state persists data between Streamlit reruns | |
| # We set default values here if they do not exist yet | |
| if "chunks" not in st.session_state: st.session_state.chunks = None | |
| if "index" not in st.session_state: st.session_state.index = None | |
| if "chat_history" not in st.session_state: st.session_state.chat_history = [] | |
| if "pdf_name" not in st.session_state: st.session_state.pdf_name = None | |
| # ── Sidebar ─────────────────────────────────────────────────── | |
| with st.sidebar: | |
| st.title("📄 RAG Assistant") | |
| st.markdown("Chat with any PDF using AI — fully local embeddings, no data sent except your question.") | |
| st.divider() | |
| # Groq API key input | |
| # type="password" hides the key so it is not visible on screen | |
| st.subheader("Step 1 — Enter API Key") | |
| api_key = st.text_input( | |
| "Groq API Key", | |
| type = "password", | |
| placeholder = "gsk_..." | |
| ) | |
| st.caption("Get your free key at console.groq.com") | |
| st.divider() | |
| # PDF upload section | |
| st.subheader("Step 2 — Upload a PDF") | |
| uploaded_file = st.file_uploader( | |
| "Choose a PDF file", | |
| type = ["pdf"] | |
| ) | |
| # Process the PDF when user uploads one | |
| if uploaded_file is not None: | |
| # Only re-process if it is a new PDF (not the same one again) | |
| if st.session_state.pdf_name != uploaded_file.name: | |
| with st.spinner("Reading and indexing your PDF..."): | |
| # Save the uploaded file temporarily to disk | |
| # Streamlit gives us the file as bytes, we write it to disk | |
| # so our load_pdf() function can open it with fitz | |
| temp_path = f"/tmp/{uploaded_file.name}" | |
| with open(temp_path, "wb") as f: | |
| f.write(uploaded_file.read()) | |
| # Run the full RAG pipeline on the PDF | |
| pages = load_pdf(temp_path) # extract text | |
| chunks = chunk_text(pages, chunk_size=500, overlap=50) # split into chunks | |
| chunks = embed_chunks(chunks) # convert to vectors | |
| index = build_index(chunks) # build FAISS index | |
| # Save everything to session state so it persists | |
| st.session_state.chunks = chunks | |
| st.session_state.index = index | |
| st.session_state.pdf_name = uploaded_file.name | |
| # Clear old chat history when a new PDF is loaded | |
| st.session_state.chat_history = [] | |
| st.success(f"Ready! Loaded {len(pages)} pages, {len(chunks)} chunks.") | |
| # Show currently loaded PDF name | |
| if st.session_state.pdf_name: | |
| st.info(f"Current PDF: {st.session_state.pdf_name}") | |
| st.divider() | |
| st.caption("Built with PyMuPDF · SentenceTransformers · FAISS · Groq · Streamlit") | |
| # ── Main area ───────────────────────────────────────────────── | |
| st.title("Chat with your PDF") | |
| # Show a welcome message if no PDF is loaded yet | |
| if st.session_state.chunks is None: | |
| st.info("Upload a PDF in the sidebar to get started.") | |
| else: | |
| # ── Chat history display ────────────────────────────────── | |
| # Loop through all previous messages and display them | |
| for message in st.session_state.chat_history: | |
| if message["role"] == "user": | |
| # Display user message on the right | |
| with st.chat_message("user"): | |
| st.write(message["content"]) | |
| elif message["role"] == "assistant": | |
| # Display assistant answer on the left | |
| with st.chat_message("assistant"): | |
| st.write(message["content"]) | |
| # Show the source pages used to generate this answer | |
| if "sources" in message: | |
| with st.expander("View sources"): | |
| for source in message["sources"]: | |
| st.markdown(f"**Page {source['page_num']}**") | |
| st.caption(source["text"][:300] + "...") | |
| st.divider() | |
| # ── Question input ──────────────────────────────────────── | |
| # st.chat_input stays fixed at the bottom of the page | |
| query = st.chat_input("Ask a question about your PDF...") | |
| if query: | |
| # Check that API key has been entered | |
| if not api_key: | |
| st.warning("Please enter your Groq API key in the sidebar.") | |
| else: | |
| # Show the user message immediately | |
| with st.chat_message("user"): | |
| st.write(query) | |
| # Save user message to chat history | |
| st.session_state.chat_history.append({ | |
| "role": "user", | |
| "content": query | |
| }) | |
| # Generate the answer | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| # Step 1 — convert question to vector | |
| query_embedding = embed_query(query) | |
| # Step 2 — find top 3 most relevant chunks | |
| relevant_chunks = search_index( | |
| st.session_state.index, | |
| st.session_state.chunks, | |
| query_embedding, | |
| top_k = 3 | |
| ) | |
| # Step 3 — get answer from LLM | |
| answer = ask_llm(query, relevant_chunks, api_key) | |
| # Display the answer | |
| st.write(answer) | |
| # Display the sources in an expandable section | |
| with st.expander("View sources"): | |
| for source in relevant_chunks: | |
| st.markdown(f"**Page {source['page_num']}**") | |
| st.caption(source["text"][:300] + "...") | |
| st.divider() | |
| # Save assistant answer + sources to chat history | |
| st.session_state.chat_history.append({ | |
| "role": "assistant", | |
| "content": answer, | |
| "sources": relevant_chunks | |
| }) | |