Add files from GitHub repo with SQLite files tracked by LFS

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+vectordb/**/*.sqlite3 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,53 @@
+# app.py
+import streamlit as st
+from db import init_connection
+from ui_pages import login_page, create_account_page, main_page
+from chat import load_user_sessions   # import it here
+
+from dotenv import load_dotenv
+load_dotenv()
+
+def app():
+    # Initialize session state
+    default_keys = {
+        'logged_in': False,
+        'username': "",
+        'show_create_account': False,
+        'messages': [],
+        'current_chat_session': None,
+        'chat_sessions': [],
+        'last_session_id': None
+    }
+    for key, val in default_keys.items():
+        if key not in st.session_state:
+            st.session_state[key] = val
+
+    # Initialize database connection (returns dict of collections)
+    db_conn = init_connection()
+    if db_conn is None:
+        return
+
+    # Route to appropriate page
+    if st.session_state.logged_in:
+        # load_user_sessions returns (sessions, current_session, messages_stub)
+        sessions, current, messages = load_user_sessions(
+            st.session_state.username,
+            db_conn["sessions"],
+            st.session_state.get("last_session_id")
+        )
+        st.session_state.chat_sessions = sessions
+        st.session_state.current_chat_session = current
+        # load chat messages if we have a current session
+        if current:
+            from chat import load_chat_history
+            st.session_state.messages = load_chat_history(str(current["_id"]), db_conn["messages"])
+        else:
+            st.session_state.messages = []
+        main_page()
+    elif st.session_state.show_create_account:
+        create_account_page()
+    else:
+        login_page()
+
+if __name__ == "__main__":
+    app()

auth.py

@@ -0,0 +1,26 @@
+import bcrypt
+import logging
+
+logger = logging.getLogger(__name__)
+
+def check_login(username: str, password: str, users_collection) -> bool:
+    """Checks if the provided username and password are valid against MongoDB."""
+    logger.info(f"Login attempt for user: {username}")
+    user = users_collection.find_one({"username": username})
+    if user:
+        stored_hash = user["password"]
+
+        # Ensure we always have bytes for bcrypt.checkpw
+        stored_hash_bytes = stored_hash.encode('utf-8') if isinstance(stored_hash, str) else stored_hash
+
+        try:
+            if bcrypt.checkpw(password.encode('utf-8'), stored_hash_bytes):
+                logger.info(f"User '{username}' logged in successfully.")
+                return True
+            else:
+                logger.warning(f"Invalid password attempt for user: {username}")
+        except Exception as e:
+            logger.error(f"Error checking password for user {username}: {e}")
+    else:
+        logger.warning(f"Login failed, user not found: {username}")
+    return False

chat.py

@@ -0,0 +1,61 @@
+from datetime import datetime
+from bson.objectid import ObjectId
+
+def load_user_sessions(username, sessions_collection, last_session_id=None):
+    """
+    Load sessions for a user. Restore last active session if possible.
+    Returns: (sessions_list, current_session_or_None, messages_list)
+    """
+    if sessions_collection is None:
+        return [], None, []
+
+    sessions = list(sessions_collection.find({"username": username}).sort("timestamp", -1))
+    current_session = None
+    messages = []
+
+    if sessions:
+        if last_session_id:
+            try:
+                last = sessions_collection.find_one({"_id": ObjectId(last_session_id)})
+            except Exception:
+                last = None
+            if last:
+                current_session = last
+                # messages will be loaded by caller or by calling load_chat_history
+        if not current_session:
+            current_session = sessions[0]
+
+    # Note: do NOT load messages here unless you also have messages_collection.
+    # Return sessions and current_session; caller can call load_chat_history with messages_collection.
+    return sessions, current_session, messages
+
+
+def load_chat_history(session_id, messages_collection):
+    """
+    Loads messages for a given chat session from messages_collection.
+    """
+    if messages_collection is None:
+        return []
+    try:
+        msgs = list(messages_collection.find({"session_id": session_id}).sort("timestamp", 1))
+        return [{"role": m.get("role", "assistant"), "content": m.get("content", "")} for m in msgs]
+    except Exception:
+        return []
+
+
+def save_message(session_id, role, content, messages_collection):
+    """
+    Save a message to the chat history in messages_collection.
+    """
+    if messages_collection is None:
+        return None
+    try:
+        doc = {
+            "session_id": session_id,
+            "role": role,
+            "content": content,
+            "timestamp": datetime.utcnow(),
+        }
+        return messages_collection.insert_one(doc).inserted_id
+    except Exception:
+        return None

db.py

@@ -0,0 +1,42 @@
+from pymongo import MongoClient
+import os
+
+from dotenv import load_dotenv
+load_dotenv()
+
+
+client = None
+db = None
+users_collection = None
+sessions_collection = None
+messages_collection = None
+
+
+def init_connection():
+    """Initialize MongoDB connection and collections."""
+    global client, db, users_collection, sessions_collection, messages_collection
+
+    mongo_uri = os.getenv("MONGO_URI")
+    if not mongo_uri:
+        raise ValueError("❌ MONGO_URI not found in environment variables.")
+
+    client = MongoClient(mongo_uri)
+    db = client.get_database("law_cases_db")
+
+    users_collection = db.get_collection("users")
+    sessions_collection = db.get_collection("chat_sessions")
+    messages_collection = db.get_collection("chat_messages")
+
+    # ✅ create unique index (username + normalized chat name)
+    sessions_collection.create_index(
+        [("username", 1), ("session_name_normalized", 1)],
+        unique=True
+    )
+
+    return {
+        "client": client,
+        "db": db,
+        "users": users_collection,
+        "sessions": sessions_collection,
+        "messages": messages_collection
+    }

rag.py

@@ -0,0 +1,175 @@
+import os
+import logging
+import streamlit as st
+from dotenv import load_dotenv
+import pickle
+
+from llama_index.llms.groq import Groq
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.core import VectorStoreIndex
+from llama_index.core.retrievers import VectorIndexRetriever, RecursiveRetriever
+from llama_index.retrievers.bm25 import BM25Retriever
+from llama_index.core.tools import QueryEngineTool
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core import get_response_synthesizer
+from llama_index.core.agent import ReActAgent
+from chromadb import PersistentClient
+
+logger = logging.getLogger(__name__)
+
+@st.cache_resource
+def setup_rag_system(debug=False):
+    load_dotenv()
+    
+    groq_api_key = os.getenv("GROQ_API_KEY") or st.secrets.get("groq", {}).get("api_key")
+    if not groq_api_key:
+        st.error("GROQ API key not found. Please check your environment variables or secrets.")
+        st.stop()
+
+    # LLM
+    llm = Groq(
+        model="llama-3.1-8b-instant",
+        api_key=groq_api_key,
+        max_input_tokens=1200,
+        max_output_tokens=1200
+    )
+
+    # Embeddings
+    embedding_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    # Persisted vector DBs
+    persist_dirs = [
+        "vectordb/case_2021",
+        "vectordb/case_2022",
+        "vectordb/case_2023",
+        "vectordb/case_2024",
+        "vectordb/case_2025"
+    ]
+    for persist_dir in persist_dirs:
+        if not os.path.exists(persist_dir):
+            st.error(f"Vector database directory {persist_dir} not found.")
+            st.stop()
+
+    # Build hybrid retrievers
+    hybrid_retrievers = []
+    for persist_dir in persist_dirs:
+        # Load pickled nodes
+        nodes_path = os.path.join(persist_dir, "nodes.pkl")
+        if not os.path.exists(nodes_path):
+            st.error(f"Pickle file {nodes_path} not found.")
+            st.stop()
+
+        with open(nodes_path, "rb") as f:
+            nodes = pickle.load(f)
+
+        # Vector store
+        client = PersistentClient(path=persist_dir)
+        collection = client.get_collection("case_collection")
+        vector_store = ChromaVectorStore(chroma_collection=collection)
+        index = VectorStoreIndex.from_vector_store(vector_store=vector_store, embed_model=embedding_model)
+
+        # Retrievers
+        vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=2, retriever_mode="mmr")
+        bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=2)
+
+        hybrid_retriever = RecursiveRetriever(
+            "vector",
+            retriever_dict={"vector": vector_retriever, "bm25": bm25_retriever},
+            verbose=True
+        )
+        hybrid_retrievers.append(hybrid_retriever)
+
+    # Case metadata
+    documents_info = [
+        {
+            "name": "Quezada2021_Retriever",
+            "description": "Retrieves information from the United States Court of Appeals for the Armed Forces decision in United States v. Quezada (21-0089-MC), issued on December 20, 2021."
+        },
+        {
+            "name": "Thompson2022_Retriever",
+            "description": "Retrieves information from the United States Court of Appeals for the Armed Forces decision in United States v. Thompson (22-0098-AF), issued on November 21, 2022."
+        },
+        {
+            "name": "Brown2023_Retriever",
+            "description": "Retrieves information from the United States Court of Appeals for the Armed Forces decision in United States v. Brown (22-0249-CG), issued on October 23, 2023."
+        },
+        {
+            "name": "Smith2024_Retriever",
+            "description": "Retrieves information from the United States Court of Appeals for the Armed Forces decision in United States v. Smith (23-0207-AF), issued on November 26, 2024."
+        },
+        {
+            "name": "Lopez2025_Retriever",
+            "description": "Retrieves information from the United States Court of Appeals for the Armed Forces decision in United States v. Lopez (24-0226-CG), issued on September 2, 2025."
+        },
+    ]
+
+
+    # Create retriever → tool
+    def create_retriever_tool(retriever, llm, name, description):
+        response_synthesizer = get_response_synthesizer(
+            llm=llm, response_mode="compact", use_async=False
+        )
+        query_engine = RetrieverQueryEngine(retriever=retriever, response_synthesizer=response_synthesizer)
+        return QueryEngineTool.from_defaults(query_engine=query_engine, name=name, description=description)
+
+    retriever_tools = [
+        create_retriever_tool(hybrid_retrievers[i], llm, info["name"], info["description"])
+        for i, info in enumerate(documents_info)
+    ]
+
+    # System prompt
+    system_prompt = """
+        You are a highly specialized legal research assistant. 
+        You may ONLY answer questions that are legal in nature. 
+        This includes both:
+        - Specific case law queries from the provided case documents (2021–2025).
+        - General legal concepts, doctrines, or terminology.
+
+        Before answering, always perform this intermediate reasoning step:
+
+        1. Classify the user query:
+        - If the query relates to law, legal concepts, legal systems, court rulings, rights, duties, contracts, procedures, or legal doctrines → classify as: LEGAL_QUERY.
+        - If the query is casual conversation, mathematics, trivia, technical programming, or anything outside the legal domain → classify as: NON_LEGAL_QUERY.
+
+        2. Response rules:
+        - If LEGAL_QUERY:
+            a) If the query references specific cases between 2021–2025, use the provided case documents to retrieve and answer. Cite the case name and year.
+            b) If the query is a general legal question, answer concisely and professionally, using legal reasoning. Do NOT speculate beyond standard legal knowledge.
+        - If NON_LEGAL_QUERY:
+            Respond ONLY with: "I can only answer questions about legal cases (2021–2025) or general law queries."
+
+        3. Examples:
+        - LEGAL_QUERY (answer these):
+            • "What is the difference between civil and criminal law?"
+            • "Explain the principle of judicial review."
+            • "Summarize the ruling in United States v. Lopez (2025)."
+            • "What is mens rea in criminal law?"
+        - NON_LEGAL_QUERY (reject these):
+            • "What is 2+2?"
+            • "Who won the FIFA World Cup in 2022?"
+            • "Write me a Python script."
+            • "Tell me a joke."
+
+        4. Style & tone:
+        - Be concise, professional, and clear.
+        - Use citations ONLY when referring to case documents (case name + year).
+        - Never provide speculative or non-legal answers.
+        """
+
+
+    # ReActAgent
+    agent = ReActAgent(
+        tools=retriever_tools,
+        llm=llm,
+        verbose=True,
+        max_iterations=20,
+        system_prompt=system_prompt
+    )
+
+    logger.info("RAG system setup complete.")
+
+    if debug:
+        return agent, llm, hybrid_retrievers
+
+    return agent, llm

requirements.txt

@@ -1,3 +1,17 @@
-altair
-pandas
-streamlit
+requests
+PyMuPDF
+chromadb
+sentence-transformers
+
+llama-index
+llama-index-vector-stores-chroma
+llama-index-embeddings-huggingface
+llama-index-llms-groq
+llama-index
+llama-index-retrievers-bm25
+llama-index-storage-chat-store-mongo
+
+pymongo
+
+streamlit
+nest_asyncio

ui_pages.py

@@ -0,0 +1,577 @@
+import time
+import streamlit as st
+import re
+import asyncio
+import nest_asyncio
+
+from auth import check_login
+from chat import load_user_sessions, load_chat_history, save_message
+from rag import setup_rag_system
+from db import init_connection
+
+nest_asyncio.apply()
+
+
+def login_page():
+    # Custom CSS for login page
+    st.markdown("""
+    <style>
+    .login-container {
+        max-width: 450px;
+        margin: 2rem auto;
+        padding: 2rem;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border-radius: 15px;
+        box-shadow: 0 10px 30px rgba(0,0,0,0.3);
+    }
+    .login-title {
+        color: white;
+        text-align: center;
+        font-size: 2.5rem;
+        margin-bottom: 1rem;
+        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+    }
+    .login-caption {
+        color: #e0e0e0;
+        text-align: center;
+        margin-bottom: 2rem;
+        font-size: 1.1rem;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    st.markdown("""
+    <div class="login-container">
+        <h1 class="login-title">Legal Case RAG</h1>
+        <p class="login-caption">Ask me about legal cases (2021–2025). I'll retrieve documents and give citations.</p>
+    </div>
+    """, unsafe_allow_html=True)
+
+    # Center the form
+    col1, col2, col3 = st.columns([1, 2, 1])
+    
+    with col2:
+        with st.form("login_form"):
+            st.markdown("### Login to Continue")
+            username = st.text_input("Username", placeholder="Enter your username")
+            password = st.text_input("Password", type="password", placeholder="Enter your password")
+            login_button = st.form_submit_button("🚀 Login", use_container_width=True)
+
+            if login_button:
+                db_conn = init_connection()
+                users_collection = db_conn["users"]
+
+                if check_login(username, password, users_collection):
+                    st.session_state.logged_in = True
+                    st.session_state.username = username
+                    st.success(f"Welcome back, {username}!")
+                    time.sleep(1)
+                    st.rerun()
+                else:
+                    st.error("Invalid username or password.")
+
+        st.markdown("---")
+        if st.button("Create an account", use_container_width=True):
+            st.session_state.show_create_account = True
+            st.rerun()
+
+
+def create_account_page():
+    # Custom CSS for create account page
+    st.markdown("""
+    <style>
+    .signup-container {
+        max-width: 450px;
+        margin: 2rem auto;
+        padding: 2rem;
+        background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
+        border-radius: 15px;
+        box-shadow: 0 10px 30px rgba(0,0,0,0.3);
+    }
+    .signup-title {
+        color: white;
+        text-align: center;
+        font-size: 2.5rem;
+        margin-bottom: 1rem;
+        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+    }
+    .signup-caption {
+        color: #e0e0e0;
+        text-align: center;
+        margin-bottom: 2rem;
+        font-size: 1.1rem;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    st.markdown("""
+    <div class="signup-container">
+        <h1 class="signup-title">Create Account</h1>
+        <p class="signup-caption">Set up your account to start chatting with the legal RAG system.</p>
+    </div>
+    """, unsafe_allow_html=True)
+
+    db_conn = init_connection()
+    users_collection = db_conn["users"]
+
+    col1, col2, col3 = st.columns([1, 2, 1])
+    
+    with col2:
+        with st.form("create_account_form"):
+            st.markdown("### Account Details")
+            new_username = st.text_input("Choose a Username", placeholder="Enter desired username")
+            new_password = st.text_input("Choose a Password", type="password", placeholder="Enter secure password")
+            create_button = st.form_submit_button("Create Account", use_container_width=True)
+
+            if create_button:
+                import bcrypt
+                if not new_username or not new_password:
+                    st.error("Username and password cannot be empty.")
+                elif users_collection.find_one({"username": new_username}):
+                    st.error("Username already exists.")
+                else:
+                    salt = bcrypt.gensalt()
+                    hashed_password = bcrypt.hashpw(new_password.encode('utf-8'), salt)
+
+                    users_collection.insert_one({
+                        "username": new_username,
+                        "password": hashed_password.decode('utf-8')
+                    })
+                    st.success("Account created successfully! Please log in.")
+                    st.session_state.show_create_account = False
+                    st.rerun()
+
+        if st.button("Back to Login", use_container_width=True):
+            st.session_state.show_create_account = False
+            st.rerun()
+
+
+def main_page():
+    # Enhanced CSS for chat interface
+    st.markdown("""
+    <style>
+    /* Main chat container */
+    .chat-container {
+        max-width: 1200px;
+        margin: 0 auto;
+    }
+    
+    /* User message styling */
+    .user-message {
+        display: flex;
+        justify-content: flex-end;
+        margin-bottom: 1rem;
+    }
+    
+    .user-bubble {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        padding: 12px 16px;
+        border-radius: 18px 18px 4px 18px;
+        max-width: 70%;
+        box-shadow: 0 2px 10px rgba(102, 126, 234, 0.3);
+        position: relative;
+    }
+    
+    .user-bubble::after {
+        content: '';
+        position: absolute;
+        bottom: 0;
+        right: -8px;
+        width: 0;
+        height: 0;
+        border: 8px solid transparent;
+        border-top-color: #764ba2;
+        border-bottom: 0;
+        margin-left: -8px;
+        margin-bottom: -8px;
+    }
+    
+    /* Assistant message styling */
+    .assistant-message {
+        display: flex;
+        justify-content: flex-start;
+        margin-bottom: 1rem;
+    }
+    
+    /* Assistant message styling */
+    .assistant-bubble {
+        background: linear-gradient(135deg, #06beb6 0%, #48b1bf 100%) !important;
+        color: white;
+        padding: 12px 16px;
+        border-radius: 18px 18px 18px 4px;
+        max-width: 70%;
+        box-shadow: 0 2px 10px rgba(72, 177, 191, 0.3);
+        position: relative;
+    }
+    
+    .assistant-bubble::after {
+        content: '';
+        position: absolute;
+        bottom: 0;
+        left: -8px;
+        width: 0;
+        height: 0;
+        border: 8px solid transparent;
+        border-top-color: #48b1bf;
+        border-bottom: 0;
+        margin-right: -8px;
+        margin-bottom: -8px;
+    }
+    
+    /* Avatar styling */
+    .user-avatar {
+        width: 40px;
+        height: 40px;
+        border-radius: 50%;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        color: white;
+        font-weight: bold;
+        margin-left: 10px;
+        margin-top: 5px;
+    }
+    
+    .assistant-avatar {
+        width: 40px;
+        height: 40px;
+        border-radius: 50%;
+        background: linear-gradient(135deg, #06beb6 0%, #48b1bf 100%) !important;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        color: white;
+        font-weight: bold;
+        margin-right: 10px;
+        margin-top: 5px;
+    }
+    
+    /* Sidebar styling */
+    .sidebar-header {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        padding: 1rem;
+        border-radius: 10px;
+        margin-bottom: 1rem;
+        text-align: center;
+    }
+    
+    /* Session button styling */
+    .session-btn {
+        width: 100%;
+        margin-bottom: 5px;
+        padding: 8px;
+        border-radius: 8px;
+        border: none;
+        background: #f0f2f6;
+        cursor: pointer;
+        transition: all 0.3s ease;
+    }
+    
+    .session-btn:hover {
+        background: #e0e0e0;
+        transform: translateY(-2px);
+    }
+    
+    /* Welcome message */
+    .welcome-header {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        padding: 1.5rem;
+        border-radius: 15px;
+        margin-bottom: 1rem;
+        text-align: center;
+        box-shadow: 0 5px 20px rgba(102, 126, 234, 0.3);
+    }
+    
+    /* Chat title */
+    .chat-title {
+        background: linear-gradient(135deg, #06beb6 0%, #48b1bf 100%);
+        color: white;
+        padding: 1rem;
+        border-radius: 10px;
+        margin-bottom: 1rem;
+        text-align: center;
+    }
+    
+    /* Source metadata styling */
+    .source-metadata {
+        background: rgba(72, 177, 191, 0.1);
+        border-left: 4px solid #48b1bf;
+        padding: 10px;
+        margin-top: 10px;
+        border-radius: 5px;
+    }
+    """, unsafe_allow_html=True)
+
+    # Initialize chat memory for RAG agent
+    if 'chat_memory' not in st.session_state:
+        st.session_state['chat_memory'] = []
+
+    if 'agent' not in st.session_state:
+        with st.spinner("🔄 Loading legal case documents..."):
+            agent, llm = setup_rag_system()
+            st.session_state.agent = agent
+            st.session_state.llm = llm
+
+    # Enhanced sidebar
+    with st.sidebar:
+        st.markdown("""
+        <div class="sidebar-header">
+            <h3>Chat Sessions</h3>
+        </div>
+        """, unsafe_allow_html=True)
+
+        db_conn = init_connection()
+        sessions_collection = db_conn["sessions"]
+        messages_collection = db_conn["messages"]
+
+        if 'chat_sessions' not in st.session_state or not st.session_state.chat_sessions:
+            sessions, current, messages = load_user_sessions(
+                st.session_state.username,
+                sessions_collection,
+                last_session_id=st.session_state.get("last_session_id")
+            )
+            st.session_state.chat_sessions = sessions
+            st.session_state.current_chat_session = current
+            if current:
+                st.session_state.messages = load_chat_history(str(current["_id"]), messages_collection)
+            else:
+                st.session_state.messages = []
+
+        # Create new session form
+        with st.form("new_chat_form"):
+            st.session_state.chat_memory = []
+            st.markdown("### Start New Chat")
+            session_name = st.text_input("Chat name:", placeholder="e.g., Contract Law Questions", key="new_chat_name")
+            if st.form_submit_button("Create", use_container_width=True):
+                session_name = session_name.strip()
+                if not session_name:
+                    st.error("Chat name cannot be empty.")
+                else:
+                    normalized = session_name.lower()
+                    existing = sessions_collection.find_one({
+                        "username": st.session_state.username,
+                        "session_name_normalized": normalized
+                    })
+                    if existing:
+                        st.error(f"You already have a chat named '{existing['session_name']}'. Please choose a different name.")
+                    else:
+                        new_session = {
+                            "username": st.session_state.username,
+                            "session_name": session_name,
+                            "timestamp": time.time(),
+                            "session_name_normalized": normalized
+                        }
+                        inserted_id = sessions_collection.insert_one(new_session).inserted_id
+                        st.session_state.chat_sessions.append({
+                            "_id": inserted_id,
+                            "session_name": session_name,
+                            "username": st.session_state.username,
+                            "timestamp": new_session["timestamp"]
+                        })
+                        st.session_state.current_chat_session = st.session_state.chat_sessions[-1]
+                        st.session_state.messages = []
+                        st.session_state.last_session_id = str(inserted_id)
+                        st.success(f"New chat '{session_name}' created!")
+                        time.sleep(1)
+                        st.rerun()
+                        
+
+        # Previous sessions
+        st.markdown("---")
+        st.markdown("### Previous Sessions")
+        for session in st.session_state.chat_sessions:
+            if st.button(f"{session['session_name']}", key=str(session["_id"]), use_container_width=True):
+                st.session_state.current_chat_session = session
+                st.session_state.messages = load_chat_history(str(session["_id"]), messages_collection)
+                st.session_state.last_session_id = str(session["_id"])
+                st.rerun()
+
+        # Logout button
+        st.markdown("---")
+        if st.button("Logout", use_container_width=True):
+            keys_to_clear = ['logged_in', 'username', 'current_chat_session', 'chat_sessions', 'messages', 'agent', 'llm']
+            for key in keys_to_clear:
+                st.session_state.pop(key, None)
+            st.rerun()
+        
+        st.markdown("---")
+        # Instructions expander
+        with st.expander("ℹ️ How to use this chatbot", expanded=False):
+            st.markdown("""
+            **How it works:**
+            1. Type your **legal question** in the chat box below  
+            2. I'll retrieve relevant cases from **2021–2025** and respond with context + citations  
+            3. Use the sidebar to **start a new chat** or switch between past sessions  
+            4. Chat titles are what you enter when creating a session  
+            5. Use **Logout** in the sidebar to securely end your session
+            
+            **Tips for better results:**
+            - Be specific about the legal area (e.g., "contract law", "criminal procedure")
+            - Include relevant case names if you know them
+            - Ask follow-up questions to dive deeper into specific aspects
+            """)
+
+    # Main content area
+    if st.session_state.current_chat_session:
+        # Welcome header
+        st.markdown(f"""
+        <div style="padding: 1rem; border: 0px solid #ddd; border-radius: 8px; text-align: center; margin-bottom: 1rem;">
+            <h2 style="margin-bottom: 0.5rem;">Legal Case RAG Chatbot</h2>
+            <p style="margin: 0;">Welcome, <strong>{st.session_state.username}</strong>! Create a new chat session to get started.</p>
+        </div>
+        """, unsafe_allow_html=True)
+
+        # Current chat title
+        # Current chat title (minimalist text)
+        st.markdown(
+            f"<p style='font-size: 1.1rem;'>You are opened <strong>{st.session_state.current_chat_session['session_name']}</strong> chat!</p>",
+            unsafe_allow_html=True
+        )
+
+
+        # Chat messages container
+        chat_container = st.container()
+        
+        with chat_container:
+            # Display previous messages with enhanced styling
+            for i, message in enumerate(st.session_state.messages):
+                if message["role"] == "user":
+                    st.markdown(f"""
+                    <div class="user-message">
+                        <div class="user-bubble">
+                            {message["content"]}
+                        </div>
+                        <div class="user-avatar">
+                            {st.session_state.username[0].upper()}
+                        </div>
+                    </div>
+                    """, unsafe_allow_html=True)
+                else:
+                    st.markdown(f"""
+                    <div class="assistant-message">
+                        <div class="assistant-avatar">
+                            🤖
+                        </div>
+                        <div class="assistant-bubble">
+                            {message["content"]}
+                        </div>
+                    </div>
+                    """, unsafe_allow_html=True)
+
+        # Chat input
+        if prompt := st.chat_input("Ask a question about the cases...", key="chat_input"):
+            # Display user message immediately
+            st.markdown(f"""
+            <div class="user-message">
+                <div class="user-bubble">
+                    {prompt}
+                </div>
+                <div class="user-avatar">
+                    {st.session_state.username[0].upper()}
+                </div>
+            </div>
+            """, unsafe_allow_html=True)
+            
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            st.session_state.chat_memory.append({"role": "user", "content": prompt})
+
+            # Get last 4 messages for context
+            recent_messages = st.session_state.chat_memory[-4:]
+            context_str = "\n".join([f"{m['role']}: {m['content']}" for m in recent_messages])
+            prompt_for_agent = st.session_state.agent.system_prompt + "\n" + context_str
+
+            # Show thinking indicator
+            thinking_placeholder = st.empty()
+            thinking_placeholder.markdown("""
+            <div class="assistant-message">
+                <div class="assistant-avatar">
+                    🤖
+                </div>
+                <div class="assistant-bubble">
+                    🤔 Thinking and searching through legal documents...
+                </div>
+            </div>
+            """, unsafe_allow_html=True)
+
+            try:
+                async def ask_agent(agent, prompt_for_agent):
+                    return await agent.run(prompt_for_agent, max_iterations=20)
+
+                response = asyncio.run(ask_agent(st.session_state.agent, prompt_for_agent))
+
+                # Clear thinking indicator
+                thinking_placeholder.empty()
+
+                # Main answer text
+                text_output = str(response.response) if hasattr(response, "response") else str(response)
+                
+                # Display assistant response
+                st.markdown(f"""
+                <div class="assistant-message">
+                    <div class="assistant-avatar">
+                        🤖
+                    </div>
+                    <div class="assistant-bubble">
+                        {text_output}
+                    </div>
+                </div>
+                """, unsafe_allow_html=True)
+
+                # Display source metadata
+                if hasattr(response, "tool_calls") and response.tool_calls:
+                    st.markdown("### Sources & References")
+                    seen = set()
+                    for tool_call in response.tool_calls:
+                        raw_output = getattr(tool_call.tool_output, "raw_output", None)
+                        if raw_output and hasattr(raw_output, "source_nodes"):
+                            for node in raw_output.source_nodes:
+                                meta = getattr(node, "metadata", {})
+                                if meta:
+                                    meta_key = tuple(sorted(meta.items()))
+                                    if meta_key in seen:
+                                        continue
+                                    seen.add(meta_key)
+                                    with st.expander(f"{meta.get('case_name', 'Source Metadata')}"):
+                                        st.markdown('<div class="source-metadata">', unsafe_allow_html=True)
+                                        for k, v in meta.items():
+                                            st.markdown(f"**{k.replace('_', ' ').title()}:** {v}")
+                                        st.markdown('</div>', unsafe_allow_html=True)
+
+                # Save messages
+                st.session_state.messages.append({"role": "assistant", "content": text_output})
+                st.session_state.chat_memory.append({"role": "assistant", "content": text_output})
+                save_message(str(st.session_state.current_chat_session['_id']), "user", prompt, messages_collection)
+                save_message(str(st.session_state.current_chat_session['_id']), "assistant", text_output, messages_collection)
+
+            except Exception as e:
+                thinking_placeholder.empty()
+                error_msg = f"Error: {str(e)}"
+                st.markdown(f"""
+                <div class="assistant-message">
+                    <div class="assistant-avatar">
+                        🤖
+                    </div>
+                    <div class="assistant-bubble">
+                        {error_msg}
+                    </div>
+                </div>
+                """, unsafe_allow_html=True)
+                
+                st.session_state.messages.append({"role": "assistant", "content": error_msg})
+                st.session_state.chat_memory.append({"role": "assistant", "content": error_msg})
+
+    else:
+        # Welcome header (minimalist)
+        st.markdown(f"""
+        <div style="padding: 1rem; border: 0px solid #ddd; border-radius: 8px; text-align: center; margin-bottom: 1rem;">
+            <h2 style="margin-bottom: 0.5rem;">⚖️ Legal Case RAG Chatbot</h2>
+            <p style="margin: 0;">Welcome, <strong>{st.session_state.username}</strong>! Create a new chat session to get started.</p>
+        </div>
+        """, unsafe_allow_html=True)
+
+        
+        st.info("Use the sidebar to create a new chat session or select from your previous conversations.")

vectordb/case_2021/bd4665fb-a7e9-4393-b81c-df50acadc5cf/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8146ecc3e4c3a36ea9b3edc3778630c452f483990ec942d38e8006f4661e430
+size 16760000

vectordb/case_2021/bd4665fb-a7e9-4393-b81c-df50acadc5cf/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f1e924efbb5e1af5201e3fbab86a97f5c195c311abe651eeec525884e5e449
+size 100

vectordb/case_2021/bd4665fb-a7e9-4393-b81c-df50acadc5cf/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5ca7fcde8f398f682e3376c51b79eb941307eec087e2abd19f05aac6de7bffd
+size 40000

vectordb/case_2021/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:691cb71ea2a3a57ef49708fb68fd9b4043b80735fa6e907b97cfc3d4a8d92042
+size 532480

vectordb/case_2021/docstore.json

@@ -0,0 +1 @@
+{}

vectordb/case_2021/graph_store.json

@@ -0,0 +1 @@
+{"graph_dict": {}}

vectordb/case_2021/image__vector_store.json

@@ -0,0 +1 @@
+{"embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {}}

vectordb/case_2021/index_store.json

@@ -0,0 +1 @@
+{"index_store/data": {"d6d79df6-f3a1-43ff-9efd-c9c4bfe91f5c": {"__type__": "vector_store", "__data__": "{\"index_id\": \"d6d79df6-f3a1-43ff-9efd-c9c4bfe91f5c\", \"summary\": null, \"nodes_dict\": {}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"}}}

vectordb/case_2021/nodes.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd68e5695476b3e99d73393b06902d05013080e0ab13e99cdc5379b6dd08fa34
+size 31887

vectordb/case_2022/be5602b6-82eb-4aaa-a2ab-d74d0808b2bc/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8146ecc3e4c3a36ea9b3edc3778630c452f483990ec942d38e8006f4661e430
+size 16760000

vectordb/case_2022/be5602b6-82eb-4aaa-a2ab-d74d0808b2bc/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f1e924efbb5e1af5201e3fbab86a97f5c195c311abe651eeec525884e5e449
+size 100

vectordb/case_2022/be5602b6-82eb-4aaa-a2ab-d74d0808b2bc/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:545e678980d704088bda20501b0c9e144292247126f7e823b28b8ce3c0977acc
+size 40000

vectordb/case_2022/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7db436c8ebb2a942fc8f8a7ed7b07a163761c997834cca4cf3d411ad6f55a1b8
+size 442368

vectordb/case_2022/docstore.json

@@ -0,0 +1 @@
+{}

vectordb/case_2022/graph_store.json

@@ -0,0 +1 @@
+{"graph_dict": {}}

vectordb/case_2022/image__vector_store.json

@@ -0,0 +1 @@
+{"embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {}}

vectordb/case_2022/index_store.json

@@ -0,0 +1 @@
+{"index_store/data": {"5affab20-5ff2-473c-a2e0-cbd1804a6d80": {"__type__": "vector_store", "__data__": "{\"index_id\": \"5affab20-5ff2-473c-a2e0-cbd1804a6d80\", \"summary\": null, \"nodes_dict\": {}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"}}}

vectordb/case_2022/nodes.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:827273e48edb92530726dbef022828e0c3e740073c1e0f3296287721e6fb125d
+size 22831

vectordb/case_2023/402be202-644a-4979-9c56-fb043affb335/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8146ecc3e4c3a36ea9b3edc3778630c452f483990ec942d38e8006f4661e430
+size 16760000

vectordb/case_2023/402be202-644a-4979-9c56-fb043affb335/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f1e924efbb5e1af5201e3fbab86a97f5c195c311abe651eeec525884e5e449
+size 100

vectordb/case_2023/402be202-644a-4979-9c56-fb043affb335/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d883f9e5aeec6c1c57dd6cdd3c593159a52497fdea9b5d50bb7fb59dccdf4efd
+size 40000

vectordb/case_2023/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c163cb56c1a8fe621b46a25b40cbe8569af276b31bc6518442742f7bc018ba34
+size 1077248

vectordb/case_2023/docstore.json

@@ -0,0 +1 @@
+{}

vectordb/case_2023/graph_store.json

@@ -0,0 +1 @@
+{"graph_dict": {}}

vectordb/case_2023/image__vector_store.json

@@ -0,0 +1 @@
+{"embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {}}

vectordb/case_2023/index_store.json

@@ -0,0 +1 @@
+{"index_store/data": {"7631bb0c-d9b1-4913-9816-643cdbdd44ab": {"__type__": "vector_store", "__data__": "{\"index_id\": \"7631bb0c-d9b1-4913-9816-643cdbdd44ab\", \"summary\": null, \"nodes_dict\": {}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"}}}

vectordb/case_2023/nodes.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc1fa857bd3db4ea56adb79b18b8589ce7477780320a8cfbd7914fc659350140
+size 67724

vectordb/case_2024/17c59c4b-b633-4586-a7de-8dd2d7c5f509/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8146ecc3e4c3a36ea9b3edc3778630c452f483990ec942d38e8006f4661e430
+size 16760000

vectordb/case_2024/17c59c4b-b633-4586-a7de-8dd2d7c5f509/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f1e924efbb5e1af5201e3fbab86a97f5c195c311abe651eeec525884e5e449
+size 100

vectordb/case_2024/17c59c4b-b633-4586-a7de-8dd2d7c5f509/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:385f4459c65a661a0e4dc359eec0cb51586519c40e85b19f10d44c772717e9c1
+size 40000

vectordb/case_2024/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ee79f054daf27d865a6111370b16f3c23b0cbad1cdc736cc28787809f546165
+size 622592

vectordb/case_2024/docstore.json

@@ -0,0 +1 @@
+{}

vectordb/case_2024/graph_store.json

@@ -0,0 +1 @@
+{"graph_dict": {}}

vectordb/case_2024/image__vector_store.json

@@ -0,0 +1 @@
+{"embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {}}

vectordb/case_2024/index_store.json

@@ -0,0 +1 @@
+{"index_store/data": {"56f9663b-c71a-4880-b449-3a5b73b9b2a1": {"__type__": "vector_store", "__data__": "{\"index_id\": \"56f9663b-c71a-4880-b449-3a5b73b9b2a1\", \"summary\": null, \"nodes_dict\": {}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"}}}

vectordb/case_2024/nodes.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f8520ad85efc803e2a451799d781d83ad44f0ccd4edee9334710cd4614ef794
+size 40056

vectordb/case_2025/b5602bd0-a583-4091-8690-a44e83836a6d/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8146ecc3e4c3a36ea9b3edc3778630c452f483990ec942d38e8006f4661e430
+size 16760000

vectordb/case_2025/b5602bd0-a583-4091-8690-a44e83836a6d/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f1e924efbb5e1af5201e3fbab86a97f5c195c311abe651eeec525884e5e449
+size 100