You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
langchain/docs/docs/integrations/vectorstores/tidb_vector.ipynb

684 lines
26 KiB
Plaintext

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# TiDB Vector\n",
"\n",
"> [TiDB Cloud](https://tidbcloud.com/), is a comprehensive Database-as-a-Service (DBaaS) solution, that provides dedicated and serverless options. TiDB Serverless is now integrating a built-in vector search into the MySQL landscape. With this enhancement, you can seamlessly develop AI applications using TiDB Serverless without the need for a new database or additional technical stacks. Be among the first to experience it by joining the waitlist for the private beta at https://tidb.cloud/ai.\n",
"\n",
"This notebook provides a detailed guide on utilizing the TiDB Vector functionality, showcasing its features and practical applications."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setting up environments"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Begin by installing the necessary packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install langchain\n",
"%pip install langchain-openai\n",
"%pip install pymysql\n",
"%pip install tidb-vector"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Configure both the OpenAI and TiDB host settings that you will need. In this notebook, we will follow the standard connection method provided by TiDB Cloud to establish a secure and efficient database connection."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Here we useimport getpass\n",
"import getpass\n",
"import os\n",
"\n",
"os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n",
"# copy from tidb cloud console\n",
"tidb_connection_string_template = \"mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true\"\n",
"# tidb_connection_string_template = \"mysql+pymysql://root:<PASSWORD>@34.212.137.91:4000/test\"\n",
"tidb_password = getpass.getpass(\"Input your TiDB password:\")\n",
"tidb_connection_string = tidb_connection_string_template.replace(\n",
" \"<PASSWORD>\", tidb_password\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Prepare the following data"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"from langchain.text_splitter import CharacterTextSplitter\n",
"from langchain_community.document_loaders import TextLoader\n",
"from langchain_community.vectorstores import TiDBVectorStore\n",
"from langchain_openai import OpenAIEmbeddings"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"loader = TextLoader(\"../../modules/state_of_the_union.txt\")\n",
"documents = loader.load()\n",
"text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
"docs = text_splitter.split_documents(documents)\n",
"\n",
"embeddings = OpenAIEmbeddings()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Semantic similarity search\n",
"\n",
"TiDB supports both cosine and Euclidean distances ('cosine', 'l2'), with 'cosine' being the default choice.\n",
"\n",
"The code snippet below creates a table named `TABLE_NAME` in TiDB, optimized for vector searching. Upon successful execution of this code, you will be able to view and access the `TABLE_NAME` table directly within your TiDB database."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"TABLE_NAME = \"semantic_embeddings\"\n",
"db = TiDBVectorStore.from_documents(\n",
" documents=docs,\n",
" embedding=embeddings,\n",
" table_name=TABLE_NAME,\n",
" connection_string=tidb_connection_string,\n",
" distance_strategy=\"cosine\", # default, another option is \"l2\"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"query = \"What did the president say about Ketanji Brown Jackson\"\n",
"docs_with_score = db.similarity_search_with_score(query, k=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Please note that a lower cosine distance indicates higher similarity."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--------------------------------------------------------------------------------\n",
"Score: 0.18459301498220004\n",
"Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while youre at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
"\n",
"Tonight, Id like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
"\n",
"One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
"\n",
"And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nations top legal minds, who will continue Justice Breyers legacy of excellence.\n",
"--------------------------------------------------------------------------------\n",
"--------------------------------------------------------------------------------\n",
"Score: 0.2172729943284636\n",
"A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since shes been nominated, shes received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n",
"\n",
"And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n",
"\n",
"We can do both. At our border, weve installed new technology like cutting-edge scanners to better detect drug smuggling. \n",
"\n",
"Weve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \n",
"\n",
"Were putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n",
"\n",
"Were securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.\n",
"--------------------------------------------------------------------------------\n",
"--------------------------------------------------------------------------------\n",
"Score: 0.2262166799003692\n",
"And for our LGBTQ+ Americans, lets finally get the bipartisan Equality Act to my desk. The onslaught of state laws targeting transgender Americans and their families is wrong. \n",
"\n",
"As I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. \n",
"\n",
"While it often appears that we never agree, that isnt true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice. \n",
"\n",
"And soon, well strengthen the Violence Against Women Act that I first wrote three decades ago. It is important for us to show the nation that we can come together and do big things. \n",
"\n",
"So tonight Im offering a Unity Agenda for the Nation. Four big things we can do together. \n",
"\n",
"First, beat the opioid epidemic.\n",
"--------------------------------------------------------------------------------\n"
]
}
],
"source": [
"for doc, score in docs_with_score:\n",
" print(\"-\" * 80)\n",
" print(\"Score: \", score)\n",
" print(doc.page_content)\n",
" print(\"-\" * 80)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, the similarity_search_with_relevance_scores method can be used to obtain relevance scores, where a higher score indicates greater similarity."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--------------------------------------------------------------------------------\n",
"Score: 0.8154069850178\n",
"Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while youre at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
"\n",
"Tonight, Id like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
"\n",
"One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
"\n",
"And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nations top legal minds, who will continue Justice Breyers legacy of excellence.\n",
"--------------------------------------------------------------------------------\n",
"--------------------------------------------------------------------------------\n",
"Score: 0.7827270056715364\n",
"A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since shes been nominated, shes received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n",
"\n",
"And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n",
"\n",
"We can do both. At our border, weve installed new technology like cutting-edge scanners to better detect drug smuggling. \n",
"\n",
"Weve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \n",
"\n",
"Were putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n",
"\n",
"Were securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.\n",
"--------------------------------------------------------------------------------\n"
]
}
],
"source": [
"docs_with_relevance_score = db.similarity_search_with_relevance_scores(query, k=2)\n",
"for doc, score in docs_with_relevance_score:\n",
" print(\"-\" * 80)\n",
" print(\"Score: \", score)\n",
" print(doc.page_content)\n",
" print(\"-\" * 80)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Filter with metadata\n",
"\n",
"perform searches using metadata filters to retrieve a specific number of nearest-neighbor results that align with the applied filters.\n",
"\n",
"## Supported metadata types\n",
"\n",
"Each vector in the TiDB Vector Store can be paired with metadata, structured as key-value pairs within a JSON object. The keys are strings, and the values can be of the following types:\n",
"\n",
"- String\n",
"- Number (integer or floating point)\n",
"- Booleans (true, false)\n",
"\n",
"For instance, consider the following valid metadata payloads:\n",
"\n",
"```json\n",
"{\n",
" \"page\": 12,\n",
" \"book_tile\": \"Siddhartha\"\n",
"}\n",
"```\n",
"\n",
"## Metadata filter syntax\n",
"\n",
"The available filters include:\n",
"\n",
"- $or - Selects vectors that meet any one of the given conditions.\n",
"- $and - Selects vectors that meet all of the given conditions.\n",
"- $eq - Equal to\n",
"- $ne - Not equal to\n",
"- $gt - Greater than\n",
"- $gte - Greater than or equal to\n",
"- $lt - Less than\n",
"- $lte - Less than or equal to\n",
"- $in - In array\n",
"- $nin - Not in array\n",
"\n",
"Assuming one vector with metada:\n",
"```json\n",
"{\n",
" \"page\": 12,\n",
" \"book_tile\": \"Siddhartha\"\n",
"}\n",
"```\n",
"\n",
"The following metadata filters will match the vector\n",
"\n",
"```json\n",
"{\"page\": 12}\n",
"\n",
"{\"page\":{\"$eq\": 12}}\n",
"\n",
"{\"page\":{\"$in\": [11, 12, 13]}}\n",
"\n",
"{\"page\":{\"$nin\": [13]}}\n",
"\n",
"{\"page\":{\"$lt\": 11}}\n",
"\n",
"{\n",
" \"$or\": [{\"page\": 11}, {\"page\": 12}],\n",
" \"$and\": [{\"page\": 12}, {\"page\": 13}],\n",
"}\n",
"```\n",
"\n",
"Please note that each key-value pair in the metadata filters is treated as a separate filter clause, and these clauses are combined using the AND logical operator."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[UUID('c782cb02-8eec-45be-a31f-fdb78914f0a7'),\n",
" UUID('08dcd2ba-9f16-4f29-a9b7-18141f8edae3')]"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"db.add_texts(\n",
" texts=[\n",
" \"TiDB Vector offers advanced, high-speed vector processing capabilities, enhancing AI workflows with efficient data handling and analytics support.\",\n",
" \"TiDB Vector, starting as low as $10 per month for basic usage\",\n",
" ],\n",
" metadatas=[\n",
" {\"title\": \"TiDB Vector functionality\"},\n",
" {\"title\": \"TiDB Vector Pricing\"},\n",
" ],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--------------------------------------------------------------------------------\n",
"Score: 0.12761409169211535\n",
"TiDB Vector offers advanced, high-speed vector processing capabilities, enhancing AI workflows with efficient data handling and analytics support.\n",
"--------------------------------------------------------------------------------\n"
]
}
],
"source": [
"docs_with_score = db.similarity_search_with_score(\n",
" \"Introduction to TiDB Vector\", filter={\"title\": \"TiDB Vector functionality\"}, k=4\n",
")\n",
"for doc, score in docs_with_score:\n",
" print(\"-\" * 80)\n",
" print(\"Score: \", score)\n",
" print(doc.page_content)\n",
" print(\"-\" * 80)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Using as a Retriever"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In Langchain, a retriever is an interface that retrieves documents in response to an unstructured query, offering a broader functionality than a vector store. The code below demonstrates how to utilize TiDB Vector as a retriever."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--------------------------------------------------------------------------------\n",
"Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while youre at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
"\n",
"Tonight, Id like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
"\n",
"One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
"\n",
"And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nations top legal minds, who will continue Justice Breyers legacy of excellence.\n",
"--------------------------------------------------------------------------------\n"
]
}
],
"source": [
"retriever = db.as_retriever(\n",
" search_type=\"similarity_score_threshold\",\n",
" search_kwargs={\"k\": 3, \"score_threshold\": 0.8},\n",
")\n",
"docs_retrieved = retriever.invoke(query)\n",
"for doc in docs_retrieved:\n",
" print(\"-\" * 80)\n",
" print(doc.page_content)\n",
" print(\"-\" * 80)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Advanced Use Case Scenario\n",
"\n",
"Let's look a advanced use case - a travel agent is crafting a custom travel report for clients who desire airports with specific amenities such as clean lounges and vegetarian options. The process involves:\n",
"- A semantic search within airport reviews to extract airport codes meeting these amenities.\n",
"- A subsequent SQL query that joins these codes with route information, detailing airlines and destinations aligned with the clients' preferences.\n",
"\n",
"First, let's prepare some airpod related data"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[UUID('6dab390f-acd9-4c7d-b252-616606fbc89b'),\n",
" UUID('9e811801-0e6b-4893-8886-60f4fb67ce69'),\n",
" UUID('f426747c-0f7b-4c62-97ed-3eeb7c8dd76e')]"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# create table to store airplan data\n",
"db.tidb_vector_client.execute(\n",
" \"\"\"CREATE TABLE airplan_routes (\n",
" id INT AUTO_INCREMENT PRIMARY KEY,\n",
" airport_code VARCHAR(10),\n",
" airline_code VARCHAR(10),\n",
" destination_code VARCHAR(10),\n",
" route_details TEXT,\n",
" duration TIME,\n",
" frequency INT,\n",
" airplane_type VARCHAR(50),\n",
" price DECIMAL(10, 2),\n",
" layover TEXT\n",
" );\"\"\"\n",
")\n",
"\n",
"# insert some data into Routes and our vector table\n",
"db.tidb_vector_client.execute(\n",
" \"\"\"INSERT INTO airplan_routes (\n",
" airport_code,\n",
" airline_code,\n",
" destination_code,\n",
" route_details,\n",
" duration,\n",
" frequency,\n",
" airplane_type,\n",
" price,\n",
" layover\n",
" ) VALUES \n",
" ('JFK', 'DL', 'LAX', 'Non-stop from JFK to LAX.', '06:00:00', 5, 'Boeing 777', 299.99, 'None'),\n",
" ('LAX', 'AA', 'ORD', 'Direct LAX to ORD route.', '04:00:00', 3, 'Airbus A320', 149.99, 'None'),\n",
" ('EFGH', 'UA', 'SEA', 'Daily flights from SFO to SEA.', '02:30:00', 7, 'Boeing 737', 129.99, 'None');\n",
" \"\"\"\n",
")\n",
"db.add_texts(\n",
" texts=[\n",
" \"Clean lounges and excellent vegetarian dining options. Highly recommended.\",\n",
" \"Comfortable seating in lounge areas and diverse food selections, including vegetarian.\",\n",
" \"Small airport with basic facilities.\",\n",
" ],\n",
" metadatas=[\n",
" {\"airport_code\": \"JFK\"},\n",
" {\"airport_code\": \"LAX\"},\n",
" {\"airport_code\": \"EFGH\"},\n",
" ],\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finding Airports with Clean Facilities and Vegetarian Options via Vector Search"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--------------------------------------------------------------------------------\n",
"Clean lounges and excellent vegetarian dining options. Highly recommended.\n",
"{'airport_code': 'JFK'}\n",
"--------------------------------------------------------------------------------\n",
"--------------------------------------------------------------------------------\n",
"Comfortable seating in lounge areas and diverse food selections, including vegetarian.\n",
"{'airport_code': 'LAX'}\n",
"--------------------------------------------------------------------------------\n"
]
}
],
"source": [
"retriever = db.as_retriever(\n",
" search_type=\"similarity_score_threshold\",\n",
" search_kwargs={\"k\": 3, \"score_threshold\": 0.85},\n",
")\n",
"semantic_query = \"Could you recommend a US airport with clean lounges and good vegetarian dining options?\"\n",
"reviews = retriever.invoke(semantic_query)\n",
"for r in reviews:\n",
" print(\"-\" * 80)\n",
" print(r.page_content)\n",
" print(r.metadata)\n",
" print(\"-\" * 80)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[(1, 'JFK', 'DL', 'LAX', 'Non-stop from JFK to LAX.', datetime.timedelta(seconds=21600), 5, 'Boeing 777', Decimal('299.99'), 'None'),\n",
" (2, 'LAX', 'AA', 'ORD', 'Direct LAX to ORD route.', datetime.timedelta(seconds=14400), 3, 'Airbus A320', Decimal('149.99'), 'None')]"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Extracting airport codes from the metadata\n",
"airport_codes = [review.metadata[\"airport_code\"] for review in reviews]\n",
"\n",
"# Executing a query to get the airport details\n",
"search_query = \"SELECT * FROM airplan_routes WHERE airport_code IN :codes\"\n",
"params = {\"codes\": tuple(airport_codes)}\n",
"\n",
"airport_details = db.tidb_vector_client.execute(search_query, params)\n",
"airport_details.get(\"result\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Alternatively, we can streamline the process by utilizing a single SQL query to accomplish the search in one step."
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[(0.1219207353407008, 1, 'JFK', 'DL', 'LAX', 'Non-stop from JFK to LAX.', datetime.timedelta(seconds=21600), 5, 'Boeing 777', Decimal('299.99'), 'None', 'Clean lounges and excellent vegetarian dining options. Highly recommended.'),\n",
" (0.14613754359804654, 2, 'LAX', 'AA', 'ORD', 'Direct LAX to ORD route.', datetime.timedelta(seconds=14400), 3, 'Airbus A320', Decimal('149.99'), 'None', 'Comfortable seating in lounge areas and diverse food selections, including vegetarian.'),\n",
" (0.19840519342700513, 3, 'EFGH', 'UA', 'SEA', 'Daily flights from SFO to SEA.', datetime.timedelta(seconds=9000), 7, 'Boeing 737', Decimal('129.99'), 'None', 'Small airport with basic facilities.')]"
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"search_query = f\"\"\"\n",
" SELECT \n",
" VEC_Cosine_Distance(se.embedding, :query_vector) as distance, \n",
" ar.*,\n",
" se.document as airport_review\n",
" FROM \n",
" airplan_routes ar\n",
" JOIN \n",
" {TABLE_NAME} se ON ar.airport_code = JSON_UNQUOTE(JSON_EXTRACT(se.meta, '$.airport_code'))\n",
" ORDER BY distance ASC \n",
" LIMIT 5;\n",
"\"\"\"\n",
"query_vector = embeddings.embed_query(semantic_query)\n",
"params = {\"query_vector\": str(query_vector)}\n",
"airport_details = db.tidb_vector_client.execute(search_query, params)\n",
"airport_details.get(\"result\")"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'success': True, 'result': 0, 'error': None}"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# clean up\n",
"db.tidb_vector_client.execute(\"DROP TABLE airplan_routes\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Delete\n",
"\n",
"You can remove the TiDB Vector Store by using the `.drop_vectorstore()` method."
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
"db.drop_vectorstore()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "langchain",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.13"
}
},
"nbformat": 4,
"nbformat_minor": 2
}