lab/i_shine.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", key="pdf_url", on_change=lambda: st.session_state.update(trigger_download=True))

    # Button OR Enter key will trigger download
    if st.button("Download and Process PDF") or st.session_state.get("trigger_download", False):
        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)

                    # Reset states
                    st.session_state.pdf_loaded = False
                    st.session_state.chunked = False
                    st.session_state.vector_created = False
                    st.session_state.trigger_download = False  # Reset trigger

                    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 = None  

# Check if a valid PDF URL has been entered (but not processed yet)
pdf_url_entered = bool(st.session_state.get("pdf_url"))  # Checks if text is in the input box

# No PDF Provided Yet
if not st.session_state.pdf_path and not pdf_url_entered:
    st.info("📥 **Please upload a PDF or enter a valid URL to proceed.**")

# PDF URL Exists but Not Processed Yet (Only show if URL exists but hasn't been downloaded)
elif pdf_url_entered and not st.session_state.pdf_loaded:
    st.warning("⚠️ **PDF URL detected! Click 'Download and Process PDF' to proceed.**")

# Processing in Progress 
elif st.session_state.get("trigger_download", False) and (
    not st.session_state.pdf_loaded or not st.session_state.chunked or not st.session_state.vector_created
):
    st.info("⏳ **Processing your document... Please wait.**")

# ✅ Step 4: Processing Complete, Ready for Questions
elif st.session_state.pdf_loaded and st.session_state.chunked and st.session_state.vector_created:
    st.success("🎉 **Processing complete! You can now ask questions.**")
    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))
        
        #st.write("📄 **Relevant Passages Retrieved:**")
        #for idx, doc in enumerate(contexts):
            #st.write(f"**Excerpt {idx+1}:** {doc.page_content}")

        
        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="""
            Your 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 a helpful 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})
        # 📌 Display Final Answer (Always Visible)
        st.markdown("## 💡 **Final Answer**")
        st.success(final_output["final_response"] if final_output["final_response"] else "⚠️ No clear answer found based on retrieved content.")

        st.markdown("---")

        # 🟥 Relevance Analysis (context_relevancy_evaluation_chain Statement)
        with st.expander("🟥 **Relevance Analysis (context_relevancy_evaluation_chain Statement)**"):
            if final_output["relevancy_response"]:  # Don't restrict to list, print whatever is there
                st.write(final_output["relevancy_response"])
            else:
                st.warning("⚠️ No relevance analysis available.")

        st.markdown("---")

        # 🟦 Picked Relevant Contexts (pick_relevant_context_chain Statement)
        with st.expander("🟦 **Picked Relevant Contexts (pick_relevant_context_chain Statement)**"):
            if final_output["context_number"]:  
                st.write(final_output["context_number"])
            else:
                st.warning("⚠️ No relevant context numbers available.")

        st.markdown("---")

        # 🟥 Extracted Relevant Contexts (relevant_contexts_chain Statement)
        with st.expander("🟥 **Extracted Relevant Contexts (relevant_contexts_chain Statement)**"):
            if final_output["relevant_contexts"]:
                st.write(final_output["relevant_contexts"])
            else:
                st.warning("⚠️ No relevant contexts available.")