RAG 示例

本示例演示如何使用 pytidb（TiDB 官方 Python SDK）构建一个最小化的 RAG 应用。

该应用使用 Ollama 在本地生成向量嵌入，使用 Streamlit 构建 Web UI，并使用 pytidb 搭建 RAG 流水线。

RAG application built with PyTiDB

使用 PyTiDB 构建的 RAG 应用

前提条件

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

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

运行方法

第 1 步：准备推理 API

使用 Ollama CLI 拉取向量嵌入模型和 LLM 模型：

ollama pull mxbai-embed-large
ollama pull gemma3:4b
ollama run gemma3:4b

验证 /embed 和 /generate 端点是否正在运行：

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

curl http://localhost:11434/api/generate -d '{
  "model": "gemma3:4b",
  "prompt": "Hello, Who are you?"
}'

第 2 步：克隆仓库

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

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

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

第 4 步：设置环境变量

在 TiDB Cloud 控制台中，进入 My TiDB 页面，然后点击目标 TiDB Cloud Starter 实例的名称，进入其实例概览页。
点击右上角的 Connect。此时会显示连接对话框，其中列出了连接参数。
根据连接参数按如下方式设置环境变量：

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

第 5 步：运行 Streamlit 应用

streamlit run main.py

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

故障排查

`502 Bad Gateway` 错误

尝试禁用全局代理设置。

源代码：在 GitHub 上查看