sondosomarr/DevOps_RAG_Assistant

1import streamlit as st2import os3import base644from src.ingestion import ingest_documents5from src.generation import ask_question6 7# Sidebar starts collapsed8st.set_page_config(page_title="RAGOps", page_icon="⚙️", initial_sidebar_state="collapsed")9 10DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")11 12def list_uploaded_pdfs():13    if not os.path.exists(DATA_DIR):14        return []15    return [f for f in os.listdir(DATA_DIR) if f.endswith('.pdf')]16 17def scroll_to_bottom():18    """Smoothly scroll the chat container to the bottom."""19    st.components.v1.html(20        """21        <script>22        function scrollChat() {23            const chatContainer = document.querySelector('.stChatMessageContainer');24            if (chatContainer) {25                chatContainer.style.scrollBehavior = 'smooth';26                chatContainer.scrollTop = chatContainer.scrollHeight;27            }28        }29        setTimeout(scrollChat, 150);30        </script>31        """,32        height=0,33    )34 35def main():36    # ----- CUSTOM CSS (Blue theme + UI polish) -----37    st.markdown(38        """39        <style>40        :root {41            --devops-blue: #007bff;42            --devops-blue-hover: #0056b3;43            --light-bg: #f8f9fa;44            --border-radius: 12px;45        }46        /* Multiselect tags – blue */47        .stMultiSelect [data-baseweb="tag"] {48            background-color: var(--devops-blue) !important;49            color: white !important;50            border-radius: 16px !important;51            padding: 4px 10px !important;52            font-weight: 500;53        }54        .stMultiSelect [data-baseweb="tag"] button {55            color: white !important;56            opacity: 0.8;57        }58        .stMultiSelect [data-baseweb="tag"] button:hover {59            opacity: 1;60        }61        /* Chat messages */62        .stChatMessage {63            border-radius: 18px !important;64            padding: 10px 16px !important;65            max-width: 75%;66            margin-bottom: 12px;67            line-height: 1.5;68            box-shadow: 0 2px 5px rgba(0,0,0,0.05);69        }70        [data-testid="chat-message-user"] {71            background-color: #e6f2ff;  /* light blue */72            align-self: flex-end;73            margin-left: auto;74            border-bottom-right-radius: 4px !important;75        }76        [data-testid="chat-message-assistant"] {77            background-color: #ffffff;78            border: 1px solid #e9ecef;79            align-self: flex-start;80            border-bottom-left-radius: 4px !important;81        }82        /* Smooth scrolling container */83        .stChatMessageContainer {84            scroll-behavior: smooth;85        }86        /* Sidebar cards */87        .sidebar-section {88            background-color: white;89            border-radius: 12px;90            padding: 16px;91            margin-bottom: 20px;92            box-shadow: 0 2px 8px rgba(0,0,0,0.03);93            border: 1px solid #f0f0f0;94        }95        /* Buttons */96        .stButton button {97            background-color: var(--devops-blue);98            color: white;99            border: none;100            border-radius: 8px;101            padding: 0.5rem 1rem;102            font-weight: 500;103            transition: all 0.2s ease;104        }105        .stButton button:hover {106            background-color: var(--devops-blue-hover);107            box-shadow: 0 4px 10px rgba(0,123,255,0.2);108        }109        /* File uploader button */110        .stFileUploader button {111            background-color: white;112            color: var(--devops-blue);113            border: 1px solid var(--devops-blue);114        }115        .stFileUploader button:hover {116            background-color: #e6f2ff;117        }118        /* Input focus */119        .stChatInput textarea:focus,120        .stMultiSelect [data-baseweb="select"]:focus-within {121            border-color: var(--devops-blue) !important;122            box-shadow: 0 0 0 1px var(--devops-blue) !important;123        }124        /* Info/warning boxes */125        .stAlert {126            border-left-color: var(--devops-blue) !important;127        }128        /* Blue placeholder text */129        div.stChatInput textarea::placeholder {130            color: #007bff !important;131            opacity: 0.8;132        }133        /* Chat input border on focus – blue */134        div.stChatInput textarea:focus {135            border-color: #007bff !important;136            box-shadow: 0 0 0 1px #007bff !important;137            outline: none !important;138        }139        /* Footer */140        .footer {141            text-align: center;142            margin-top: 30px;143            padding: 12px;144            font-size: 0.85rem;145            color: #6c757d;146            border-top: 1px solid #e9ecef;147        }148        /* Hero container – centers image both horizontally and vertically */149        .hero-container {150            display: flex;151            flex-direction: column;152            align-items: center;153            justify-content: center;154            min-height: 40vh;  /* Adjust this value to control vertical centering */155            width: 100%;156            margin-bottom: 1rem;157        }158        .hero-container img {159            max-width: 100%;160            height: auto;161        }162        </style>163        """,164        unsafe_allow_html=True,165    )166 167    # ----- HERO SECTION with centered image (base64 to avoid MediaFileStorageError) -----168    logo_path = "assets/logo.png"  # <-- place your image here169 170    if os.path.exists(logo_path):171        with open(logo_path, "rb") as img_file:172            img_base64 = base64.b64encode(img_file.read()).decode()173 174        st.markdown(175            f"""176            <div class="hero-container">177                <img src="data:image/jpeg;base64,{img_base64}" width="350">178            </div>179            """,180            unsafe_allow_html=True,181        )182    else:183        # Fallback emoji if image is missing184        st.markdown(185            """186            <div class="hero-container">187                <div style="font-size:60px;">🖼️</div>188            </div>189            """,190            unsafe_allow_html=True,191        )192 193    # Subtitle directly below hero section194    st.markdown("  ")195 196    # Initialize chat history197    if "messages" not in st.session_state:198        st.session_state.messages = []199 200    # ----- SIDEBAR (with cards) -----201    with st.sidebar:202        # Document Management card203        st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)204        st.markdown("#### 📁 **Document Management**")205        st.caption("Upload your DevOps PDFs (max 200MB each).")206 207        uploaded_files = st.file_uploader(208            "Choose PDF files",209            type="pdf",210            accept_multiple_files=True,211            help="Drag and drop or click to select multiple PDFs."212        )213 214        if st.button("⚙️ Process Documents", use_container_width=True):215            if uploaded_files:216                if not os.path.exists(DATA_DIR):217                    os.makedirs(DATA_DIR)218                with st.spinner("Saving files and building vector index..."):219                    for file in uploaded_files:220                        file_path = os.path.join(DATA_DIR, file.name)221                        with open(file_path, "wb") as f:222                            f.write(file.getbuffer())223                    try:224                        ingest_documents()225                        st.success("✅ Documents processed successfully!")226                    except Exception as e:227                        st.error(f"Error during processing: {e}")228            else:229                st.warning("Please upload some PDFs first.")230        st.markdown('</div>', unsafe_allow_html=True)231 232        # Active Documents card233        st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)234        st.markdown("#### 📄 **Active Documents**")235        st.caption("Select which documents to query against. Unchecked ones will be ignored.")236 237        available_docs = list_uploaded_pdfs()238        active_docs = st.multiselect(239            "Documents to search",240            options=available_docs,241            default=available_docs,242            help="Only the selected PDFs will be used when answering your questions."243        )244 245        if active_docs:246            st.info(f"🔍 Searching **{len(active_docs)}** document(s).")247        else:248            st.warning("⚠️ No documents selected – please select at least one.")249        st.markdown('</div>', unsafe_allow_html=True)250 251    # ----- CHAT HISTORY -----252    for message in st.session_state.messages:253        with st.chat_message(message["role"]):254            st.markdown(message["content"])255            if "sources" in message and message["sources"]:256                with st.expander("📚 Retrieved Sources"):257                    for src in message["sources"]:258                        st.markdown(259                            f"""260                            <div style="background-color: #f8f9fa; border-radius: 8px; padding: 10px; margin-bottom: 8px; border-left: 4px solid #007bff;">261                                <b>{src.get('doc_title', 'Unknown')}</b>  ·  Page {src.get('page', '?')}262                            </div>263                            """,264                            unsafe_allow_html=True,265                        )266 267    # ----- CHAT INPUT -----268    if prompt := st.chat_input("What is your DevOps question?"):269        if not active_docs:270            st.warning("Please select at least one Active Document from the sidebar before asking a question.")271            st.stop()272 273        # Add user message274        st.session_state.messages.append({"role": "user", "content": prompt})275        with st.chat_message("user"):276            st.markdown(prompt)277        scroll_to_bottom()278 279        # Generate assistant response280        with st.chat_message("assistant"):281            with st.spinner("💭 Thinking..."):282                try:283                    response, sources = ask_question(prompt, active_docs)284 285                    st.markdown(response)286 287                    # Store response288                    st.session_state.messages.append({289                        "role": "assistant",290                        "content": response,291                        "sources": sources292                    })293 294                    if sources:295                        with st.expander("📚 Retrieved Sources"):296                            for src in sources:297                                st.markdown(298                                    f"""299                                    <div style="background-color: #f8f9fa; border-radius: 8px; padding: 10px; margin-bottom: 8px; border-left: 4px solid #007bff;">300                                        <b>{src.get('doc_title', 'Unknown')}</b>  ·  Page {src.get('page', '?')}301                                    </div>302                                    """,303                                    unsafe_allow_html=True,304                                )305 306                    scroll_to_bottom()307 308                except Exception as e:309                    st.error(f"Error generating response: {e}")310 311    # ----- FOOTER -----312    st.markdown(313        '<div class="footer">Powered by <b>Qwen 2.5 7B</b> · Built with Streamlit · <a href="https://github.com/your-repo" target="_blank">GitHub</a></div>',314        unsafe_allow_html=True,315    )316 317if __name__ == "__main__":318    main()