Update app.py

app.py

@@ -7,8 +7,6 @@ from dotenv import load_dotenv
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.prompts import PromptTemplate
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
 from langchain.schema import Document
 from langchain_together import Together
 import streamlit as st
@@ -30,12 +28,13 @@ chunked_json_path = "chunked_bns_dataset.json"
 faiss_db_path = "bns_vector_db"
 
 # Streamlit page config
-st.set_page_config(page_title="LawGPT - Bharatiya Nyay Sanhita")
+st.set_page_config(page_title="LawGPT - Bharatiya Nyay Sanhita", layout="centered")
 
-# Streamlit UI
 st.title("Law4Her: Bharatiya Nyay Sanhita")
-st.markdown("This app provides answers to legal questions based on the **Bharatiya Nyay Sanhita (BNS)**.")
-st.image("https://res.cloudinary.com/dzzhbgbnp/image/upload/v1736073326/lawforher_logo1_yznqxr.png", use_container_width=True)
+st.markdown("This app provides answers to legal questions based on the *Bharatiya Nyay Sanhita (BNS)*.")
+st.image("https://res.cloudinary.com/dzzhbgbnp/image/upload/v1736073326/lawforher_logo1_yznqxr.png",
+         use_container_width=True)
+
 
 # Load CSV and Process Dataset
 def process_csv_to_chunks(csv_file, output_json, chunk_size=512, overlap=100):
@@ -76,6 +75,7 @@ def process_csv_to_chunks(csv_file, output_json, chunk_size=512, overlap=100):
 
     return chunks
 
+
 if not os.path.exists(chunked_json_path):
     logger.info("Processing CSV to JSON...")
     chunks = process_csv_to_chunks(csv_path, chunked_json_path)
@@ -94,7 +94,8 @@ embeddings = HuggingFaceEmbeddings(
 if not os.path.exists(f"{faiss_db_path}/index.faiss"):
     logger.info("Creating FAISS vectorstore...")
     texts = [
-        Document(page_content=chunk["chunk"], metadata={"chapter": chunk["chapter"], "section_title": chunk["section_title"]})
+        Document(page_content=chunk["chunk"],
+                 metadata={"chapter": chunk["chapter"], "section_title": chunk["section_title"]})
         for chunk in chunks
     ]
     db = FAISS.from_documents(texts, embeddings)
@@ -107,103 +108,79 @@ retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})
 
 # Define Prompt Template
 prompt_template = """<s>[INST]
-You are a legal chatbot specializing in the Bharatiya Nyay Sanhita (Indian Penal Code replacement). Provide concise, context-aware answers in a conversational tone. Avoid presenting the response as a question-answer format. If the requested information is not available in the provided dataset, respond with: "The required information is not available."
-
-When providing legal information, always cite the relevant section number(s) explicitly.
+You are a legal chatbot specializing in the Bharatiya Nyay Sanhita (Indian Penal Code replacement). Provide answers **only based on the provided CONTEXT**. 
+If the requested information is not available in the CONTEXT, respond with: "The required information is not available."
 
 CONTEXT: {context}
 
-CHAT HISTORY:
-{chat_history}
-
-USER QUERY:
-{question}
+USER QUERY: {question}
 
 RESPONSE:
 </s>[INST]
 """
-prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
+prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
 
 # Initialize Together API
 llm = Together(
     model="mistralai/Mistral-7B-Instruct-v0.2",
-    temperature=0.3,  # Lower temperature for deterministic output
-    max_tokens=500,   # Limit token length
+    temperature=0.5,
+    max_tokens=1024,
     together_api_key=TOGETHER_AI_API,
 )
 
-# Conversational Retrieval Chain
-qa_chain = ConversationalRetrievalChain.from_llm(
-    llm=llm,
-    memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True),
-    retriever=retriever,
-    combine_docs_chain_kwargs={"prompt": prompt},
-)
 
-# Function to format retrieved documents into context for the prompt
-def format_context(retrieved_docs):
-    context = ""
-    for doc in retrieved_docs:
-        context += f"Chapter: {doc.metadata.get('chapter', 'N/A')}\n"
-        context += f"Section Title: {doc.metadata.get('section_title', 'N/A')}\n"
-        context += f"Content: {doc.page_content}\n\n"
-    return context
-
-# Log the full prompt sent to the LLM
-def log_full_prompt(context, chat_history, user_query):
-    full_prompt = prompt.format(context=context, question=user_query, chat_history=chat_history)
-    logger.info(f"Full Prompt Sent to LLM:\n{full_prompt}")
-    return full_prompt
-
-# Reset Chat
-def reset_conversation():
-    st.session_state.messages = [{"role": "assistant", "content": "Hi, how can I assist you today?"}]
-    st.session_state.memory.clear()
+# Chat Interface
+def generate_response(user_query):
+    try:
+        # Retrieve relevant documents
+        retrieved_docs = retriever.get_relevant_documents(user_query)
+
+        # Log retrieved documents
+        logger.info(f"User Query: {user_query}")
+        for i, doc in enumerate(retrieved_docs):
+            logger.info(
+                f"Document {i + 1}: Chapter - {doc.metadata['chapter']}, Section - {doc.metadata['section_title']}")
+            logger.info(f"Content: {doc.page_content}")
+
+        # Prepare context for LLM
+        context = "\n\n".join(
+            f"Chapter: {doc.metadata['chapter']}, Section: {doc.metadata['section_title']}\n{doc.page_content}"
+            for doc in retrieved_docs
+        )
+
+        # Construct LLM prompt input
+        prompt_input = {"context": context, "question": user_query}
 
+        # Generate response using LLM
+        logger.debug(f"Payload sent to LLM: {json.dumps(prompt_input, ensure_ascii=False, indent=2)}")
+        response = llm(prompt.format(**prompt_input))
+
+        return response
+
+    except Exception as e:
+        logger.error(f"Error generating response: {e}")
+        return "An error occurred while generating the response."
+
+
+# Streamlit Chat Interface
 if "messages" not in st.session_state:
-    st.session_state.messages = [{"role": "assistant", "content": "Hi, how can I assist you today?"}]
-    st.session_state.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    st.session_state.messages = [{"role": "assistant", "content": "Hi! How can I assist you today?"}]
 
-# Chat Interface
+# Display Chat Messages
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
-# User Input Handling
-user_input = st.chat_input("Ask a legal question about the Bharatiya Nyay Sanhita")
-
-if user_input:
+# User Input
+if user_input := st.chat_input("Type your question here..."):
+    # User message
+    st.session_state.messages.append({"role": "user", "content": user_input})
     with st.chat_message("user"):
         st.write(user_input)
 
-    st.session_state.messages.append({"role": "user", "content": user_input})
-
+    # Assistant response
     with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            try:
-                # Retrieve and log the documents
-                retrieved_docs = retriever.get_relevant_documents(user_input)
-                formatted_context = format_context(retrieved_docs)
-                logger.info(f"Retrieved Context:\n{formatted_context}")
-
-                # Retrieve chat history
-                chat_history = st.session_state.memory.load_memory_variables({}).get("chat_history", "")
-
-                # Log the full prompt
-                log_full_prompt(context=formatted_context, chat_history=chat_history, user_query=user_input)
-
-                # Generate the response
-                response = qa_chain.invoke(
-                    input=user_input,
-                    context=formatted_context,
-                    chat_history=chat_history
-                ).get("answer", "The required information is not available.")
-                
-                # Append and display the response
-                st.session_state.messages.append({"role": "assistant", "content": response})
-                st.write(response)
-            except Exception as e:
-                st.error(f"Error: {e}")
-                logger.error(f"Error during response generation: {e}")
-
-st.button("Reset Chat", on_click=reset_conversation)
+        with st.spinner("Generating response..."):
+            response = generate_response(user_input)
+            st.write(response)
+        st.session_state.messages.append({"role": "assistant", "content": response})