Tutorial: How To Chat with Your Notion Workspace Using Langchain

This tutorial will outline how you can chat with your Notion workspace in under 100 lines of code using Langchain and Streamlit.

We will:

1. Load the Notion workspace
2. Split the contents into chunks
3. Create embeddings with OpenAI
4. Store those chunks & embeddings in a vector database
5. Create our question-answer LLM chain and setup a cache
6. Put everything in a simple Streamlit UI

Installing Requirements

We will build our app with Python. Here's a list of requirements, start off by installing them. Create your requirements.txt file and install it via pip:

langchain==0.0.152
openai
tiktoken
streamlit_chat
chromadb
streamlit-extras

pip install -r requirements.txt

Also, here are all the imports you will need:

from langchain import LLMChain
from langchain.cache import SQLiteCache
from langchain.chains import RetrievalQA
from langchain.document_loaders import NotionDirectoryLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter 
from langchain.vectorstores import Chroma
from streamlit_chat import message
import langchain
import os
import streamlit as st

An Intro to Streamlit Session State

We are using Streamlit to build our UI. The state is stored in a map called st.session_state. We'll be using this throughout the whole code, but before that we need to initialise each state variable like this:

def init_session_state():
    if "docs" not in st.session_state:
        st.session_state["docs"] = None

    if "db" not in st.session_state:
        st.session_state["db"] = None

    if "generated" not in st.session_state:
        st.session_state["generated"] = []

    if "past" not in st.session_state:
        st.session_state["past"] = []

Exporting and Loading Notion Workspace

To load our Notion workspace, we first need to export it. Go to the root page of your Notion workspace and start exporting. Make sure to choose the following sections:

Then unzip the directory and rename it to "Notion_DB" - and we're ready to load the data and split it up:

1. We only need to load the data if we haven't already done so (checking the session state).
2. Once we've loaded the documents with NotionDirectoryLoader, we need to split them up into smaller chunks. This is because we can only send along a limited amount of context to OpenAI.
3. We do the splitting up with the RecursiveCharacterTextSplitter. You can play around with the CHUNK_SIZE and CHUNK_OVERLAP variables to your liking - for reference, I've set them to 768 and 24 respectively.
4. Finally, we store the chunks in the Streamlit session state.

def load_notion_database():
    if st.session_state["docs"] is not None:
        return

    loader = NotionDirectoryLoader("Notion_DB")

    try:
        documents     = loader.load()
    except:
        return None

    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=globals.CHUNK_SIZE, chunk_overlap=globals.CHUNK_OVERLAP)
    texts         = text_splitter.split_documents(documents)

    st.session_state["docs"]   = texts

Create Embeddings and Chroma Vectorstorage

Next, we need to create embeddings and store them in a so-called vectorstorage. We'll do that lazily once the first chatbot query comes in.

1. Make sure you set your OPENAI_API_KEY_QUERY and DATA_HOME variables. If you don't have access to OpenAI you can alternatively use HuggingFace embeddings - for more, follow the Langchain docs here. DATA_HOME is just the folder where the Chroma vectorstorage should be stored.
2. Similar to before, we only need to do this if the vectorstorage database hasn't been loaded into the session state.
3. We only create and store the embeddings once to save on OpenAI costs. We do this by checking whether the DATA_HOME folder exists. If you want to re-create the embeddings (eg because you have a newer Notion export), you'll need to delete the DATA_HOME folder.
4. If the DATA_HOME folder doesn't exist, we're creating the Chroma vectorstorage from the chunks that we created before, making use of the OpenAI embeddings. Then we persist the storage to the disk...
5. ... to reload from on the next run of the application.

def query_index(query):
    if st.session_state.db is None:
        embeddings = OpenAIEmbeddings(openai_api_key=globals.OPENAI_API_KEY_QUERY)

        if os.path.exists(globals.DATA_HOME):
            st.session_state["db"] = Chroma(embedding_function=embeddings, 
                                            persist_directory=globals.DATA_HOME)
        else:
            st.session_state["db"] = Chroma.from_documents(st.session_state.docs, 
                                                           embeddings,
                                                           persist_directory=globals.DATA_HOME)
            st.session_state["db"].persist()

Querying Notion with RetrievalQA and Setting Up a Cache for OpenAI Requests

Next, we will continue the query_index method and send our query to OpenAI with the RetrievalQA chain:

1. Use the Chroma vectorstorage as retriever.
2. Create the RetrievalQA chain. Essentially, this chain will - upon execution - perform a similarity search (query in our previously created chunks) for us and send the found context to OpenAI along with the query. Make sure to set your OpenAI API Key.
3. To save on OpenAI costs, we are also introducing a simple SQLiteCache. In case we get the same query twice, we will simply use the cached response instead of sending the data to OpenAI again.
4. Finally, we execute the the RetrievalQA chain and return the result.

    retriever = st.session_state["db"].as_retriever()

    qa = RetrievalQA.from_chain_type(llm=OpenAI(model_name="gpt-3.5-turbo", openai_api_key=globals.OPENAI_API_KEY_QUERY), chain_type="stuff", retriever=retriever)
    langchain.llm_cache = SQLiteCache(database_path=globals.CACHE_PATH)
    result = qa({'query': query})
    return result

Setting Up the Streamlit UI

First, let's handle the user input:

1. This function is the event handler function and is called whenever we ask a new question in the UI.
2. Container in this case is just the group that we put all of our UI elements in.
3. We use st.spinner to create a "loading" indicator while we are querying the index.
4. Finally, we are putting both the chat input and the result into our session state from which we'll render the chat UI.

def handle_chat_user_input(container):
    output = None
    with container:
        with st.spinner("AI is working on answering your question..."):
            output = query_index(st.session_state.chat_input)['result']

        st.session_state['past'].append(st.session_state.chat_input)
        st.session_state['generated'].append(output)

Next, we're simply setting up the Streamlit page header:

def setup_page():
    st.set_page_config(page_title="Notion Chatbot", page_icon=":robt:")
    st.header("Quickly get information from your Notion")

And now we'll render the chat based on the session state variables:

1. We are making use of the "message" function from the streamlit_chat package.
2. We are simply looping through the session state - which at this point will contain both the chat input and the response from the AI.
3. Finally, we are creating the input form that users can use to input the question. Clicking on the submit button is what triggers the event handling function we defined above.

def render_app():
    container_chat = st.container()
    render_chat(container_chat)
    
def render_chat(container):
    with container:
        if st.session_state['generated']:
            for i in range(len(st.session_state['generated'])):
                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="icons")
                message(st.session_state["generated"][i], key=str(i), avatar_style="identicon")

        with st.form(key='chat', clear_on_submit=True):
            st.text_input("chat with the Notion content:", label_visibility="visible", placeholder="ask me anything", key='chat_input')
            st.form_submit_button(label='Send', on_click=handle_chat_user_input, use_container_width=True, kwargs=dict(container=container))

Calling All The Functions

Lastly, just simply call all the relevant functions we've created:

setup_page()
init_session_state()
load_notion_database()
render_app()

And that's it! Congratulations 🎉

Now simply run your app with:

streamlit run app.py