Update app.py

app.py

@@ -1,5 +1,6 @@
-import time
+
 import os
+import time
 import streamlit as st
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
@@ -7,17 +8,16 @@ from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferWindowMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain_together import Together
-
 from footer import footer
 
-# Set the Streamlit page configuration and theme
 st.set_page_config(page_title="BharatLAW", layout="centered")
 
-# Display the logo image
+# Display logo
 col1, col2, col3 = st.columns([1, 30, 1])
 with col2:
-    st.image("https://github.com/Nike-one/BharatLAW/blob/master/images/banner.png?raw=true", use_column_width=True)
+    st.image("https://github.com/Nike-one/BharatLAW/blob/master/images/banner.png?raw=true", use_container_width=True)
 
+# Hide hamburger and default footer
 def hide_hamburger_menu():
     st.markdown("""
         <style>
@@ -28,7 +28,7 @@ def hide_hamburger_menu():
 
 hide_hamburger_menu()
 
-# Initialize session state for messages and memory
+# Session state init
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
@@ -37,23 +37,27 @@ if "memory" not in st.session_state:
 
 @st.cache_resource
 def load_embeddings():
-    """Load and cache the embeddings model."""
     return HuggingFaceEmbeddings(model_name="law-ai/InLegalBERT")
 
 embeddings = load_embeddings()
-db = FAISS.load_local("ipc_embed_db", embeddings, allow_dangerous_deserialization=True)
+
+# ✅ Add FAISS index check before loading
+index_dir = "ipc_embed_db"
+index_faiss = os.path.join(index_dir, "index.faiss")
+index_pkl = os.path.join(index_dir, "index.pkl")
+
+if os.path.exists(index_faiss) and os.path.exists(index_pkl):
+    db = FAISS.load_local(index_dir, embeddings, allow_dangerous_deserialization=True)
+else:
+    st.error("❌ FAISS index files not found! Please run the FAISS indexing script first.")
+    st.stop()
+
 db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})
 
+# Prompt template for QA
 prompt_template = """
 <s>[INST]
-As a legal chatbot specializing in the Indian Penal Code, you are tasked with providing highly accurate and contextually appropriate responses. Ensure your answers meet these criteria:
-- Respond in a bullet-point format to clearly delineate distinct aspects of the legal query.
-- Each point should accurately reflect the breadth of the legal provision in question, avoiding over-specificity unless directly relevant to the user's query.
-- Clarify the general applicability of the legal rules or sections mentioned, highlighting any common misconceptions or frequently misunderstood aspects.
-- Limit responses to essential information that directly addresses the user's question, providing concise yet comprehensive explanations.
-- Avoid assuming specific contexts or details not provided in the query, focusing on delivering universally applicable legal interpretations unless otherwise specified.
-- Conclude with a brief summary that captures the essence of the legal discussion and corrects any common misinterpretations related to the topic.
-
+As a legal chatbot specializing in the Indian Penal Code, you are tasked with providing highly accurate and contextually appropriate responses...
 CONTEXT: {context}
 CHAT HISTORY: {chat_history}
 QUESTION: {question}
@@ -66,21 +70,17 @@ ANSWER:
 </s>[INST]
 """
 
-
-
-prompt = PromptTemplate(template=prompt_template,
-                        input_variables=['context', 'question', 'chat_history'])
+prompt = PromptTemplate(template=prompt_template, input_variables=['context', 'question', 'chat_history'])
 
 api_key = os.getenv('TOGETHER_API_KEY')
 if not api_key:
     st.error("API key for Together is missing. Please set the TOGETHER_API_KEY environment variable.")
 
 llm = Together(model="mistralai/Mixtral-8x22B-Instruct-v0.1", temperature=0.5, max_tokens=1024, together_api_key=api_key)
-
 qa = ConversationalRetrievalChain.from_llm(llm=llm, memory=st.session_state.memory, retriever=db_retriever, combine_docs_chain_kwargs={'prompt': prompt})
 
+# Helper functions
 def extract_answer(full_response):
-    """Extracts the answer from the LLM's full response by removing the instructional text."""
     try:
         answer_start = full_response.find("Response:")
         if answer_start != -1:
@@ -94,29 +94,27 @@ def reset_conversation():
     st.session_state.messages = []
     st.session_state.memory.clear()
 
+# Chat Interface
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
-
 input_prompt = st.chat_input("Say something...")
 if input_prompt:
     with st.chat_message("user"):
         st.markdown(f"**You:** {input_prompt}")
-
     st.session_state.messages.append({"role": "user", "content": input_prompt})
+
     with st.chat_message("assistant"):
         with st.spinner("Thinking 💡..."):
             result = qa.invoke(input=input_prompt)
             message_placeholder = st.empty()
             answer = extract_answer(result["answer"])
 
-            # Initialize the response message
             full_response = "⚠️ **_Gentle reminder: We generally ensure precise information, but do double-check._** \n\n\n"
             for chunk in answer:
-                # Simulate typing by appending chunks of the response over time
                 full_response += chunk
-                time.sleep(0.02)  # Adjust the sleep time to control the "typing" speed
+                time.sleep(0.02)
                 message_placeholder.markdown(full_response + " |", unsafe_allow_html=True)
 
         st.session_state.messages.append({"role": "assistant", "content": answer})
@@ -124,9 +122,5 @@ if input_prompt:
         if st.button('🗑️ Reset All Chat', on_click=reset_conversation):
             st.experimental_rerun()
 
-
-
-# Define the CSS to style the footer
+# Footer styling
 footer()
-
-