lab/content_key_issue.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 prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth

# 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:", ["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")
    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:")
    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))
    # ----------------- Run Individual Chains Explicitly -----------------

    context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query","context"],template=relevancy_prompt)

    if "context_relevancy_evaluation_chain" in st.session_state:
        del st.session_state["context_relevancy_evaluation_chain"]



    context_relevancy_evaluation_chain = LLMChain( llm=llm_judge, prompt=PromptTemplate(input_variables=["retriever_query","context"],template=relevancy_prompt), output_key="relevancy_response" )

    st.write("🚀 Debugging Expected Keys for `context_relevancy_evaluation_chain`:")
    st.write(context_relevancy_evaluation_chain.input_keys)

    response_crisis = context_relevancy_evaluation_chain.invoke({"context":context,"retriever_query":query})

    relevant_prompt = relevant_prompt = PromptTemplate(input_variables=["relevancy_response"], template= relevant_context_picker_prompt)

    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']})

    context_prompt = PromptTemplate(input_variables=["context_number"],template=response_synth)

    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})

    #temp
    st.subheader("Relevant Contexts")
    st.json(contexts['relevant_contexts']) 

    response_chain = LLMChain(llm=rag_llm,prompt=final_prompt,output_key="final_response")


    #temp
    st.subheader("Response Chain")
    st.json(response_chain)

    #response = chain.invoke({"query":query,"context":contexts['relevant_contexts']})

    #temp
    #st.subheader("blah response")
    #st.json(response.content)

    # 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 from Context Management chain")
    st.json(final_output)
    
    st.subheader("Context of Final Output from Context Management chain")
    st.json(final_output['context'])
    
    st.header("Relevancy Response")
    st.json(final_output['relevancy_response'])
    
    st.subheader("Relevant Context")
    st.json(final_output['relevant_contexts'])
    
    response = chain.invoke({"query":query,"context":final_output['relevant_contexts']})
    
    st.subheader("Final Response")
    st.json(response.content)

    # ----------------- Display All Outputs -----------------
    #st.subheader("response_crisis")
    #st.json((response_crisis))

    #st.subheader("response_crisis['relevancy_response']")
    #st.json((response_crisis['relevancy_response']))


    
    #st.markdown("### Context Relevancy Evaluation")
    #st.json(response_crisis["relevancy_response"])

    #st.markdown("### Picked Relevant Contexts")
    #st.json(relevant_response["context_number"])

    #st.markdown("### Extracted Relevant Contexts")
    #st.json(contexts["relevant_contexts"])

    #st.subheader("context_relevancy_evaluation_chain Statement")
    #st.json(final_response["relevancy_response"])

    #st.subheader("pick_relevant_context_chain Statement")
    #st.json(final_response["context_number"])

    #st.subheader("relevant_contexts_chain Statement")
    #st.json(final_response["relevant_contexts"])

    #st.subheader("RAG Response Statement")
    #st.json(final_response["final_response"])