In the evolving landscape of conversational AI, Retrieval-Augmented Generation (RAG) has emerged as a powerful approach to enhance chatbot performance. By combining the strengths of retrieval and generation, RAG-based systems provide more accurate and contextually relevant responses. In this article, we will explore how to build a RAG-based chatbot using LlamaIndex and TiDB Vector Search, a MySQL-compatible database.<\/p>\n\n\n\n
Retrieval-Augmented Generation (RAG) leverages both retrieval mechanisms and generative models. The retrieval mechanism fetches relevant documents or data snippets in response to a query, while the generative model creates human-like responses based on the retrieved information. This combination ensures that the chatbot has access to precise data and can generate coherent and contextually appropriate responses.<\/p>\n\n\n\n
Before we dive into the code, ensure you have the following environment variables set up:<\/p>\n\n\n\n
TIDB_USERNAME<\/code>: Your TiDB username.<\/li>\n\n\n\nTIDB_PASSWORD<\/code>: Your TiDB password.<\/li>\n\n\n\nTIDB_HOST<\/code>: The host address of your TiDB instance.<\/li>\n\n\n\nOPENAI_API_KEY<\/code>: Your OpenAI API key.<\/li>\n<\/ul>\n\n\n\nPrerequisites<\/strong><\/p>\n\n\n\n\n- A running TiDB Cloud \uc2a4\ud0c0\ud130<\/a> cluster with vector search enabled.<\/li>\n\n\n\n
- Python 3.8 or later.<\/li>\n\n\n\n
- OpenAI API key<\/a>.<\/li>\n<\/ul>\n\n\n\n
<\/span>Running the Example<\/span><\/h2>\n\n\n\nClone the Repository<\/h3>\n\n\n\n
First, clone the repository containing the example code:<\/p>\n\n\n\n
git clone https:\/\/github.com\/pingcap\/tidb-vector-python.git\ncd tidb-vector-python\/examples\/llamaindex-tidb-vector<\/code><\/pre>\n\n\n\nCreate a Virtual Environment<\/h3>\n\n\n\n
Next, create and activate a virtual environment:<\/p>\n\n\n\n
python3 -m venv .venv\nsource .venv\/bin\/activate<\/code><\/pre>\n\n\n\nInstall Dependencies<\/h3>\n\n\n\n
Install the required Python packages:<\/p>\n\n\n\n
pip install -r requirements.txt<\/code><\/pre>\n\n\n\nSet Environment Variables<\/h3>\n\n\n\n
Set the necessary environment variables with your credentials:<\/p>\n\n\n\n
export OPENAI_API_KEY=\"sk-*******\"\nexport TIDB_HOST=\"gateway01.*******.shared.aws.tidbcloud.com\"\nexport TIDB_USERNAME=\"****.root\"\nexport TIDB_PASSWORD=\"****\"<\/code><\/pre>\n\n\n\n<\/span>Code Walkthrough<\/span><\/h2>\n\n\n\nBelow is the complete code to build our RAG-based chatbot.<\/p>\n\n\n\n
#!\/usr\/bin\/env python\nimport os\n\nimport click\nfrom sqlalchemy import URL\nfrom llama_index.core import VectorStoreIndex, StorageContext\nfrom llama_index.vector_stores.tidbvector import TiDBVectorStore # type: ignore\nfrom llama_index.readers.web import SimpleWebPageReader\n\n# Define TiDB connection URL\ntidb_connection_url = URL(\n \"mysql+pymysql\",\n username=os.environ['TIDB_USERNAME'],\n password=os.environ['TIDB_PASSWORD'],\n host=os.environ['TIDB_HOST'],\n port=4000,\n database=\"test\",\n query={\"ssl_verify_cert\": True, \"ssl_verify_identity\": True},\n)\n\n# Initialize TiDB Vector Store\ntidbvec = TiDBVectorStore(\n connection_string=tidb_connection_url,\n table_name=\"llama_index_rag_test\",\n distance_strategy=\"cosine\",\n vector_dimension=1536, # The dimension is decided by the model\n drop_existing_table=False,\n)\n\n# Create VectorStoreIndex and StorageContext\ntidb_vec_index = VectorStoreIndex.from_vector_store(tidbvec)\nstorage_context = StorageContext.from_defaults(vector_store=tidbvec)\nquery_engine = tidb_vec_index.as_query_engine(streaming=True)\n\n# Function to prepare data\ndef do_prepare_data(url):\n documents = SimpleWebPageReader(html_to_text=True).load_data([url,])\n tidb_vec_index.from_documents(documents, storage_context=storage_context, show_progress=True)\n\n# Default URL for data loading\n_default_url = 'https:\/\/docs.pingcap.com\/tidb\/stable\/overview'\n\n@click.command()\n@click.option('--url', default=_default_url,\n help=f'URL you want to talk to, default={_default_url}')\ndef chat_with_url(url):\n do_prepare_data(url)\n while True:\n question = click.prompt(\"Enter your question\")\n response = query_engine.query(question)\n click.echo(response)\n\nif __name__ == '__main__':\n chat_with_url()<\/code><\/pre>\n\n\n\nExplanation of the Code<\/h3>\n\n\n\n\n- Importing Libraries<\/strong>: We start by importing the necessary libraries.
os<\/code> is used for environment variables, click<\/code> for command-line interaction, and various modules from llama_index<\/code> \uadf8\ub9ac\uace0 sqlalchemy<\/code> for handling the vector store and database connection.<\/li>\n\n\n\n- Defining TiDB Connection URL<\/strong>: We create a connection URL using the
URL<\/code> class from sqlalchemy<\/code>. This URL includes the TiDB credentials and connection details.<\/li>\n\n\n\n- Initializing TiDB Vector Store<\/strong>: We instantiate
TiDBVectorStore<\/code> with the connection string, table name, distance strategy (cosine similarity), vector dimension, and an option to drop the existing table.<\/li>\n\n\n\n- Creating VectorStoreIndex and StorageContext<\/strong>: We create an index from the vector store and a storage context for managing data storage.<\/li>\n\n\n\n
- Data Preparation Function<\/strong>: The
do_prepare_data<\/code> function loads data from a given URL, converts it to text, and stores it in the vector index.<\/li>\n\n\n\n- Command-Line Interaction<\/strong>: Using
click<\/code>, we define a command-line interface to allow users to specify a URL for data loading and interact with the chatbot by entering questions.<\/li>\n\n\n\n- Main Function<\/strong>: The
chat_with_url<\/code> function prepares data from the specified URL and enters a loop where it prompts the user for questions and returns responses from the query engine.<\/li>\n<\/ol>\n\n\n\nRunning the Chatbot<\/h3>\n\n\n\n
To run the chatbot, save the code to a file, for example, chat_with_url.py<\/code>, and execute it in your terminal:<\/p>\n\n\n\npython chat_with_url.py<\/code><\/pre>\n\n\n\n