Upload 7 files

src/collections_setup.py

@@ -0,0 +1,107 @@
+import chromadb
+import streamlit as st
+import fitz 
+import os
+from chromadb.utils import embedding_functions
+from text_processing import lines_chunking, paragraphs_chunking
+
+
+def get_chroma_client():
+    """
+    Get an ephemeral ChromaDB client for session-based RAG.
+    Data is automatically deleted when user closes browser/session ends.
+    """
+    return chromadb.EphemeralClient()
+
+
+@st.cache_resource
+def initialize_chroma_client():
+    """
+    Initialize ChromaDB client and store in Streamlit's resource cache.
+    This ensures one client per Streamlit session.
+    """
+    return get_chroma_client()
+
+
+@st.cache_resource
+def initialize_chromadb(embedding_model):
+    """
+    Initialize ChromaDB client and embedding function.
+    Both are cached to avoid recreating on every rerun.
+    """
+    # Get the cached client
+    client = initialize_chroma_client()
+
+    # Initialize an embedding function (using a Sentence Transformer model)
+    embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name=embedding_model
+    )
+
+    return client, embedding_func
+
+
+def initialize_collection(client, embedding_func, collection_name):
+    """
+    Initialize a collection in ChromaDB.
+    """
+    collection = client.get_or_create_collection(
+        name=collection_name,
+        embedding_function=embedding_func,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+    return collection
+
+
+def update_collection(collection, files_to_add_to_collection):
+    """
+    Update collection with new uploaded files.
+    Returns updated collection and session state.
+    """
+    for file_to_add in files_to_add_to_collection:
+
+        current_file = next(
+            (file for file in st.session_state.get('uploaded_files_raw', []) 
+            if file.name == file_to_add),None)
+        
+        if current_file is None:
+            st.error(f"File '{file_to_add}' not found in uploaded files.")
+            continue  
+
+        # Read file content
+        try:
+            if current_file.type == "text/plain":  # Handling TXT files
+                file_text = current_file.getvalue().decode("utf-8")
+            elif current_file.type == "application/pdf":  # Handling PDFs
+                with fitz.open(stream=current_file.getvalue(), filetype="pdf") as pdf_document:
+                    file_text = "\n".join([page.get_text("text") for page in pdf_document])
+            else:
+                st.warning(f"Unsupported file type: {current_file.name} type:{current_file.type}")
+                continue
+
+            # Tokenize text into chunks
+            max_words = 200
+            chunks = lines_chunking(file_text, max_words=max_words)
+            
+            if not chunks:  # Skip if no chunks generated
+                st.warning(f"No content extracted from {current_file.name}")
+                continue
+
+            # Store chunks in the collection
+            filename = current_file.name
+            collection.add(
+                documents=chunks,
+                ids=[f"id{filename[:-4]}.{j}" for j in range(len(chunks))],
+                metadatas=[{"source": filename, "part": n} for n in range(len(chunks))],
+            )
+            
+            st.session_state.collections_files_name.append(filename)
+            st.success(f"Added {len(chunks)} chunks from {filename}")
+            
+        except Exception as e:
+            st.error(f"Error processing {current_file.name}: {str(e)}")
+            # Remove from session state if processing failed
+            st.session_state.uploaded_files_name.remove(filename)
+    
+    return collection
+    

src/mylogging.py

@@ -0,0 +1,45 @@
+import streamlit as st
+import logging
+import io
+
+
+def configure_logging():
+    """
+    Configure logging.
+    """
+    log_stream = io.StringIO()  
+    handler = logging.StreamHandler(log_stream)
+    handler.setFormatter(logging.Formatter('%(message)s'))
+    handler.setLevel(logging.WARNING) 
+    
+    logger = logging.getLogger()
+    logger.setLevel(logging.WARNING)
+    logger.addHandler(handler)
+    
+    return logger, log_stream
+
+
+def toggle_logging(level, logger):
+    """
+    Toggle logging level.
+    """
+    if level == 'DEBUG':
+        logger.setLevel(logging.DEBUG)
+    elif level == 'INFO':
+        logger.setLevel(logging.INFO)
+    elif level == 'WARNING':
+        logger.setLevel(logging.WARNING)
+    else:
+        logger.warning(f"Unknown logging level: {level}. Using WARNING as default.")
+        logger.setLevel(logging.WARNING)
+    for handler in logger.handlers:
+        handler.setLevel(logger.level)
+
+
+def display_logs(log_stream):
+    """
+    Display logs in the app
+    """
+    log_stream.seek(0)  
+    logs = log_stream.read()  
+    st.text(logs)  

src/run.py

@@ -0,0 +1,128 @@
+import streamlit as st
+import os
+from utils import load_background_image, apply_style, configure_page, breaks, file_uploader, initialise_session_state
+from mylogging import configure_logging, toggle_logging, display_logs
+from collections_setup import initialize_chromadb, initialize_collection, update_collection
+from runpod_setup import get_relevant_text, generate_answer, get_contextual_prompt
+
+if __name__ == "__main__":
+
+    configure_page()
+    apply_style()
+    load_background_image()
+    initialise_session_state()
+    breaks(2)
+    st.write(
+        """
+    Welcome to this Streamlit app that demonstrates Retrieval-Augmented Generation (RAG) using a **Mistral-7B model hosted on Runpod** and **ChromaDB** for retrieval.
+
+    With this app, you can:
+    - Upload multiple PDF or text files to build a contextual knowledge base,
+    - Ask custom questions based on your uploaded documents, and
+    - Generate informed responses using a lightweight, hosted LLM.
+
+    **Note:** All uploaded files and generated embeddings are stored **in memory only** and will be **lost when the app is closed or restarted**. No data is persisted between sessions.
+    """
+    )
+    
+    # # Disable Chroma telemetry
+    os.environ["CHROMA_TELEMETRY_ENABLED"] = "False"
+    
+    # Initialize logger
+    logger, log_stream = configure_logging()
+    st.markdown(
+        """
+        <style>
+        /* This targets the selectbox container */
+        div[data-baseweb="select"] {
+            max-width: 150px;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    ) 
+    st.divider()
+
+    # ---- Logging Setup ----
+    use_logging = False
+    if use_logging:
+        logging_level = st.selectbox("Select logging level", ['INFO', 'DEBUG', 'WARNING'], index=2)
+        toggle_logging(logging_level, logger)
+
+    # ---- Vector Store Setup ----
+    # Initialize ChromaDB and collection
+    EMBEDDING_MODEL = "all-MiniLM-L6-v2"  
+    client, embedding_func = initialize_chromadb(EMBEDDING_MODEL)
+    collection_name = "my_collection"
+    collection = initialize_collection(client, embedding_func, collection_name)
+
+    # Upload files
+    st.markdown(
+    '<h3>Upload Files</h3>',
+    unsafe_allow_html=True)
+    st.html("""
+    Uploaded files are processed to build a contextual knowledge base for the RAG model.<br>
+    When you submit a prompt, the model retrieves relevant information from these documents to generate responses.
+    """)
+    col1_, _, col2_ = st.columns([.4, .1, .5])
+    with col1_:
+        file_uploader()
+
+    # Get the current uploaded filenames
+    logger.debug(f"\n\t-- Currently uploaded files: {st.session_state.get('uploaded_files_name', 'None')}")
+
+
+    # Update collection with uploaded files
+    files_to_add_to_collection= [file_name for file_name in st.session_state.get("uploaded_files_name", []) if file_name not in st.session_state.get("collections_files_name", [])]
+    logger.debug(f"\n\t-- Files not in collection: {files_to_add_to_collection}")
+
+    if files_to_add_to_collection:
+        collection = update_collection(collection, files_to_add_to_collection)
+
+    # Update the session state
+    logger.debug(f"Collection count: {collection.count()}")
+    logger.debug(f"\n\t-- Collection data currently uploaded:")
+    data_head = collection.get(limit=5)
+    for i, (metadata, document) in enumerate(zip(data_head["metadatas"], data_head["documents"]), start=1):
+        logger.debug(f"Item {i}:")
+        logger.debug(f"Metadata: {metadata}")
+        logger.debug(f"Document: {document}")
+        logger.debug("-" * 40)
+
+    # ---- Response Generation ----
+    # Streamlit UI
+    st.divider()
+    col1, _, col2 = st.columns([.6, .01, 1])
+    with col1:
+        st.subheader("Enter your prompt")
+        query = st.text_area("", height=200)
+        generate_clicked = st.button("Generate Response")
+    if generate_clicked:
+        if query.strip():
+            # Get the number of available documents in ChromaDB
+            available_docs = collection.count()
+
+            if available_docs > 0:
+                # Ensure n_results doesn't exceed available_docs
+                n_results = min(2, available_docs)
+                relevant_text = get_relevant_text(collection, query=query, nresults=n_results)
+            else:
+                relevant_text = ""  # No documents available, so no additional context
+                st.warning("No knowledge base available. Generating response based only on the prompt.")
+
+            logger.debug("\n\t-- Relevant text retrieved:")
+            logger.debug(relevant_text)
+
+            with st.spinner("Generating response..."):
+                context_query = get_contextual_prompt(query, relevant_text)
+                response = generate_answer(context_query, max_tokens=200)
+
+            with col2:
+                st.subheader("Response:")
+                st.text_area("", value=response, height=200)
+        else:
+            logger.debug("No query provided; skipping relevant text retrieval.")
+            st.warning("Please enter a prompt.")
+
+    if use_logging:
+        display_logs(log_stream)

src/runpod_setup.py

@@ -0,0 +1,76 @@
+import requests
+import os
+from dotenv import load_dotenv
+from pathlib import Path
+
+# Load .env from project root
+load_dotenv(dotenv_path=Path(__file__).resolve().parents[1] / ".env")
+
+API_KEY = os.getenv("RUNPOD_API_KEY")
+ENDPOINT = os.getenv("RUNPOD_ENDPOINT")
+
+HEADERS = {
+    "Authorization": f"Bearer {API_KEY}",
+    "Content-Type": "application/json"
+}
+    
+
+def get_relevant_text(collection, query='', nresults=3, sim_th=None):
+    """
+    Get relevant text from a collection for a given query
+    """
+    query_result = collection.query(query_texts=query, n_results=nresults)
+    docs = query_result.get('documents')[0]
+    if sim_th is not None:
+        similarities = [1 - d for d in query_result.get("distances")[0]]
+        relevant_docs = [d for d, s in zip(docs, similarities) if s >= sim_th]
+        return ''.join(relevant_docs)
+    return ''.join([doc for doc in docs if doc is not None])
+
+
+def get_contextual_prompt(question, context):
+    """
+    Optimized prompt format for Mistral 7B 
+    """
+    # Option 1: Mistral Chat Template (Recommended)
+    contextual_prompt = f"""<s>[INST] You are a helpful assistant that answers questions based on the provided context. Use only the information given in the context to answer the question. If the context doesn't contain enough information, say so clearly.
+
+Context:
+{context}
+
+Question: {question} [/INST]"""
+    
+    return contextual_prompt
+
+
+def generate_answer(prompt, max_tokens=150, temperature=0.7, HEADERS=HEADERS, ENDPOINT=ENDPOINT):
+    """
+    Submit a prompt to the RunPod SYNC endpoint and get back a response string.
+    """
+    payload = {
+        "input": {
+            "prompt": prompt,
+            "max_tokens": max_tokens,
+            "temperature": temperature
+        }
+    }
+
+    try:
+        # Use /runsync instead of /run - immediate response!
+        response = requests.post(f"{ENDPOINT}/runsync", headers=HEADERS, json=payload, timeout=65)
+        response.raise_for_status()
+        result = response.json()
+        
+        print(f"[RunPod] Request completed successfully")
+        
+        if result.get("status") == "COMPLETED":
+            return result["output"]["response"]
+        else:
+            error_msg = result.get("error", "Unknown error")
+            raise RuntimeError(f"RunPod job failed: {error_msg}")
+            
+    except requests.exceptions.Timeout:
+        raise RuntimeError("Request timed out (>60s). Try reducing prompt length or max_tokens.")
+    except requests.exceptions.RequestException as e:
+        raise RuntimeError(f"RunPod API error: {e}")
+

src/text_processing.py

@@ -0,0 +1,92 @@
+import nltk
+from nltk.tokenize import sent_tokenize
+nltk.download('punkt_tab')
+nltk.download("punkt")  
+
+
+def paragraphs_chunking(text, max_words=200, max_sentence_words=50):
+    """
+    Splits text into structured chunks, preserving paragraph integrity and avoiding unnatural breaks.
+    - Uses paragraph-based splitting first.
+    - Splits long paragraphs into smaller chunks based on sentence boundaries.
+    """
+    # Split text into paragraphs first
+    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
+    
+    chunks = []
+    for para in paragraphs:
+        words = para.split()
+        
+        # If paragraph is within limit, keep as a single chunk
+        if len(words) <= max_words:
+            chunks.append(para)
+            continue
+        
+        # Sentence-based chunking for large paragraphs
+        sentences = sent_tokenize(para)
+        chunk, chunk_word_count = [], 0
+
+        for sentence in sentences:
+            sentence_word_count = len(sentence.split())
+            
+            # If adding this sentence keeps chunk within word limit, add it
+            if chunk_word_count + sentence_word_count <= max_words:
+                chunk.append(sentence)
+                chunk_word_count += sentence_word_count
+            else:
+                # Finalize current chunk and start a new one
+                chunks.append(" ".join(chunk))
+                chunk = [sentence]
+                chunk_word_count = sentence_word_count
+
+        # Append any remaining chunk
+        if chunk:
+            chunks.append(" ".join(chunk))
+
+    return chunks
+
+
+def lines_chunking(text, max_words=200):
+    """
+    Splits text into structured chunks, preserving paragraph integrity and avoiding unnatural breaks.
+    - Uses paragraph-based splitting first.
+    - Splits long paragraphs into smaller chunks based on sentence boundaries.
+    """
+    # Split text into lines
+    lines = text.splitlines()
+
+    # Group lines into paragraphs
+    paragraphs = []
+    current_paragraph = []
+    for line in lines:
+        if line.strip():  
+            current_paragraph.append(line.strip())
+        else:  # Empty line indicates end of paragraph
+            if current_paragraph:
+                paragraphs.append(" ".join(current_paragraph))
+                current_paragraph = []
+    if current_paragraph: 
+        paragraphs.append(" ".join(current_paragraph))
+
+    # Process paragraphs
+    chunks = []
+    for para in paragraphs:
+        words = para.split()
+        if len(words) <= max_words:
+            chunks.append(para)
+        else:
+            sentences = sent_tokenize(para)
+            chunk, chunk_word_count = [], 0
+            for sentence in sentences:
+                sentence_word_count = len(sentence.split())
+                if chunk_word_count + sentence_word_count <= max_words:
+                    chunk.append(sentence)
+                    chunk_word_count += sentence_word_count
+                else:
+                    chunks.append(" ".join(chunk))
+                    chunk = [sentence]
+                    chunk_word_count = sentence_word_count
+            if chunk:
+                chunks.append(" ".join(chunk))
+
+    return chunks

src/utils.py

@@ -0,0 +1,181 @@
+import streamlit as st
+import os
+import sqlite3
+import base64
+
+
+DEFAULT_SESSION_STATE = {
+    # PDF Upload
+    'uploaded_files_name': [],
+    'collections_files_name': [],
+    'uploaded_files_raw': [],
+}
+
+
+def configure_page() -> None:
+    """
+    Configures the Streamlit page.
+    """
+    st.set_page_config(page_title="myRAG", 
+                       layout="wide", 
+                       page_icon=":rocket:")
+
+
+def apply_style():
+    st.markdown("""
+    <style>
+    @import url('https://fonts.googleapis.com/css2?family=Work+Sans:wght@400;700&display=swap');
+    html, body, .stApp,
+    .css-1v3fvcr, .css-ffhzg2, .css-1d391kg,
+    div[data-testid="stMarkdownContainer"],
+    div[data-testid="stText"],
+    div[data-testid="stTextInput"],
+    div[data-testid="stSelectbox"],
+    div[data-testid="stCheckbox"],
+    div[data-testid="stSlider"],
+    label, input, textarea, button, select,
+    .stButton, .stTextInput > div, .stMarkdown, .stCaption,
+    .streamlit-expanderHeader, .st-expander > div,
+    h1, h2, h3, h4, h5, h6,
+    .stMarkdown h1, .stMarkdown h2, .stMarkdown h3 {
+        font-family: 'Work Sans', sans-serif !important;
+    }
+    /* Ensure bold text uses the correct font weight */
+    strong, b, .stMarkdown strong, .stMarkdown b {
+        font-family: 'Work Sans', sans-serif !important;
+        font-weight: 700 !important;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
+
+def breaks(n=1):
+    """
+    Creates a line break.
+    """
+    if n == 1:
+        st.markdown("<br>",unsafe_allow_html=True)
+    elif n == 2:
+        st.markdown("<br><br>",unsafe_allow_html=True)
+    elif n == 3:
+        st.markdown("<br><br><br>",unsafe_allow_html=True)
+    else:
+        st.markdown("<br><br><br><br>",unsafe_allow_html=True)
+
+
+def get_base64_encoded_image(image_path):
+    """
+    Reads an image file and encodes it to Base64.
+    """
+    with open(image_path, "rb") as img_file:
+        return base64.b64encode(img_file.read()).decode()
+
+
+def load_background_image():
+    """
+    Loads and displays a background image with an overlaid title.
+    """
+
+    possible_paths = [
+     "../images/image6.jpg",      # Local development (from src/ folder)
+     "images/image6.jpg",         # Docker container (from /app)
+    ]
+    
+    image_path = None
+    for path in possible_paths:
+        if os.path.exists(path):
+            image_path = path
+            break
+    
+    if not image_path:
+        st.error("Could not find image6.jpg in any expected location")
+        return
+    
+    base64_image = get_base64_encoded_image(image_path)
+    
+    # Inject CSS for the background and title overlay
+    st.markdown(
+        f"""
+        <style>
+        /* Background container with image */
+        .bg-container {{
+            position: relative;
+            background-image: url("data:image/png;base64,{base64_image}");
+            background-size: container;
+            background-position: center;
+            height: 150px;  /* Adjust the height of the background */
+            width: 100%;
+            margin: 0 auto;
+            filter: contrast(110%) brightness(210%); /* Dim the brightness of the image */
+            border-radius: 100px;  /* Makes the container's corners rounded */
+            overflow: hidden;  
+        }}
+
+        /* Overlay for dimming effect */
+        .bg-container::after {{
+            content: '';
+            position: absolute;
+            top: ;
+            left: 0;
+            width: 100%;
+            height: 100%;
+            background-color: rgba(20, 10, 20, 0.44); /* Semi-transparent black overlay */
+            z-index: 1; /* Ensure the overlay is above the image */
+        }}
+
+        /* Overlay title styling */
+        .overlay-title {{
+            position: absolute;
+            top: 50%;
+            left: 50%;
+            transform: translate(-50%, -50%);
+            color: black;   /* Title color */
+            font-size: 50px;
+            font-weight: bold;
+            text-shadow: 1px 1px 3px rgba(255, 255, 255, .0); /* Shadow for better visibility */
+            text-align: center;
+            z-index: 2; /* Ensure the title is above the overlay */
+        }}
+        </style>
+        """,
+        unsafe_allow_html=True
+    )
+
+    # Create the background container with an overlaid title
+    st.markdown(
+        """
+        <div class="bg-container">
+            <div class="overlay-title">Mistral-RAG</div>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+
+def initialise_session_state():
+    """
+    Initializes the session state variables if not already set.
+    """
+    for key, default_val in DEFAULT_SESSION_STATE.items():
+        if key not in st.session_state:
+            st.session_state[key] = default_val
+
+
+def file_uploader():
+    uploaded_files = st.file_uploader(
+        "",
+        type=["txt", "pdf"], 
+        accept_multiple_files=True)  
+    
+    if uploaded_files:  # Check if list is not empty
+        for file in uploaded_files:  # Process each file
+            if file.name not in st.session_state.uploaded_files_name:
+                # Append to session state lists safely
+                st.session_state.uploaded_files_name.append(file.name)
+                st.session_state.uploaded_files_raw.append(file)
+                st.success(f"Added new file: {file.name}")
+   
+    else:
+        st.info("Please upload a PDF file to proceed.")
+
+