Update app.py

app.py

@@ -1,18 +1,9 @@
+import logging
 from dataclasses import dataclass
 from typing import Literal
 import streamlit as st
-from langchain_pinecone.vectorstores import PineconeVectorStore
-from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpoint
-from langchain.prompts import PromptTemplate
-from pinecone import Pinecone #, ServerlessSpec
-from langchain_community.chat_message_histories import ChatMessageHistory
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
-from dotenv import load_dotenv
-import os
-
-# Load environment variables from the .env file
-load_dotenv()
+from utils import process
+from chat.bot import ChatBot
 
 @dataclass
 class Message:
@@ -20,81 +11,31 @@ class Message:
  origin: Literal["👤 Human", "👨🏻‍⚖️ Ai"]
  message: str
 
-
-def download_hugging_face_embeddings():
- embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
- return embeddings
-
-
 def initialize_session_state():
+ """Initialize session state variables."""
  if "history" not in st.session_state:
  st.session_state.history = []
  if "conversation" not in st.session_state:
- embeddings = download_hugging_face_embeddings()
- pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
- index = pc.Index("il-legal")
- docsearch = PineconeVectorStore.from_existing_index(index_name="il-legal", embedding=embeddings)
- 
- repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
- llm = HuggingFaceEndpoint(
- repo_id=repo_id,
- model_kwargs={"huggingface_api_token":os.getenv("HUGGINGFACEHUB_API_TOKEN")},
- temperature=0.5,
- top_k=10,
- )
-
- prompt_template = """
- You are a trained bot to guide people about Illinois Crimnal Law Statutes and the Safe-T Act. You will answer user's query with your knowledge and the context provided. 
- If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
- Do not say thank you and tell you are an AI Assistant and be open about everything.
- Use the following pieces of context to answer the users question.
- Context: {context}
- Question: {question}
- Only return the helpful answer below and nothing else.
- Helpful answer:
- """
-
- PROMPT = PromptTemplate(
- template=prompt_template, 
- input_variables=["context", "question"])
- 
- #chain_type_kwargs = {"prompt": PROMPT}
- message_history = ChatMessageHistory()
- memory = ConversationBufferMemory(
- memory_key="chat_history",
- output_key="answer",
- chat_memory=message_history,
- return_messages=True,
- )
- retrieval_chain = ConversationalRetrievalChain.from_llm(
- llm=llm,
- chain_type="stuff",
- retriever=docsearch.as_retriever(
- search_kwargs={
- 'filter': {'source': 'user_id'},
- }),
- return_source_documents=True,
- combine_docs_chain_kwargs={"prompt": PROMPT},
- memory= memory
- )
-
+ retrieval_chain, chroma_collection, langchain_chroma = ChatBot()
  st.session_state.conversation = retrieval_chain
+ st.session_state.chroma_collection = chroma_collection
+ st.session_state.langchain_chroma = langchain_chroma
+
+def on_submit(user_input):
+ """Handle user input and generate response."""
+ if user_input:
+ response = st.session_state.conversation({
+ "question":user_input
+ })
+ llm_response = response['answer']
+ st.session_state.history.append(
+ Message("🗣️ Human", user_input)
+ )
+ st.session_state.history.append(
+ Message("🧑‍⚖️ AI Lawyer", llm_response)
+ )
 
-
-def on_click_callback():
- human_prompt = st.session_state.human_prompt
- st.session_state.human_prompt=""
- response = st.session_state.conversation(
- human_prompt
- )
- llm_response = response['answer']
- st.session_state.history.append(
- Message("👤 Human", human_prompt)
- )
- st.session_state.history.append(
- Message("👨🏻‍⚖️ Ai", llm_response)
- )
-
+ st.rerun()
 
 initialize_session_state()
 
@@ -109,7 +50,7 @@ st.markdown(
  
  - Answer questions on various aspects of Illinois criminal law.
  - Guide you through legal processes relevant to Illinois.
- - Provide information on your rights and responsibilities as per Illinois legal standards.
+ - Provide information on your rights and responsibilities regarding Illinois legal standards.
  
  ⚖️ **Disclaimer:**
  
@@ -117,29 +58,52 @@ st.markdown(
  
  🤖 **Getting Started:**
  
- Feel free to ask any legal question related to Illinois law, using keywords like "pre-trial release," "motions," or "procedure." I'm here to assist you!
+ Feel free to ask any legal question related to Illinois criminal law. I'm here to assist you!
+ If you have any documents pertinent to your case to better assist you, please upload them below.
  Let's get started! How may I help you today?
  """
 )
 
 chat_placeholder = st.container()
-prompt_placeholder = st.form("chat-form")
 
 with chat_placeholder:
  for chat in st.session_state.history:
  st.markdown(f"{chat.origin} : {chat.message}")
 
-with prompt_placeholder:
- st.markdown("**Chat**")
- cols = st.columns((6, 1))
- cols[0].text_input(
- "Chat",
- label_visibility="collapsed",
- key="human_prompt",
- )
- cols[1].form_submit_button(
- "Submit",
- type="primary",
- on_click=on_click_callback,
- )
+user_question = st.chat_input("Enter your question here...")
+
+# File upload and processing
+uploaded_file = st.file_uploader("Upload your legal document", type="pdf")
+
+if uploaded_file is not None:
+ try:
+ uploaded_file.seek(0)
+ text = process.extract_text_from_pdf(uploaded_file)
+ chunks = process.chunk_text(text)
+ st.session_state.user_chunks = chunks
+ st.success(f"Uploaded {uploaded_file.name} successfully with {len(chunks)} chunks")
+
+ # Add chunks to Chroma
+ ids = [f"doc_{i}" for i in range(len(chunks))]
+ metadatas = [{"source": "user_upload"} for _ in chunks] #range(len(chunks))],
+ st.session_state.chroma_collection.add(
+ documents=chunks,
+ ids=ids,
+ metadatas=metadatas
+ )
+
+ # Add chunks to LangChain Chroma wrapper
+ st.session_state.langchain_chroma.add_texts(
+ texts=chunks,
+ metadatas=metadatas
+ )
+
+ st.success("Document processed and vectorized successfully!")
+
+ except Exception as e:
+ logging.exception(f"An error occurred while processing {uploaded_file.name}: {str(e)}")
+ st.error(f"An error occurred while processing {uploaded_file.name}: {str(e)}")
+
 
+if user_question:
+ on_submit(user_question)