Update app.py

app.py

@@ -1,27 +1,51 @@
+import streamlit as st
+import os
+import requests
+import tempfile
+import chromadb  
+from langchain.document_loaders import PDFPlumberLoader
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_experimental.text_splitter import SemanticChunker
+from langchain_chroma import Chroma
+from langchain.chains import LLMChain, SequentialChain
+from langchain.prompts import PromptTemplate
+from langchain_groq import ChatGroq
+from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
 
+# ----------------- Streamlit UI Setup -----------------
+st.set_page_config(page_title="Blah", layout="wide")
+st.image("https://huggingface.co/front/assets/huggingface_logo-noborder.svg", width=150)  
+    st.title("Blah-1")
 
-# Set API Keys
-os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
-
-# Load LLM models
-llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
-rag_llm = ChatGroq(model="mixtral-8x7b-32768")
 
-llm_judge.verbose = True
-rag_llm.verbose = True
+# ----------------- API Keys -----------------
+os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-# Clear ChromaDB cache to fix tenant issue
+# ----------------- Clear ChromaDB Cache -----------------
 chromadb.api.client.SharedSystemClient.clear_system_cache()
 
-st.title("Blah - 1")
-
+# ----------------- Initialize Session State -----------------
+if "pdf_loaded" not in st.session_state:
+    st.session_state.pdf_loaded = False
+if "chunked" not in st.session_state:
+    st.session_state.chunked = False
+if "vector_created" not in st.session_state:
+    st.session_state.vector_created = False
+if "processed_chunks" not in st.session_state:
+    st.session_state.processed_chunks = None
+if "vector_store" not in st.session_state:
+    st.session_state.vector_store = None
 
+# ----------------- Load Models -----------------
+llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
+rag_llm = ChatGroq(model="mixtral-8x7b-32768")
 
-# Step 1: Choose PDF Source
-pdf_source = st.radio("Upload or provide a link to a PDF:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
+# ----------------- PDF Selection (Upload or URL) -----------------
+st.sidebar.subheader("📂 PDF Selection")
+pdf_source = st.radio("Choose a PDF source:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
 
 if pdf_source == "Upload a PDF file":
-    uploaded_file = st.file_uploader("Upload your PDF file", type="pdf")
+    uploaded_file = st.sidebar.file_uploader("Upload your PDF file", type=["pdf"])
     if uploaded_file:
         st.session_state.pdf_path = "temp.pdf"
         with open(st.session_state.pdf_path, "wb") as f:
@@ -31,9 +55,9 @@ if pdf_source == "Upload a PDF file":
         st.session_state.vector_created = False
 
 elif pdf_source == "Enter a PDF URL":
-    pdf_url = st.text_input("Enter PDF URL:")
-    if pdf_url and not st.session_state.get("pdf_loaded", False):
-        with st.spinner("Downloading PDF..."):
+    pdf_url = st.sidebar.text_input("Enter PDF URL:")
+    if pdf_url and not st.session_state.pdf_loaded:
+        with st.spinner("🔄 Downloading PDF..."):
             try:
                 response = requests.get(pdf_url)
                 if response.status_code == 200:
@@ -49,3 +73,79 @@ elif pdf_source == "Enter a PDF URL":
             except Exception as e:
                 st.error(f"Error downloading PDF: {e}")
 
+# ----------------- Process PDF -----------------
+if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
+    with st.spinner("🔄 Processing document... Please wait."):
+        loader = PDFPlumberLoader(st.session_state.pdf_path)
+        docs = loader.load()
+
+        # Embedding Model
+        model_name = "nomic-ai/modernbert-embed-base"
+        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"})
+
+        # Split into Chunks
+        text_splitter = SemanticChunker(embedding_model)
+        document_chunks = text_splitter.split_documents(docs)
+
+        # Store chunks in session state
+        st.session_state.processed_chunks = document_chunks
+        st.session_state.pdf_loaded = True
+        st.success("✅ Document processed and chunked successfully!")
+
+# ----------------- Setup Vector Store  -----------------
+if not st.session_state.vector_created and st.session_state.processed_chunks:
+    with st.spinner("🔄 Initializing Vector Store..."):
+        vector_store = Chroma(
+            collection_name="deepseek_collection",
+            collection_metadata={"hnsw:space": "cosine"},
+            embedding_function=embedding_model,
+            persist_directory="./chroma_langchain_db"
+        )
+        vector_store.add_documents(st.session_state.processed_chunks)
+        st.session_state.vector_store = vector_store
+        st.session_state.vector_created = True
+        st.success("✅ Vector store initialized successfully!")
+
+# ----------------- Query Input -----------------
+query = st.text_input("🔍 Ask a question about the document:")
+
+if query:
+    with st.spinner("🔄 Retrieving relevant context..."):
+        retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+        retrieved_docs = retriever.invoke(query)
+        context = [d.page_content for d in retrieved_docs]
+        st.success("✅ Context retrieved successfully!")
+
+    # ----------------- Full SequentialChain Execution -----------------
+    with st.spinner("🔄 Running full pipeline..."):
+        context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt)
+        relevant_prompt = PromptTemplate(input_variables=["relevancy_response"], template=relevant_context_picker_prompt)
+        context_prompt = PromptTemplate(input_variables=["context_number", "context"], template=response_synth)
+        final_prompt = PromptTemplate(input_variables=["query", "context"], template=rag_prompt)
+
+        context_relevancy_chain = LLMChain(llm=llm_judge, prompt=context_relevancy_checker_prompt, output_key="relevancy_response")
+        relevant_context_chain = LLMChain(llm=llm_judge, prompt=relevant_prompt, output_key="context_number")
+        relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts")
+        response_chain = LLMChain(llm=rag_llm, prompt=final_prompt, output_key="final_response")
+
+        context_management_chain = SequentialChain(
+            chains=[context_relevancy_chain, relevant_context_chain, relevant_contexts_chain, response_chain],
+            input_variables=["context", "retriever_query", "query"],
+            output_variables=["relevancy_response", "context_number", "relevant_contexts", "final_response"]
+        )
+
+        final_output = context_management_chain.invoke({"context": context, "retriever_query": query, "query": query})
+        st.success("✅ Full pipeline executed successfully!")
+
+    # ----------------- Display All Outputs -----------------
+    st.subheader("🟥 Context Relevancy Evaluation")
+    st.json(final_output["relevancy_response"])
+
+    st.subheader("🟦 Picked Relevant Contexts")
+    st.json(final_output["context_number"])
+
+    st.subheader("🟥 Extracted Relevant Contexts")
+    st.json(final_output["relevant_contexts"])
+
+    st.subheader("🟥 RAG Final Response")
+    st.write(final_output["final_response"])