import os import requests import streamlit as st import pickle from langchain.chains import LLMChain from langchain.prompts import PromptTemplate from langchain_groq import ChatGroq from langchain.document_loaders import PDFPlumberLoader from langchain_experimental.text_splitter import SemanticChunker from langchain_huggingface import HuggingFaceEmbeddings from langchain_chroma import Chroma from langchain.chains import SequentialChain, LLMChain # 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 VECTOR_DB_PATH = "/tmp/chroma_db" CHUNKS_FILE = "/tmp/chunks.pkl" # Session State Initialization if "vector_store" not in st.session_state: st.session_state.vector_store = None if "documents" not in st.session_state: st.session_state.documents = None if "pdf_path" not in st.session_state: st.session_state.pdf_path = None 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 st.title("Blah-2") # Step 1: Choose PDF Source pdf_source = st.radio("Upload or provide a link to a PDF:", ["Enter a PDF URL", "Upload a PDF file"], index=0, horizontal=True) if pdf_source == "Upload a PDF file": uploaded_file = st.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.text_input("Enter PDF URL:", value = "https://arxiv.org/pdf/2406.06998") if pdf_url and not st.session_state.pdf_path: 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}") # Step 2: Load & Process PDF (Only Once) if st.session_state.pdf_path and not st.session_state.pdf_loaded: with st.spinner("Loading PDF..."): try: loader = PDFPlumberLoader(st.session_state.pdf_path) docs = loader.load() st.session_state.documents = docs st.session_state.pdf_loaded = True st.success(f"✅ **PDF Loaded!** Total Pages: {len(docs)}") except Exception as e: st.error(f"❌ Error processing PDF: {e}") # Load Cached Chunks if Available def load_chunks(): if os.path.exists(CHUNKS_FILE): with open(CHUNKS_FILE, "rb") as f: return pickle.load(f) return None if not st.session_state.chunked: # Ensure chunking only happens once cached_chunks = load_chunks() if cached_chunks: st.session_state.documents = cached_chunks st.session_state.chunked = True # Step 3: Chunking (Only Happens Once) if st.session_state.pdf_loaded and not st.session_state.chunked: with st.spinner("Chunking the document..."): try: model_name = "nomic-ai/modernbert-embed-base" embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'}) text_splitter = SemanticChunker(embedding_model) if st.session_state.documents: documents = text_splitter.split_documents(st.session_state.documents) st.session_state.documents = documents st.session_state.chunked = True # Save chunks for persistence with open(CHUNKS_FILE, "wb") as f: pickle.dump(documents, f) st.success(f"✅ **Document Chunked!** Total Chunks: {len(documents)}") except Exception as e: st.error(f"❌ Error chunking document: {e}") # Step 4: Setup Vectorstore def load_vector_store(): return Chroma(persist_directory=VECTOR_DB_PATH, collection_name="deepseek_collection", embedding_function=HuggingFaceEmbeddings(model_name="nomic-ai/modernbert-embed-base")) if st.session_state.chunked and not st.session_state.vector_created: with st.spinner("Creating vector store..."): try: if st.session_state.vector_store is None: # Prevent unnecessary reloading st.session_state.vector_store = load_vector_store() if len(st.session_state.vector_store.get()["documents"]) == 0: # Prevent duplicate insertions st.session_state.vector_store.add_documents(st.session_state.documents) num_documents = len(st.session_state.vector_store.get()["documents"]) st.session_state.vector_created = True st.success(f"✅ **Vector Store Created!** Total documents stored: {num_documents}") except Exception as e: st.error(f"❌ Error creating vector store: {e}") # Debugging Logs st.write("📄 **PDF Loaded:**", st.session_state.pdf_loaded) st.write("🔹 **Chunked:**", st.session_state.chunked) st.write("📂 **Vector Store Created:**", st.session_state.vector_created) # ----------------- 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}) contexts = retriever.invoke(query) # Debugging: Check what was retrieved st.write("Retrieved Contexts:", contexts) st.write("Number of Contexts:", len(contexts)) context = [d.page_content for d in contexts] # Debugging: Check extracted context st.write("Extracted Context (page_content):", context) st.write("Number of Extracted Contexts:", len(context)) relevancy_prompt = """You are an expert judge tasked with evaluating whether the EACH OF THE CONTEXT provided in the CONTEXT LIST is self sufficient to answer the QUERY asked. Analyze the provided QUERY AND CONTEXT to determine if each Ccontent in the CONTEXT LIST contains Relevant information to answer the QUERY. Guidelines: 1. The content must not introduce new information beyond what's provided in the QUERY. 2. Pay close attention to the subject of statements. Ensure that attributes, actions, or dates are correctly associated with the right entities (e.g., a person vs. a TV show they star in). 3. Be vigilant for subtle misattributions or conflations of information, even if the date or other details are correct. 4. Check that the content in the CONTEXT LIST doesn't oversimplify or generalize information in a way that changes the meaning of the QUERY. Analyze the text thoroughly and assign a relevancy score 0 or 1 where: - 0: The content has all the necessary information to answer the QUERY - 1: The content does not has the necessary information to answer the QUERY ``` EXAMPLE: INPUT (for context only, not to be used for faithfulness evaluation): What is the capital of France? CONTEXT: ['France is a country in Western Europe. Its capital is Paris, which is known for landmarks like the Eiffel Tower.', 'Mr. Naveen patnaik has been the chief minister of Odisha for consequetive 5 terms'] OUTPUT: The Context has sufficient information to answer the query. RESPONSE: {{"score":0}} ``` CONTENT LIST: {context} QUERY: {retriever_query} Provide your verdict in JSON format with a single key 'score' and no preamble or explanation: [{{"content:1,"score": ,"Reasoning":}}, {{"content:2,"score": ,"Reasoning":}}, ...] """ context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query","context"],template=relevancy_prompt) relevant_prompt = PromptTemplate( input_variables=["relevancy_response"], template=""" Your main task is to analyze the json structure as a part of the Relevancy Response. Review the Relevancy Response and do the following:- (1) Look at the Json Structure content (2) Analyze the 'score' key in the Json Structure content. (3) pick the value of 'content' key against those 'score' key value which has 0. Relevancy Response: {relevancy_response} Provide your verdict in JSON format with a single key 'content number' and no preamble or explanation: [{{"content":}}] """ ) context_prompt = PromptTemplate( input_variables=["context_number"], template=""" You main task is to analyze the json structure as a part of the Context Number Response and the list of Contexts provided in the 'Content List' and perform the following steps:- (1) Look at the output from the Relevant Context Picker Agent. (2) Analyze the 'content' key in the Json Structure format({{"content":<>}}). (3) Retrieve the value of 'content' key and pick up the context corresponding to that element from the Content List provided. (4) Pass the retrieved context for each corresponing element number referred in the 'Context Number Response' Context Number Response: {context_number} Content List: {context} Provide your verdict in JSON format with a two key 'relevant_content' and 'context_number' no preamble or explanation: [{{"context_number":,"relevant_content":}}, {{"context_number":,"relevant_content":}}, ... ] """ ) rag_prompt = """ You are ahelpful assistant very profiient in formulating clear and meaningful answers from the context provided.Based on the CONTEXT Provided ,Please formulate a clear concise and meaningful answer for the QUERY asked.Please refrain from making up your own answer in case the COTEXT provided is not sufficient to answer the QUERY.In such a situation please respond as 'I do not know'. QUERY: {query} CONTEXT {context} ANSWER: """ context_relevancy_evaluation_chain = LLMChain(llm=llm_judge, prompt=context_relevancy_checker_prompt, output_key="relevancy_response") response_crisis = context_relevancy_evaluation_chain.invoke({"context":context,"retriever_query":query}) pick_relevant_context_chain = LLMChain(llm=llm_judge, prompt=relevant_prompt, output_key="context_number") relevant_response = pick_relevant_context_chain.invoke({"relevancy_response":response_crisis['relevancy_response']}) relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts") contexts = relevant_contexts_chain.invoke({"context_number":relevant_response['context_number'],"context":context}) final_prompt = PromptTemplate(input_variables=["query","context"],template=rag_prompt) response_chain = LLMChain(llm=rag_llm,prompt=final_prompt,output_key="final_response") response = response_chain.invoke({"query":query,"context":contexts['relevant_contexts']}) # Orchestrate using SequentialChain context_management_chain = SequentialChain( chains=[context_relevancy_evaluation_chain ,pick_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({"context":context,"retriever_query":query,"query":query}) st.subheader('final_output["relevancy_response"]') st.json(final_output["relevancy_response"] ) st.subheader('final_output["context_number"]') st.json(final_output["context_number"]) st.subheader('final_output["relevant_contexts"]') st.json(final_output["relevant_contexts"]) st.subheader('final_output["final_response"]') st.json(final_output["final_response"])