Create content_key_issue_fixed.py

lab/content_key_issue_fixed.py

@@ -0,0 +1,292 @@
+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"])