Update app.py

app.py

@@ -26,11 +26,7 @@ chromadb.api.client.SharedSystemClient.clear_system_cache()
 
 st.title("Blah")
 
-# 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
+# **Initialize session state variables**
 if "pdf_path" not in st.session_state:
     st.session_state.pdf_path = None  
 if "pdf_loaded" not in st.session_state:
@@ -39,8 +35,14 @@ 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
+if "vector_store_path" not in st.session_state:
+    st.session_state.vector_store_path = "./chroma_langchain_db"  
+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
+# **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":
@@ -81,7 +83,7 @@ if st.session_state.pdf_path and not st.session_state.pdf_loaded:
         st.session_state.pdf_loaded = True
         st.success(f"✅ **PDF Loaded!** Total Pages: {len(docs)}")
 
-# Step 3: Chunking (Only if Not Already Done)
+# Step 3: Chunking
 if st.session_state.pdf_loaded and not st.session_state.chunked:
     with st.spinner("Chunking the document..."):
         model_name = "nomic-ai/modernbert-embed-base"
@@ -98,7 +100,8 @@ if st.session_state.chunked and not st.session_state.vector_created:
         vector_store = Chroma(
             collection_name="deepseek_collection",
             collection_metadata={"hnsw:space": "cosine"},
-            embedding_function=embedding_model
+            embedding_function=embedding_model,
+            persist_directory=st.session_state.vector_store_path 
         )
         vector_store.add_documents(st.session_state.documents)
         num_documents = len(vector_store.get()["documents"])
@@ -109,7 +112,8 @@ if st.session_state.chunked and not st.session_state.vector_created:
 # Step 5: Query Input
 if st.session_state.vector_created:
     query = st.text_input("🔍 Enter a Query:")
-    if query:
+    
+    if query and st.session_state.vector_store: 
         with st.spinner("Retrieving relevant contexts..."):
             retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
             contexts = retriever.invoke(query)
@@ -119,41 +123,11 @@ if st.session_state.vector_created:
         for i, text in enumerate(context_texts, 1):
             st.write(f"**Context {i}:** {text[:500]}...")
 
-        # 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_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})
-
-        st.subheader("🟥 Context Relevancy Evaluation")
-        st.json(relevancy_response['relevancy_response'])
-
-        # 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)
-            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']})
-
-        st.subheader("🟦 Pick Relevant Context Chain")
-        st.json(relevant_response['context_number'])
-
-        # Step 8: Retrieving Context for Response Generation
-        with st.spinner("Retrieving final context..."):
-            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})
-
-        st.subheader("🟥 Relevant Contexts Extracted")
-        st.json(final_contexts['relevant_contexts'])
-
-        # Step 9: Generate Final Response
+        # **Step 6: Generate Final Response**
         with st.spinner("Generating the final answer..."):
             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']})
+            final_response = response_chain.invoke({"query": query, "context": context_texts})
 
         st.subheader("🟥 RAG Final Response")
         st.success(final_response['final_response'])
-
-else:
-    st.warning("📄 Please upload or provide a PDF URL first.")