RAG LLM

Learn how to augment LLMs with vector search with TileDB-Vector-Search and LangChain.

How to run this tutorial

We recommend running this tutorial, as well as the other various tutorials in the Tutorials section, inside TileDB Cloud. This will allow you to quickly experiment avoiding all the installation, deployment, and configuration hassles. Sign up for the free tier, spin up a TileDB Cloud notebook with a Python kernel, and follow the tutorial instructions. If you wish to learn how to run tutorials locally on your machine, read the Tutorials: Running Locally tutorial.

One of the limitations of LLMs is that their knowledge extends only to the data that were used during their training. Public training datasets are missing private and proprietary information required for enterprise applications. They are also missing information about the world and events that happened after the dataset creation time. This problem affects all types of LLMs, including public models, proprietary models, and even those deployed and used locally (e.g., in sensitive enterprise applications).

In this tutorial, you will use TileDB-Vector-Search to allow the gpt-3.5-turbo model to answer questions about LangChain. Most ChatGPT models have limited world knowledge after 2021, which is the training data cutoff date. LangChain was created and became popular after 2021. Although we use ChatGPT 3.5, this example can be easily extended to other LLMs. You will augment gpt-3.5-turbo with TileDB-Vector-Search, via LangChain, one of the most popular large language model (LLM) application development frameworks that integrates with our TileDB-Vector-Search library. This approach is called retrieval-augmented generation (RAG).

If you wish to learn more about RAG LLMs, visit the Introduction section.

Setup

To be able to run this tutorial, you will need an OpenAI API key. In addition, if you wish to use your local machine instead of a TileDB Cloud notebook, you will need to install the following:

Conda

conda install -c conda-forge langchain==0.0.331 openai==0.28.1 tiktoken

Pip

pip install langchain==0.0.331 openai==0.28.1 tiktoken

Start by importing the necessary libraries, setting the URIs you will use throughout the tutorial, and clean up any previously generated data.

import os
import shutil

from langchain.chains import ConversationalRetrievalChain, ConversationChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders.generic import GenericLoader
from langchain.document_loaders.parsers.txt import TextParser
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores.tiledb import TileDB

# URIs to be used throughout the tutorial
langchain_repo_uri = "langchain"
index_uri = "langchain_doc_index"

# Clean up
if os.path.exists(langchain_repo_uri):
    shutil.rmtree(langchain_repo_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)

Vanilla ChatGPT 3.5 Turbo

Initialize ChatGPT 3.5 Turbo.

# NOTE: Make sure you set the OPENAI_API_KEY environment
# variable with your OpenAI API key.

# Initialize chatgpt
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
chatgpt = ConversationChain(llm=llm)

Ask ChatGPT a question about LangChain. Note that ChatGPT incorrectly describes LangChain as a language learning platform based on the blockchain, and not as a framework for developing applications powered by LLMs.

question = "What is langchain?"
print(f"User: {question}")
print(f"AI: {chatgpt.run(question)}")

User: What is langchain?

AI: Langchain is a decentralized platform that uses blockchain technology to create a secure and transparent system for language learning. It allows users to connect with language tutors, access learning materials, and track their progress in real-time. The platform also uses smart contracts to ensure fair transactions between users and tutors.

RAG ChatGPT 3.5 Turbo

Now, you will use LangChain’s documentation to augment ChatGPT so that it can correctly answer the question about the project.

First, clone the LangChain repo from GitHub:

!git clone https://github.com/langchain-ai/langchain.git

Cloning into 'langchain'...
remote: Enumerating objects: 202474, done.
remote: Counting objects: 100% (15425/15425), done.
remote: Compressing objects: 100% (1401/1401), done.
remote: Total 202474 (delta 14572), reused 14225 (delta 14024), pack-reused 187049 (from 1)
Receiving objects: 100% (202474/202474), 283.53 MiB | 48.00 MiB/s, done.
Resolving deltas: 100% (151831/151831), done.
Updating files: 100% (7795/7795), done.

Next, parse the documents in the repo.

# Parse markdown documents and split them into text chunks
documentation_path = "./langchain/docs"
loader = GenericLoader.from_filesystem(
    documentation_path, glob="**/*", suffixes=[".mdx"], parser=TextParser()
)
splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.MARKDOWN, chunk_size=1000, chunk_overlap=100
)
documents = loader.load()
print(f"Number of raw documents loaded: {len(documents)}")
documents = splitter.split_documents(documents)
documents = [d for d in documents if len(d.page_content) > 5]
texts = [d.page_content for d in documents]
metadatas = [d.metadata for d in documents]
print(f"Number of document chunks: {len(texts)}")

Number of raw documents loaded: 299
Number of document chunks: 1268

Generate the appropriate vector embeddings from the documents, which will be used to create a vector index afterwards.

# Generate embeddings for each document chunk
embedding = OpenAIEmbeddings()
text_embeddings = embedding.embed_documents(texts)
text_embedding_pairs = list(zip(texts, text_embeddings))

Create a vector index on the generated embeddings, using TileDB-Vector-Search:

# Index document chunks using a TileDB IVF_FLAT index
db = TileDB.from_embeddings(
    text_embedding_pairs,
    embedding,
    index_uri=index_uri,
    index_type="IVF_FLAT",
    metadatas=metadatas,
    allow_dangerous_deserialization=True,
)
print(
    f"Number of vector embeddings stored in TileDB-Vector-Search: {len(text_embeddings)}"
)

Number of vector embeddings stored in TileDB-Vector-Search: 1268

Now, ask the same question, augmenting ChatGPT with the vector index you created.

db = TileDB.load(
    index_uri=index_uri, embedding=embedding, allow_dangerous_deserialization=True
)
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
retriever = db.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5},
)

# Chatgpt augmented with our vector index
rag_chatgpt = ConversationalRetrievalChain.from_llm(llm, retriever=retriever)

question = "What is langchain?"
print(f"User: {question}")
print(f"AI: {rag_chatgpt.run({'question': question, 'chat_history': ''})}\n")

User: What is langchain?
AI: LangChain is a platform that implements the latest research in Natural Language Processing (NLP). It allows users to build AI applications, chatbots, and agents using advanced language models like OpenAI, Google Gemini Pro, and LLAMA2. LangChain is used in various industries like retail, ed-tech, and more, and it offers a range of online courses and tutorials to help users get started with building AI applications.

You can see that ChatGPT successfully responds with meaningful information about the LangChain project.

Clean up

Clean up the generated data.

# Clean up
if os.path.exists(langchain_repo_uri):
    shutil.rmtree(langchain_repo_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)

What’s next?

Now that you know how to augment an LLM with vector search, you should learn about how you can augment it with conversation history as well by reading the Tutorials: LLM memory tutorial.