Upload 3 files

src/app_updated.py

@@ -0,0 +1,202 @@
+import streamlit as st
+import os
+import dotenv
+import uuid
+
+# Patch sqlite3 for Streamlit Cloud compatibility
+if os.name == 'posix':
+    __import__('pysqlite3')
+    import sys
+    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
+
+from langchain.schema import HumanMessage, AIMessage
+from langchain_groq import ChatGroq
+
+from rag_methods import (
+    load_doc_to_db,
+    load_url_to_db,
+    stream_llm_response,
+    stream_llm_rag_response,
+)
+
+dotenv.load_dotenv()
+
+# --- Custom CSS Styling ---
+def apply_custom_css():
+    st.markdown("""
+    <style>
+    .main .block-container {
+        padding-top: 2rem;
+        padding-bottom: 2rem;
+    }
+    h1, h2, h3, h4 {
+        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+        font-weight: 600;
+    }
+    .app-title {
+        text-align: center;
+        color: #4361ee;
+        font-size: 2.2rem;
+        font-weight: 700;
+        margin-bottom: 1.5rem;
+        padding: 1rem;
+        border-radius: 10px;
+        background: linear-gradient(90deg, rgba(67, 97, 238, 0.1), rgba(58, 12, 163, 0.1));
+        text-shadow: 0px 0px 2px rgba(0,0,0,0.1);
+    }
+    .chat-container {
+        border-radius: 10px;
+        padding: 10px;
+        margin-bottom: 1rem;
+    }
+    .message-container {
+        padding: 0.8rem;
+        margin-bottom: 0.8rem;
+        border-radius: 8px;
+    }
+    .user-message {
+        background-color: rgba(67, 97, 238, 0.15);
+        border-left: 4px solid #4361ee;
+    }
+    .assistant-message {
+        background-color: rgba(58, 12, 163, 0.1);
+        border-left: 4px solid #3a0ca3;
+    }
+    .document-list {
+        background-color: rgba(67, 97, 238, 0.05);
+        border-radius: 8px;
+        padding: 0.7rem;
+    }
+    .upload-container {
+        border: 2px dashed rgba(67, 97, 238, 0.5);
+        border-radius: 10px;
+        padding: 1rem;
+        margin-bottom: 1rem;
+        text-align: center;
+    }
+    .status-indicator {
+        font-size: 0.85rem;
+        font-weight: 600;
+        padding: 0.3rem 0.7rem;
+        border-radius: 20px;
+        display: inline-block;
+        margin-bottom: 0.5rem;
+    }
+    .status-active {
+        background-color: rgba(46, 196, 182, 0.2);
+        color: #2EC4B6;
+    }
+    .status-inactive {
+        background-color: rgba(231, 111, 81, 0.2);
+        color: #E76F51;
+    }
+    @media screen and (max-width: 768px) {
+        .app-title {
+            font-size: 1.8rem;
+            padding: 0.7rem;
+        }
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
+# --- Page Setup ---
+st.set_page_config(
+    page_title="RAG-Xpert: An Enhanced RAG Framework",
+    page_icon="📚",
+    layout="centered",
+    initial_sidebar_state="expanded"
+)
+
+apply_custom_css()
+
+st.markdown('<h1 class="app-title">📚 RAG-Xpert: An Enhanced Retrieval-Augmented Generation Framework 🤖</h1>', unsafe_allow_html=True)
+
+# --- Session Initialization ---
+if "session_id" not in st.session_state:
+    st.session_state.session_id = str(uuid.uuid4())
+if "rag_sources" not in st.session_state:
+    st.session_state.rag_sources = []
+if "messages" not in st.session_state:
+    st.session_state.messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there! How can I assist you today?"}
+    ]
+
+# --- Sidebar ---
+with st.sidebar:
+    st.markdown("""
+        <div style="
+            text-align: center; 
+            padding: 1rem 0; 
+            margin-bottom: 1.5rem; 
+            background: linear-gradient(to right, #4361ee22, #3a0ca322); 
+            border-radius: 10px;">
+            <div style="font-size: 0.85rem; color: #888;">Developed By</div>
+            <div style="font-size: 1.2rem; font-weight: 700; color: #4361ee;">Uditanshu Pandey</div>
+        </div>
+    """, unsafe_allow_html=True)
+
+    is_vector_db_loaded = "vector_db" in st.session_state and st.session_state.vector_db is not None
+    rag_status = st.toggle("Enable Knowledge Enhancement (RAG)", value=is_vector_db_loaded, key="use_rag", disabled=not is_vector_db_loaded)
+
+    if rag_status:
+        st.markdown('<div class="status-indicator status-active">RAG Mode: Active ✓</div>', unsafe_allow_html=True)
+    else:
+        st.markdown('<div class="status-indicator status-inactive">RAG Mode: Inactive ✗</div>', unsafe_allow_html=True)
+
+    st.toggle("Show Retrieved Context", key="debug_mode", value=False)
+    st.button("🧹 Clear Chat History", on_click=lambda: st.session_state.messages.clear(), type="primary")
+
+    st.markdown("<h3 style='text-align: center; color: #4361ee; margin-top: 1.5rem;'>📚 Knowledge Sources</h3>", unsafe_allow_html=True)
+    st.markdown('<div class="upload-container">', unsafe_allow_html=True)
+    st.file_uploader("📄 Upload Documents", type=["pdf", "txt", "docx", "md"], accept_multiple_files=True, on_change=load_doc_to_db, key="rag_docs")
+    st.markdown('</div>', unsafe_allow_html=True)
+
+    st.text_input("🌐 Add Webpage URL", placeholder="https://example.com", on_change=load_url_to_db, key="rag_url")
+
+    doc_count = len(st.session_state.rag_sources) if is_vector_db_loaded else 0
+    with st.expander(f"📑 Knowledge Base ({doc_count} sources)"):
+        if doc_count:
+            st.markdown('<div class="document-list">', unsafe_allow_html=True)
+            for i, source in enumerate(st.session_state.rag_sources):
+                st.markdown(f"**{i+1}.** {source}")
+            st.markdown('</div>', unsafe_allow_html=True)
+        else:
+            st.info("No documents added yet. Upload files or add URLs to enhance the assistant's knowledge.")
+
+# --- Initialize LLM ---
+llm_stream = ChatGroq(
+    model_name="meta-llama/llama-4-scout-17b-16e-instruct",
+    api_key=os.getenv("GROQ_API_KEY"),
+    temperature=0.4,
+    max_tokens=1024,
+)
+
+# --- Chat Display ---
+st.markdown('<div class="chat-container">', unsafe_allow_html=True)
+for message in st.session_state.messages:
+    avatar = "👤" if message["role"] == "user" else "🤖"
+    css_class = "user-message" if message["role"] == "user" else "assistant-message"
+    with st.chat_message(message["role"], avatar=avatar):
+        st.markdown(f'<div class="message-container {css_class}">{message["content"]}</div>', unsafe_allow_html=True)
+st.markdown('</div>', unsafe_allow_html=True)
+
+# --- User Input Handling ---
+if prompt := st.chat_input("Ask me anything..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user", avatar="👤"):
+        st.markdown(f'<div class="message-container user-message">{prompt}</div>', unsafe_allow_html=True)
+
+    with st.chat_message("assistant", avatar="🤖"):
+        thinking_placeholder = st.empty()
+        thinking_placeholder.info("Thinking... Please wait a moment.")
+        messages = [
+            HumanMessage(content=m["content"]) if m["role"] == "user" else AIMessage(content=m["content"])
+            for m in st.session_state.messages
+        ]
+        if not st.session_state.use_rag:
+            thinking_placeholder.empty()
+            st.write_stream(stream_llm_response(llm_stream, messages))
+        else:
+            thinking_placeholder.info("Searching knowledge base... Please wait.")
+            st.write_stream(stream_llm_rag_response(llm_stream, messages))

src/rag_methods.py

@@ -0,0 +1,174 @@
+import os
+import dotenv
+from time import time
+import streamlit as st
+
+from langchain_community.document_loaders.text import TextLoader
+from langchain_community.document_loaders import (
+    WebBaseLoader, 
+    PyPDFLoader, 
+    Docx2txtLoader,
+)
+from langchain_community.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.chains import create_history_aware_retriever, create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+
+dotenv.load_dotenv()
+
+os.environ["USER_AGENT"] = "myagent"
+DB_DOCS_LIMIT = 10
+
+# Stream non-RAG LLM response
+def stream_llm_response(llm_stream, messages):
+    response_message = ""
+    for chunk in llm_stream.stream(messages):
+        response_message += chunk.content
+        yield chunk
+    st.session_state.messages.append({"role": "assistant", "content": response_message})
+
+# --- Document Loading and Indexing ---
+def load_doc_to_db():
+    if "rag_docs" in st.session_state and st.session_state.rag_docs:
+        docs = []
+        for doc_file in st.session_state.rag_docs:
+            if doc_file.name not in st.session_state.rag_sources:
+                if len(st.session_state.rag_sources) < DB_DOCS_LIMIT:
+                    os.makedirs("source_files", exist_ok=True)
+                    file_path = f"./source_files/{doc_file.name}"
+                    with open(file_path, "wb") as file:
+                        file.write(doc_file.read())
+                    try:
+                        if doc_file.type == "application/pdf":
+                            loader = PyPDFLoader(file_path)
+                        elif doc_file.name.endswith(".docx"):
+                            loader = Docx2txtLoader(file_path)
+                        elif doc_file.type in ["text/plain", "text/markdown"]:
+                            loader = TextLoader(file_path)
+                        else:
+                            st.warning(f"Unsupported document type: {doc_file.type}")
+                            continue
+                        docs.extend(loader.load())
+                        st.session_state.rag_sources.append(doc_file.name)
+                    except Exception as e:
+                        st.toast(f"Error loading document {doc_file.name}: {e}", icon="⚠️")
+                    finally:
+                        os.remove(file_path)
+                else:
+                    st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
+        if docs:
+            _split_and_load_docs(docs)
+            st.toast(f"Documents loaded successfully.", icon="✅")
+
+def load_url_to_db():
+    if "rag_url" in st.session_state and st.session_state.rag_url:
+        url = st.session_state.rag_url
+        docs = []
+        if url not in st.session_state.rag_sources:
+            if len(st.session_state.rag_sources) < DB_DOCS_LIMIT:
+                try:
+                    loader = WebBaseLoader(url)
+                    docs.extend(loader.load())
+                    st.session_state.rag_sources.append(url)
+                except Exception as e:
+                    st.error(f"Error loading from URL {url}: {e}")
+                if docs:
+                    _split_and_load_docs(docs)
+                    st.toast(f"Loaded content from URL: {url}", icon="✅")
+            else:
+                st.error(f"Max documents reached ({DB_DOCS_LIMIT}).")
+
+def initialize_vector_db(docs):
+    # Initialize HuggingFace embeddings
+    embedding = HuggingFaceEmbeddings(
+        model_name="BAAI/bge-large-en-v1.5",
+        model_kwargs={'device': 'cpu'},
+        encode_kwargs={'normalize_embeddings': False}
+    )
+
+    # Shared persistent directory for long-term storage
+    persist_dir = "./chroma_persistent_db"
+    collection_name = "persistent_collection"
+
+    # Create the persistent Chroma vector store
+    vector_db = Chroma.from_documents(
+        documents=docs,
+        embedding=embedding,
+        persist_directory=persist_dir,
+        collection_name=collection_name
+    )
+
+    # Persist to disk
+    vector_db.persist()
+
+    return vector_db
+
+
+def _split_and_load_docs(docs):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+    )
+
+    chunks = text_splitter.split_documents(docs)
+
+    if "vector_db" not in st.session_state:
+        st.session_state.vector_db = initialize_vector_db(chunks)
+    else:
+        st.session_state.vector_db.add_documents(chunks)
+        st.session_state.vector_db.persist()  # Save changes
+
+# --- RAG Chain ---
+
+def _get_context_retriever_chain(vector_db, llm):
+    retriever = vector_db.as_retriever()
+    prompt = ChatPromptTemplate.from_messages([
+        MessagesPlaceholder(variable_name="messages"),
+        ("user", "{input}"),
+        ("user", "Given the above conversation, generate a search query to find relevant information.")
+    ])
+    return create_history_aware_retriever(llm, retriever, prompt)
+
+def get_conversational_rag_chain(llm):
+    retriever_chain = _get_context_retriever_chain(st.session_state.vector_db, llm)
+    prompt = ChatPromptTemplate.from_messages([
+        ("system",
+         """You are a helpful assistant answering the user's queries using the provided context if available.\n
+         {context}"""),
+        MessagesPlaceholder(variable_name="messages"),
+        ("user", "{input}")
+    ])
+    stuff_documents_chain = create_stuff_documents_chain(llm, prompt)
+    return create_retrieval_chain(retriever_chain, stuff_documents_chain)
+
+# Stream RAG LLM response
+def stream_llm_rag_response(llm_stream, messages):
+    rag_chain = get_conversational_rag_chain(llm_stream)
+
+    # Extract latest user input and prior messages
+    input_text = messages[-1].content
+    history = messages[:-1]
+
+    # --- DEBUG: Show context retrieved ---
+    if st.session_state.get("debug_mode"):
+        retriever = st.session_state.vector_db.as_retriever()
+        retrieved_docs = retriever.get_relevant_documents(input_text)
+        st.markdown("### 🔍 Retrieved Context (Debug Mode)")
+        for i, doc in enumerate(retrieved_docs):
+            st.markdown(f"**Chunk {i+1}:**\n```\n{doc.page_content.strip()}\n```")
+
+    response_message = "*(RAG Response)*\n"
+    response = rag_chain.stream({
+        "messages": history,
+        "input": input_text
+    })
+
+    for chunk in response:
+        if 'answer' in chunk:
+            response_message += chunk['answer']
+            yield chunk['answer']
+
+    st.session_state.messages.append({"role": "assistant", "content": response_message})
+