[探索TiDB Vector：如何在AI应用中实现高效向量搜索]

最新推荐文章于 2024-10-01 19:47:46 发布

akhfuiigabv

最新推荐文章于 2024-10-01 19:47:46 发布

阅读量312

点赞数 5

文章标签： tidb 人工智能 python

本文链接：https://blog.csdn.net/akhfuiigabv/article/details/142413857

版权

引言

TiDB Cloud推出的TiDB Serverless为AI应用集成了内建的向量搜索功能。通过这一增强，开发者无需额外的数据库或技术栈，即可在MySQL环境中无缝开发AI应用。本篇文章将详细介绍如何利用TiDB Vector功能，展示其特性及实际应用。

主要内容

环境搭建

首先，我们需要安装所需的Python包。

%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)

准备数据

加载数据并进行文本切分，然后生成向量嵌入。

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的表，用于优化向量搜索。

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.",
        "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
)

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