How can I fix the 'st.session_state has no attribute "retriever"' error in my LangChain RAG app with Chroma?

I'm building a document Q&A assistant using LangChain with Chroma DB and Streamlit. When running the app, I encounter the following error:

An error occurred: st.session_state has no attribute "retriever". Did you forget to initialize it? More info: https://docs.streamlit.io/develop/concepts/architecture/session-state#initialization

I verified that the session state values are initialized correctly (as confirmed by my print statements), and the same app works fine without the Chroma DB integration. I'm unsure why the error occurs only when using Chroma. Could someone guide me on what might be missing or misconfigured?

Below is the complete code I'm using:

PYTHONCopied
import streamlit as st
from langchain.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.output_parsers import StrOutputParser
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.runnables import RunnablePassthrough
from typing import List
from langchain_core.documents import Document
import os

# Create a custom callback handler for streaming
class StreamHandler(BaseCallbackHandler):
    def __init__(self, container):
        self.container = container
        self.text = ""

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)

def get_chat_history(messages):
    """Format chat history from Streamlit messages for the prompt."""
    chat_history = []
    for message in messages:
        if message["role"] != "assistant":  # Skip welcome message
            chat_history.append(f"Human: {message['content']}")
        if message["role"] == "assistant" and message["content"] != "Hello, I'm running on Ollama. How can I help you today?":
            chat_history.append(f"Assistant: {message['content']}")
    return "\n".join(chat_history)

def format_docs(docs: List[Document]) -> str:
    """Format retrieved documents into a single string."""
    return "\n\n".join(doc.page_content for doc in docs)

def initialize_session_state():
    if "messages" not in st.session_state:
        st.session_state.messages = []
        st.write("Messages initialized.")

    if "llm" not in st.session_state:
        st.session_state.llm = OllamaLLM(model="llama3.2:3b", temperature=0)
        st.write("LLM initialized.")

    if "embeddings" not in st.session_state:
        try:
            st.session_state.embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
            st.write("Embeddings initialized.")
        except Exception as e:
            st.error(f"Failed to initialize embeddings: {e}")

    if "vector_store" not in st.session_state:
        try:
            st.session_state.vector_store = Chroma(
                collection_name="India_employee_handbook_collection",
                embedding_function=st.session_state.embeddings,
                persist_directory="./chroma_langchain_db"
            )
            st.write("Vector store initialized.")
        except Exception as e:
            st.error(f"Failed to initialize vector store: {e}")

    if "retriever" not in st.session_state:
        try:
            st.session_state.retriever = st.session_state.vector_store.as_retriever(
                search_type="mmr",
                search_kwargs={"k": 2, "fetch_k": 3}
            )
            st.write("Retriever initialized.")
        except Exception as e:
            st.error(f"Failed to initialize retriever: {e}")

    if "prompt" not in st.session_state:
        st.session_state.prompt = ChatPromptTemplate.from_template("""
        You are a helpful assistant. Use the following context and chat history to answer the question.
        If you don't find the answer in the context, just say that you don't know.

        Context:
        {context}

        Chat History:
        {chat_history}

        Human: {question}
        Assistant: """)
        st.write("Prompt initialized.")

    if "chain" not in st.session_state:
        try:
            st.session_state.chain = (
                {
                    "context": lambda x: format_docs(st.session_state.retriever.get_relevant_documents(x["question"])),
                    "question": lambda x: x["question"],
                    "chat_history": lambda x: x["chat_history"]
                }
                | st.session_state.prompt
                | st.session_state.llm
                | StrOutputParser()
            )
            st.write("Chain initialized.")
        except Exception as e:
            st.error(f"Failed to initialize chain: {e}")

def main():
    st.title("Document Q&A Assistant with MMR Search")

    # Ensure session state is initialized
    if "retriever" not in st.session_state:
        initialize_session_state()

    # Display welcome message if it's the first message
    if not st.session_state.messages:
        welcome_msg = "Hello, I'm a document Q&A assistant using MMR search. I can help answer questions about the employee handbook. How can I help you today?"
        st.write(welcome_msg)
        st.session_state.messages.append({"role": "assistant", "content": welcome_msg})

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input
    if prompt := st.chat_input("Ask a question about the document..."):
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})

        # Display assistant response
        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            stream_handler = StreamHandler(response_placeholder)

            try:
                # Get chat history
                chat_history = get_chat_history(st.session_state.messages[:-1])

                # Check and initialize retriever again if necessary
                if "retriever" not in st.session_state:
                    st.error("Retriever is not initialized. Reinitializing session state.")
                    initialize_session_state()

                # Get response using LCEL chain
                response = st.session_state.chain.invoke(
                    {
                        "question": prompt,
                        "chat_history": chat_history
                    },
                    config={"callbacks": [stream_handler]}
                )

                # Add assistant's response to message history
                st.session_state.messages.append({"role": "assistant", "content": response})

            except Exception as e:
                error_msg = f"An error occurred: {str(e)}"
                st.error(error_msg)
                st.session_state.messages.append({"role": "assistant", "content": error_msg})


if __name__ == "__main__":
    main()

I appreciate any insights or suggestions to resolve this attribute error effectively.