向量搜索示例

本示例演示如何使用 TiDB 和本地嵌入模型构建一个语义搜索应用。它使用向量搜索根据语义查找相似项（而不仅仅是关键字）。

该应用使用 Ollama 在本地生成向量嵌入，使用 Streamlit 构建 Web UI，并使用 pytidb（TiDB 的官方 Python SDK）来构建 RAG 流水线。

使用向量嵌入的语义搜索

前提条件

开始之前，请确保你具备以下条件：

• Python (>=3.10)：安装 Python 3.10 或以上版本。
• 一个 TiDB Cloud Starter 实例：你可以在 TiDB Cloud 上创建一个免费的 TiDB Cloud Starter 实例。
• Ollama：从 Ollama 安装。

运行方法

第 1 步：使用 Ollama 启动嵌入服务

拉取嵌入模型：

ollama pull mxbai-embed-large

验证嵌入服务是否正在运行：

curl http://localhost:11434/api/embed -d '{
  "model": "mxbai-embed-large",
  "input": "Llamas are members of the camelid family"
}'

第 2 步：克隆仓库

git clone https://github.com/pingcap/pytidb.git
cd pytidb/examples/vector_search/

第 3 步：安装所需软件包并设置环境

python -m venv .venv
source .venv/bin/activate
pip install -r reqs.txt

第 4 步：设置环境变量

1. 在 TiDB Cloud 控制台 中，进入 My TiDB 页面，然后点击目标 TiDB Cloud Starter 实例的名称，进入其实例概览页面。

2. 点击右上角的 Connect。此时会显示连接对话框，其中列出了连接参数。

3. 根据连接参数按如下方式设置环境变量：

   cat > .env <<EOF
   TIDB_HOST={gateway-region}.prod.aws.tidbcloud.com
   TIDB_PORT=4000
   TIDB_USERNAME={prefix}.root
   TIDB_PASSWORD={password}
   TIDB_DATABASE=pytidb_vector_search
   EOF
   

第 5 步：运行 Streamlit 应用

streamlit run app.py

打开浏览器并访问 http://localhost:8501。

相关资源

• 源代码：在 GitHub 上查看