lab/content_key_issue_fixed.py

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": <your score either 0 or 1>,"Reasoning":<why you have chose the score as 0 or 1>}},
        {{"content:2,"score": <your score either 0 or 1>,"Reasoning":<why you have chose the score as 0 or 1>}},
        ...]
        """
        
        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":<content number>}}]
            """
        )
        
        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":<<content_number>>}}).
            (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":<content1>,"relevant_content":<content corresponing to that element 1 in the Content List>}},
            {{"context_number":<content4>,"relevant_content":<content corresponing to that element 4 in the Content List>}},
            ...
            ]
            """
        )

        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"])