import streamlit as st import os import requests 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") # ----------------- API Keys ----------------- os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "") # ----------------- Clear ChromaDB Cache ----------------- chromadb.api.client.SharedSystemClient.clear_system_cache() # ----------------- 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") # Enable verbose logging for debugging llm_judge.verbose = True rag_llm.verbose = 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.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: f.write(uploaded_file.getbuffer()) st.session_state.pdf_loaded = False st.session_state.chunked = False st.session_state.vector_created = False elif pdf_source == "Enter a PDF URL": 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: st.session_state.pdf_path = "temp.pdf" with open(st.session_state.pdf_path, "wb") as f: f.write(response.content) st.session_state.pdf_loaded = False st.session_state.chunked = False st.session_state.vector_created = False st.success("✅ PDF Downloaded Successfully!") else: st.error("❌ Failed to download PDF. Check the 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 (HF on CPU) model_name = "nomic-ai/modernbert-embed-base" embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}) # Prevent unnecessary re-chunking if not st.session_state.chunked: text_splitter = SemanticChunker(embedding_model) document_chunks = text_splitter.split_documents(docs) st.session_state.processed_chunks = document_chunks st.session_state.chunked = True 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..."): st.session_state.vector_store = Chroma( collection_name="deepseek_collection", collection_metadata={"hnsw:space": "cosine"}, embedding_function=embedding_model ) st.session_state.vector_store.add_documents(st.session_state.processed_chunks) 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..."): final_output = SequentialChain( chains=[ LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt), output_key="relevancy_response"), LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["relevancy_response"], template=relevant_context_picker_prompt), output_key="context_number"), LLMChain(llm=llm_judge, prompt=PromptTemplate(input_variables=["context_number", "context"], template=response_synth), output_key="relevant_contexts"), LLMChain(llm=rag_llm, prompt=PromptTemplate(input_variables=["query", "context"], template=rag_prompt), output_key="final_response") ], input_variables=["context", "retriever_query", "query"], output_variables=["relevancy_response", "context_number", "relevant_contexts", "final_response"] ).invoke({"context": context, "retriever_query": query, "query": query}) # ----------------- Display All Outputs ----------------- st.markdown("### 🟥 Context Relevancy Evaluation") st.json(final_output["relevancy_response"]) st.markdown("### 🟦 Picked Relevant Contexts") st.json(final_output["context_number"]) st.markdown("### 🟥 Extracted Relevant Contexts") st.json(final_output["relevant_contexts"]) st.markdown("## 🟥 RAG Final Response") st.write(final_output["final_response"]) # ----------------- Streamlit-Friendly Debugging (Replacing print statements) ----------------- st.markdown("### Debug Logs:") st.text("\n-------- 🟥 Context Relevancy Evaluation Statement 🟥 --------\n") st.json(final_output["relevancy_response"]) st.text("\n-------- 🟦 Picked Relevant Context Statement 🟦 --------\n") st.json(final_output["context_number"]) st.text("\n-------- 🟥 Relevant Contexts Statement 🟥 --------\n") st.json(final_output["relevant_contexts"]) st.text("\n-------- 🟥 RAG Response Statement 🟥 --------\n") st.write(final_output["final_response"])