Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import re | |
| from pathlib import Path | |
| import streamlit as st | |
| from rag.ingestion_pipeline import process_doc, create_vector_store | |
| from rag.retrieval_pipeline import retrieval_pipeline, load_vector_store, chat_history | |
| import hashlib | |
| DB_DIR = Path("db") | |
| INDEX_PATH = DB_DIR / "index.json" | |
| st.title("RAG Assistant") | |
| def load_index(): | |
| """ Load the stored filename and his path. """ | |
| if INDEX_PATH.exists(): | |
| return json.loads(INDEX_PATH.read_text()) | |
| return {} | |
| def save_index(data): | |
| """ Save the filename mapping to disk """ | |
| # Persist the filename -> db_path mapping on disk. | |
| DB_DIR.mkdir(parents=True, exist_ok=True) | |
| INDEX_PATH.write_text(json.dumps(data, indent=2)) | |
| def file_hash(uploaded_file) -> str : | |
| """ Return the content of the file uploaded in a short MD5 hash """ | |
| hash_md5 = hashlib.md5() | |
| hash_md5.update(uploaded_file.getbuffer()) | |
| return hash_md5.hexdigest()[:10] | |
| def safe_name(file_name: str, uploaded_file) -> str: | |
| """ Convert The uploaded filename into a unique, safe and clean folder name""" | |
| # Convert the uploaded filename into a safe folder name. | |
| # This avoids spaces and special characters in the DB path. | |
| clean_name = re.sub(r"[^a-zA-Z0-9_-]", "_", Path(file_name).stem) | |
| file_name_hashed = file_hash(uploaded_file) | |
| return f"{clean_name}_{file_name_hashed}" | |
| def get_db_path(filename: str, uploaded_file) -> str: | |
| """ Return the local Chroma database path for a specific uploaded name.""" | |
| # Each filename gets its own Chroma folder. | |
| # This avoids deleting an open DB and causing readonly/locked errors. | |
| return str(DB_DIR / f"chroma_db_{safe_name(filename, uploaded_file)}") | |
| # Load the filename registry once at startup. | |
| if "file_index" not in st.session_state: | |
| st.session_state.file_index = load_index() | |
| if "vector_store" not in st.session_state: | |
| st.session_state.vector_store = None | |
| st.session_state.current_document = None | |
| uploaded = st.file_uploader("Upload a PDF or TXT file", type=["pdf", "txt"]) | |
| # If a different file is selected, clear the current in-memory store. | |
| if uploaded and uploaded.name != st.session_state.current_document: | |
| st.info("A new file was selected. Process it to load or create its document store.") | |
| st.session_state.vector_store = None | |
| # Process only when the button is clicked. | |
| if uploaded and st.button("Process Document"): | |
| file_name = uploaded.name | |
| registry_key = safe_name(file_name, uploaded) | |
| db_path = st.session_state.file_index.get(registry_key, get_db_path(file_name, uploaded)) | |
| # If this file was already processed before, reuse its existing vector store. | |
| if os.path.exists(db_path): | |
| st.session_state.vector_store = load_vector_store(db_path) | |
| st.session_state.current_document = file_name | |
| st.success("Existing vector store loaded for this file.") | |
| else: | |
| # New file: parse, chunk, and store it in a new Chroma folder. | |
| chunked_text = process_doc(uploaded, file_hash(uploaded)) | |
| if not chunked_text: | |
| st.warning("No text chunks were created from the uploaded file.") | |
| else: | |
| # Clear chat history because this is a different document. | |
| chat_history.clear() | |
| with st.spinner("Creating vector store..."): | |
| st.session_state.vector_store = create_vector_store( | |
| chunked_text, | |
| path=db_path, | |
| ) | |
| # Save the filename -> db_path mapping so we can reuse it later. | |
| st.session_state.file_index[registry_key] = db_path | |
| save_index(st.session_state.file_index) | |
| st.session_state.current_document = file_name | |
| st.success(f"Split into {len(chunked_text)} chunks") | |
| st.success("Document processed and saved.") | |
| if st.session_state.vector_store is None: | |
| st.info("Upload and process a document before asking questions.") | |
| question = st.text_input( | |
| "Ask a question about your document", | |
| disabled=st.session_state.vector_store is None, | |
| ) | |
| if question and st.session_state.vector_store: | |
| answer = retrieval_pipeline(question, st.session_state.vector_store) | |
| st.write(answer) | |