|
import logging |
|
from dataclasses import dataclass |
|
from typing import Literal |
|
import streamlit as st |
|
from utils import process |
|
from chat.bot import ChatBot |
|
|
|
@dataclass |
|
class Message: |
|
"""Class for keeping track of a chat message.""" |
|
origin: Literal["👤 Human", "👨🏻⚖️ Ai"] |
|
message: str |
|
|
|
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: |
|
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) |
|
) |
|
|
|
st.rerun() |
|
|
|
initialize_session_state() |
|
|
|
st.title("IL-Legal Advisor Chatbot") |
|
|
|
st.markdown( |
|
""" |
|
👋 **Welcome to IL-Legal Advisor!** |
|
I'm here to assist you with your legal queries within the framework of Illinois criminal law. Whether you're navigating through specific legal issues or seeking general advice, I'm here to help. |
|
|
|
📚 **How I Can Assist:** |
|
|
|
- Answer questions on various aspects of Illinois criminal law. |
|
- Guide you through legal processes relevant to Illinois. |
|
- Provide information on your rights and responsibilities regarding Illinois legal standards. |
|
|
|
⚖️ **Disclaimer:** |
|
|
|
While I can provide general information, it may be necessary to consult with a qualified Illinois attorney for advice tailored to your specific situation. |
|
|
|
🤖 **Getting Started:** |
|
|
|
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() |
|
|
|
with chat_placeholder: |
|
for chat in st.session_state.history: |
|
st.markdown(f"{chat.origin} : {chat.message}") |
|
|
|
user_question = st.chat_input("Enter your question here...") |
|
|
|
|
|
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") |
|
|
|
|
|
ids = [f"doc_{i}" for i in range(len(chunks))] |
|
metadatas = [{"source": "user_upload"} for _ in chunks] |
|
st.session_state.chroma_collection.add( |
|
documents=chunks, |
|
ids=ids, |
|
metadatas=metadatas |
|
) |
|
|
|
|
|
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) |
|
|
