Update app.py

app.py

@@ -1,24 +1,104 @@
+import os
+import time
+import streamlit as st
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain_together import Together
+from footer import footer  # Ensure this module is present in the working directory
 
-import gradio
-from transformers import pipeline
+# Set Streamlit configuration
+st.set_page_config(page_title="AI Legal App", layout="centered")
 
-import groq
-import gradio
+# Display a logo or banner (replace with a local image or URL)
+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)
 
-groq.api_key = "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"
+def hide_hamburger_menu():
+    st.markdown("""
+        <style>
+            #MainMenu {visibility: hidden;}
+            footer {visibility: hidden;}
+        </style>
+        """, unsafe_allow_html=True)
 
-messages = [{"role": "system", "content": "You are a Indian Lawyer and Gave Advice according to Indian constitution"}]
+hide_hamburger_menu()
 
-def CustomChatGPT(user_input):
-    messages.append({"role": "user", "content": user_input})
-    response = groq.ChatCompletion.create(
-        model = "gpt-3.5-turbo",
-        messages = messages
-    )
-    ChatGPT_reply = response["choices"][0]["message"]["content"]
-    messages.append({"role": "assistant", "content": ChatGPT_reply})
-    return ChatGPT_reply
+# Initialize session state
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-iface = gradio.Interface(fn=CustomChatGPT, inputs = "text", outputs = "text", title = "AI Chat Bot for Legal Assistance This is Free version and if it doesnt work kindly contact at Email ID contact2abhaygupta@gmail.com")
-print(" This is Free version and if it doesnt work kindly contact at Email ID contact2abhaygupta@gmail.com")
-iface.launch(inline=False)
+if "memory" not in st.session_state:
+    st.session_state.memory = ConversationBufferWindowMemory(k=2, memory_key="chat_history", return_messages=True)
+
+@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)
+db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+
+prompt_template = """
+<s>[INST]
+As a legal chatbot specializing in Indian law, your responses must be concise and accurate:
+- Provide bullet points summarizing key legal aspects.
+- Avoid assumptions or overly specific advice unless requested.
+- Clarify any common misconceptions.
+- Keep responses aligned with general legal principles.
+CONTEXT: {context}
+CHAT HISTORY: {chat_history}
+QUESTION: {question}
+ANSWER:
+</s>[INST]
+"""
+
+prompt = PromptTemplate(template=prompt_template,
+                        input_variables=['context', 'question', 'chat_history'])
+
+api_key = os.getenv('TOGETHER_API_KEY')
+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})
+
+def extract_answer(full_response):
+    """Extracts the assistant's answer from the response."""
+    return full_response.strip()
+
+def reset_conversation():
+    st.session_state.messages = []
+    st.session_state.memory.clear()
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+input_prompt = st.chat_input("Ask your legal query...")
+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("Analyzing..."):
+            result = qa.invoke(input=input_prompt)
+            message_placeholder = st.empty()
+            answer = extract_answer(result["answer"])
+
+            # Simulated typing effect
+            response = ""
+            for char in answer:
+                response += char
+                time.sleep(0.02)
+                message_placeholder.markdown(response + " |", unsafe_allow_html=True)
+
+        st.session_state.messages.append({"role": "assistant", "content": answer})
+
+        if st.button('🗑️ Reset Chat', on_click=reset_conversation):
+            st.experimental_rerun()
+
+footer()