TiDB Python SDK

[!NOTE]
该 Python 包正在快速开发中，API 可能会发生变化，建议在导入时使用固定版本，例如 pytidb==0.0.6

TiDB AI 的 Python SDK：一个统一的数据平台，旨在帮助开发者构建下一代 AI 应用。

• 🔍 支持多种搜索模式：向量搜索、全文搜索、混合搜索
• 🔄 自动生成嵌入
• 🎯 高级过滤功能
• 🥇 使用 Reranker 调整搜索结果
• 💱 事务支持

文档: https://pingcap.github.io/ai/

快速入门指南: Jupyter Notebook

安装 TiDB MCP 服务器 (文档: https://pingcap.github.io/ai/integrations/mcp):

安装

bash
pip install pytidb

如果你想使用内置的嵌入函数和重新排序器。

pip install "pytidb[models]"

如果你想将查询结果转换为 pandas DataFrame。

pip install pandas

连接到 TiDB Cloud

前往 tidbcloud.com 创建一个免费的 TiDB 集群。

python
import os
from pytidb import TiDBClient

db = TiDBClient.connect(
host=os.getenv("TIDB_HOST"),
port=int(os.getenv("TIDB_PORT")),
username=os.getenv("TIDB_USERNAME"),
password=os.getenv("TIDB_PASSWORD"),
database=os.getenv("TIDB_DATABASE"),
)

特性亮点

🤖 自动嵌入

PyTiDB 会自动将文本字段（例如 text）嵌入，并将向量嵌入保存到向量字段（例如 text_vec）。

创建带有嵌入函数的表：

python
from pytidb.schema import TableModel, Field
from pytidb.embeddings import EmbeddingFunction

text_embed = EmbeddingFunction("openai/text-embedding-3-small")

class Chunk(TableModel, table=True):
tablename = "chunks"

id: int = Field(primary_key=True)
text: str = Field()
text_vec: list[float] = text_embed.VectorField(
    source_field="text"
)  # 👈 定义向量字段。
user_id: int = Field()

table = db.create_table(schema=Chunk)

批量插入数据：

python
table.bulk_insert(
[
Chunk(id=2, text="bar", user_id=2), # 👈 文本字段将被自动嵌入并保存到
Chunk(id=3, text="baz", user_id=3), # text_vec 字段。
Chunk(id=4, text="qux", user_id=4), # 自动完成。
]
)

🔍 搜索

向量搜索

向量搜索基于语义相似性帮助你找到最相关的记录，因此你不需要在查询中明确包含所有关键词。

python
df = (
table.search("") # 👈 查询将被自动嵌入。
.filter({"user_id": 2})
.limit(2)
.to_pandas()
)

完整的示例，请参见 向量搜索 演示。

全文搜索

全文搜索通过分词查询并匹配确切的关键词来找到最相关的记录。

python
if not table.has_fts_index("text"):
table.create_fts_index("text") # 👈 在文本列上创建全文索引。

df = (
table.search("", search_type="fulltext")
.limit(2)
.to_pandas()
)对于完整的示例，请参阅全文搜索演示。

混合搜索

混合搜索结合了向量搜索和全文搜索，以提供更准确和相关的搜索结果。

python
from pytidb.rerankers import Reranker

jinaai = Reranker(model_name="jina_ai/jina-reranker-m0")

df = (
table.search("", search_type="hybrid")
.rerank(jinaai, "text") # 👈 使用 jinaai 模型对查询结果进行重新排序。
.limit(2)
.to_pandas()
)

对于完整的示例，请参阅混合搜索演示。

高级过滤

PyTiDB 支持各种操作符以实现灵活的过滤：

⛓ 结构化数据与非结构化数据的连接

python
from pytidb import Session
from pytidb.sql import select

创建一个表来存储用户数据：

class User(TableModel, table=True):
tablename = "users"

id: int = Field(primary_key=True)
name: str = Field(max_length=20)

with Session(engine) as session:
query = (
select(Chunk).join(User, Chunk.user_id == User.id).where(User.name == "Alice")
)
chunks = session.exec(query).all()

[(c.id, c.text, c.user_id) for c in chunks]

💱 事务支持

PyTiDB 支持事务管理，因此您可以避免竞态条件并确保数据一致性。

python
with db.session() as session:
initial_total_balance = db.query("SELECT SUM(balance) FROM players").scalar()

# 从玩家1向玩家2转账10个硬币
db.execute("UPDATE players SET balance = balance - 10 WHERE id = 1")
db.execute("UPDATE players SET balance = balance + 10 WHERE id = 2")

session.commit()
# 或者 session.rollback()

final_total_balance = db.query("SELECT SUM(balance) FROM players").scalar()
assert final_total_balance == initial_total_balance