Integrate Vector Search with LlamaIndex
This tutorial demonstrates how to integrate the vector search feature in TiDB Cloud with LlamaIndex.
Prerequisites
To complete this tutorial, you need:
- Python 3.8 or higher installed.
- Jupyter Notebook installed.
- Git installed.
- A TiDB Serverless cluster. Follow creating a TiDB Serverless cluster to create your own TiDB Cloud cluster if you don't have one.
Get started
This section provides step-by-step instructions for integrating TiDB Vector Search with LlamaIndex to perform semantic searches.
Step 1. Create a new Jupyter Notebook file
In your preferred directory, create a new Jupyter Notebook file named integrate_with_llamaindex.ipynb
:
touch integrate_with_llamaindex.ipynb
Step 2. Install required dependencies
In your project directory, run the following command to install the required packages:
pip install llama-index-vector-stores-tidbvector
pip install llama-index
Open the integrate_with_llamaindex.ipynb
file in Jupyter Notebook and add the following code to import the required packages:
import textwrap
from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.tidbvector import TiDBVectorStore
Step 3. Set up your environment
Step 3.1 Obtain the connection string to the TiDB cluster
Navigate to the Clusters page, and then click the name of your target cluster to go to its overview page.
Click Connect in the upper-right corner. A connection dialog is displayed.
Ensure the configurations in the connection dialog match your operating environment.
- Connection Type is set to
Public
. - Branch is set to
main
. - Connect With is set to
SQLAlchemy
. - Operating System matches your environment.
- Connection Type is set to
Click the PyMySQL tab and copy the connection string.
Step 3.2 Configure environment variables
To establish a secure and efficient database connection, use the standard connection method provided by TiDB Cloud.
This document uses OpenAI as the embedding model provider. In this step, you need to provide the connection string obtained from step 3.1 and your OpenAI API key.
To configure the environment variables, run the following code. You will be prompted to enter your connection string and OpenAI API key:
import getpass
import os
tidb_connection_url = getpass.getpass(
"TiDB connection URL (format - mysql+pymysql://root@127.0.0.1:4000/test): "
)
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
Step 4. Load the sample document
Step 4.1 Download the sample document
In your project directory, create a directory named data/paul_graham/
and download the sample document paul_graham_essay.txt
from the run-llama/llama_index GitHub repository.
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
Step 4.2 Load the document
Load the sample document from data/paul_graham/paul_graham_essay.txt
using the SimpleDirectoryReader
class.
documents = SimpleDirectoryReader("./data/paul_graham").load_data()
print("Document ID:", documents[0].doc_id)
for index, document in enumerate(documents):
document.metadata = {"book": "paul_graham"}
Step 5. Embed and store document vectors
Step 5.1 Initialize the TiDB vector store
The following code creates a table named paul_graham_test
in TiDB, which is optimized for vector search.
tidbvec = TiDBVectorStore(
connection_string=tidb_connection_url,
table_name="paul_graham_test",
distance_strategy="cosine",
vector_dimension=1536,
drop_existing_table=False,
)
Upon successful execution, you can directly view and access the paul_graham_test
table in your TiDB database.
Step 5.2 Generate and store embeddings
The following code parses the documents, generates embeddings, and stores them in the TiDB vector store.
storage_context = StorageContext.from_defaults(vector_store=tidbvec)
index = VectorStoreIndex.from_documents(
documents, storage_context=storage_context, show_progress=True
)
The expected output is as follows:
Parsing nodes: 100%|██████████| 1/1 [00:00<00:00, 8.76it/s]
Generating embeddings: 100%|██████████| 21/21 [00:02<00:00, 8.22it/s]
Step 6. Perform a vector search
The following creates a query engine based on the TiDB vector store and performs a semantic similarity search.
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do?")
print(textwrap.fill(str(response), 100))
The expected output is as follows:
The author worked on writing, programming, building microcomputers, giving talks at conferences,
publishing essays online, developing spam filters, painting, hosting dinner parties, and purchasing
a building for office use.
Step 7. Search with metadata filters
To refine your searches, you can use metadata filters to retrieve specific nearest-neighbor results that match the applied filters.
Query with book != "paul_graham"
filter
The following example excludes results where the book
metadata field is "paul_graham"
:
from llama_index.core.vector_stores.types import (
MetadataFilter,
MetadataFilters,
)
query_engine = index.as_query_engine(
filters=MetadataFilters(
filters=[
MetadataFilter(key="book", value="paul_graham", operator="!="),
]
),
similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))
The expected output is as follows:
Empty Response
Query with book == "paul_graham"
filter
The following example filters results to include only documents where the book
metadata field is "paul_graham"
:
from llama_index.core.vector_stores.types import (
MetadataFilter,
MetadataFilters,
)
query_engine = index.as_query_engine(
filters=MetadataFilters(
filters=[
MetadataFilter(key="book", value="paul_graham", operator="=="),
]
),
similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))
The expected output is as follows:
The author learned programming on an IBM 1401 using an early version of Fortran in 9th grade, then
later transitioned to working with microcomputers like the TRS-80 and Apple II. Additionally, the
author studied philosophy in college but found it unfulfilling, leading to a switch to studying AI.
Later on, the author attended art school in both the US and Italy, where they observed a lack of
substantial teaching in the painting department.
Step 8. Delete documents
Delete the first document from the index:
tidbvec.delete(documents[0].doc_id)
Check whether the documents had been deleted:
query_engine = index.as_query_engine()
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))
The expected output is as follows:
Empty Response