Update app.py

app.py

@@ -23,18 +23,22 @@ rag_llm.verbose = True
 
 st.title("❓")
 
+# Initialize session state variables
+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
+
 # Step 1: Choose PDF Source
-#### Initialize pdf_path
-pdf_path = None
 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:
-        with open("temp.pdf", "wb") as f:
-            f.write(uploaded_file.getbuffer())
         pdf_path = "temp.pdf"
+        with open(pdf_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        st.success("✅ PDF Uploaded Successfully!")
 
 elif pdf_source == "Enter a PDF URL":
     pdf_url = st.text_input("Enter PDF URL:")
@@ -43,9 +47,9 @@ elif pdf_source == "Enter a PDF URL":
             try:
                 response = requests.get(pdf_url)
                 if response.status_code == 200:
-                    with open("temp.pdf", "wb") as f:
-                        f.write(response.content)
                     pdf_path = "temp.pdf"
+                    with open(pdf_path, "wb") as f:
+                        f.write(response.content)
                     st.success("✅ PDF Downloaded Successfully!")
                 else:
                     st.error("❌ Failed to download PDF. Check the URL.")
@@ -56,40 +60,40 @@ elif pdf_source == "Enter a PDF URL":
 else:
     pdf_path = None
 
-# Step 2: Process PDF
-if pdf_path:
-    with st.spinner("Loading PDF..."):
+# Step 2: Process PDF and Create Vector Store (Only if Not Processed)
+if pdf_path and st.session_state.vector_store is None:
+    with st.spinner("Loading and processing PDF..."):
         loader = PDFPlumberLoader(pdf_path)
         docs = loader.load()
-    
-    st.success(f"✅ **PDF Loaded!** Total Pages: {len(docs)}")
-
-    # Step 3: Chunking
-    with st.spinner("Chunking the document..."):
-        model_name = "nomic-ai/modernbert-embed-base"
-        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'})
-        text_splitter = SemanticChunker(embedding_model)
-        documents = text_splitter.split_documents(docs)
-
-    st.success(f"✅ **Document Chunked!** Total Chunks: {len(documents)}")
-
-    # Step 4: Setup Vectorstore
-    with st.spinner("Creating vector store..."):
-        vector_store = Chroma(
-            collection_name="deepseek_collection",
-            collection_metadata={"hnsw:space": "cosine"},
-            embedding_function=embedding_model
-        )
-        vector_store.add_documents(documents)
-        num_documents = len(vector_store.get()["documents"])
-
-    st.success(f"✅ **Vector Store Created!** Total documents stored: {num_documents}")
-
-    # Step 5: Query Input
+        st.success(f"✅ **PDF Loaded!** Total Pages: {len(docs)}")
+
+        # Step 3: Chunking
+        with st.spinner("Chunking the document..."):
+            model_name = "nomic-ai/modernbert-embed-base"
+            embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'})
+            text_splitter = SemanticChunker(embedding_model)
+            documents = text_splitter.split_documents(docs)
+            st.session_state.documents = documents  # Store in session state
+            st.success(f"✅ **Document Chunked!** Total Chunks: {len(documents)}")
+
+        # Step 4: Setup Vectorstore
+        with st.spinner("Creating vector store..."):
+            vector_store = Chroma(
+                collection_name="deepseek_collection",
+                collection_metadata={"hnsw:space": "cosine"},
+                embedding_function=embedding_model
+            )
+            vector_store.add_documents(documents)
+            num_documents = len(vector_store.get()["documents"])
+            st.session_state.vector_store = vector_store  # Store vector store in session state
+            st.success(f"✅ **Vector Store Created!** Total documents stored: {num_documents}")
+
+# Step 5: Query Input (Only allow if vector store exists)
+if st.session_state.vector_store:
     query = st.text_input("🔍 Enter a Query:")
     if query:
         with st.spinner("Retrieving relevant contexts..."):
-            retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+            retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
             contexts = retriever.invoke(query)
             context_texts = [doc.page_content for doc in contexts]
 
@@ -99,9 +103,7 @@ if pdf_path:
 
         # Step 6: Context Relevancy Checker
         with st.spinner("Evaluating context relevancy..."):
-            context_relevancy_checker_prompt = PromptTemplate(
-                input_variables=["retriever_query", "context"], template=relevancy_prompt
-            )
+            context_relevancy_checker_prompt = PromptTemplate(input_variables=["retriever_query", "context"], template=relevancy_prompt)
             context_relevancy_chain = LLMChain(llm=llm_judge, prompt=context_relevancy_checker_prompt, output_key="relevancy_response")
             relevancy_response = context_relevancy_chain.invoke({"context": context_texts, "retriever_query": query})
 
@@ -110,9 +112,7 @@ if pdf_path:
 
         # Step 7: Selecting Relevant Contexts
         with st.spinner("Selecting the most relevant contexts..."):
-            relevant_prompt = PromptTemplate(
-                input_variables=["relevancy_response"], template=relevant_context_picker_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": relevancy_response['relevancy_response']})
 
@@ -121,9 +121,7 @@ if pdf_path:
 
         # Step 8: Retrieving Context for Response Generation
         with st.spinner("Retrieving final context..."):
-            context_prompt = PromptTemplate(
-                input_variables=["context_number", "context"], template=response_synth
-            )
+            context_prompt = PromptTemplate(input_variables=["context_number", "context"], template=response_synth)
             relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts")
             final_contexts = relevant_contexts_chain.invoke({"context_number": relevant_response['context_number'], "context": context_texts})
 
@@ -132,9 +130,7 @@ if pdf_path:
 
         # Step 9: Generate Final Response
         with st.spinner("Generating the final answer..."):
-            final_prompt = PromptTemplate(
-                input_variables=["query", "context"], template=rag_prompt
-            )
+            final_prompt = PromptTemplate(input_variables=["query", "context"], template=rag_prompt)
             response_chain = LLMChain(llm=rag_llm, prompt=final_prompt, output_key="final_response")
             final_response = response_chain.invoke({"query": query, "context": final_contexts['relevant_contexts']})
 
@@ -148,4 +144,6 @@ if pdf_path:
             "Relevant Contexts": relevant_response["context_number"],
             "Extracted Contexts": final_contexts["relevant_contexts"],
             "Final Answer": final_response["final_response"]
-        })
+        })
+else:
+    st.warning("📄 Please upload or provide a PDF URL first.")