07_TiDB AI应用开发：使用PyTiDB SDK快速上手

07_TiDB AI 应用开发使用 PyTiDB SDK 快速上手标签TiDBPyTiDBPython SDK向量搜索AI应用开发Auto Embedding对话记忆关键词PyTiDB SDK、TiDBClient、向量搜索Python、全文搜索Python、混合搜索、Auto Embedding、对话记忆、Document模型、AI应用开发、TiDB Python一、为什么需要 PyTiDB 这样的高级 SDK很长一段时间我在 TiDB 上做向量搜索都是直接用 pymysql 拼 SQL——毕竟 SQL 功底好写起来也不费力。但有一天同事接手我写的代码问了我一个问题“这个VEC_COSINE_DISTANCE参数格式是[1,2,3]还是[1,2,3]需要 escape 吗”我愣了一下。这是一个本不应该存在的问题。pymysql 处理向量列时你需要手动把 Python 列表序列化成 TiDB 接受的字符串格式查询结果中向量字段也需要手动反序列化还要自己管理嵌入 API 调用、批量写入、连接池……每一件事单独看都不难但加在一起代码变得又长又脆。PyTiDB 的设计目标就是消掉这些胶水代码让你专注在业务逻辑上。本文从零带你掌握 PyTiDB 的核心用法。二、PyTiDB 核心概念2.1 三层架构PyTiDB 架构图 --------------------------- | 你的应用代码 | --------------------------- | | | TiDBClient | 数据库连接管理 | | | | -- Table | 表操作抽象向量/全文/混合搜索 | | | | -- Document | 文档数据模型text embedding meta | | | EmbeddingFunction | 嵌入模型集成OpenAI/本地模型 | | --------------------------- | TiDB 数据库后端 | ---------------------------2.2 Document 数据模型Document 是 PyTiDB 的核心数据模型一个 Document 对应知识库中的一条记录frompytidb.schemaimportDocument docDocument(textTiDB 是一款开源的分布式 NewSQL 数据库,# 文本内容text_vec[0.12,-0.34,0.87,...],# 预先计算的向量可选meta{# 任意元数据JSONcategory:数据库,version:8.5,source:官方文档,page:1})如果没有提供text_vec且配置了 Auto EmbeddingPyTiDB 会在插入时自动调用嵌入 API 生成向量。三、快速开始连接与基础 CRUD3.1 安装pipinstallpytidb[openai]# 包含 OpenAI 嵌入支持# 或pipinstallpytidb[all]# 包含所有可选依赖OpenAI Jina MCP等3.2 连接数据库frompytidbimportTiDBClient# 方式一参数连接推荐dbTiDBClient.connect(hostgateway01.eu-central-1.prod.aws.tidbcloud.com,port4000,usernameyour_username.root,passwordyour_password,databaseai_demo,# TiDB Cloud Serverless 需要 SSLssl{ca:/path/to/ca.pem})# 方式二连接字符串dbTiDBClient.connect(database_url(mysqlpymysql://user:passhost:4000/db?ssl_verify_certtruessl_verify_identitytruessl_ca/path/to/ca.pem))print(连接成功TiDB 版本,db.tidb_version)3.3 创建表和插入数据frompytidb.schemaimportTableModel,Column,TEXT,VectorType# 方式一使用内置文档模型最简单tabledb.get_table(my_documents)# 如果表不存在会自动创建默认结构# id (bigint), text (text), text_vec (vector), meta (json)# 插入文档手动指定向量table.bulk_insert([Document(text苹果是一种水果,text_vec[1.0,2.0,1.0,0.5],meta{category:水果,season:秋}),Document(text香蕉富含钾元素,text_vec[1.0,2.0,1.5,0.3],meta{category:水果,season:全年}),Document(text牧羊犬是一种工作犬,text_vec[0.1,0.2,0.1,0.9],meta{category:动物,type:犬类}),])print(f当前表中共有{table.count()}条记录)3.4 查询单条记录# 按 ID 查询recordtable.get(id1)print(record.text,record.meta)# 列出所有记录调试用生产慎用all_recordstable.list()forrinall_records:print(f[{r.id}]{r.text})# 按元数据过滤fruitstable.list(filters{category:水果})四、搜索模式详解4.1 向量搜索# 输入向量找最相似的 K 条记录query_vec[1.0,2.0,1.2,0.4]# 某种水果相关的向量results(table.search(query_vec)# 输入查询向量.limit(3)# 返回前 3 条.to_list())forrinresults:print(f[距离:{r[_distance]:.4f}]{r[text]})# 输出示例# [距离: 0.0320] 苹果是一种水果# [距离: 0.0845] 香蕉富含钾元素# [距离: 0.8920] 牧羊犬是一种工作犬4.2 向量搜索 元数据过滤# 只在水果类别中搜索results(table.search(query_vec).filter({category:水果})# 先过滤再搜索.limit(5).to_list())4.3 全文搜索# 关键词全文搜索BM25 相关性排序results(table.search(钾元素 营养,search_typefulltext).limit(5).to_list())forrinresults:print(f[分数:{r.get(_score,0):.4f}]{r[text]})4.4 混合搜索推荐 RAG 使用# 混合搜索向量 全文自动 RRF 融合results(table.search(水果营养价值,search_typehybrid).limit(5).to_list())4.5 距离计算选择# 默认使用余弦距离可以指定其他距离类型frompytidb.distanceimportDistanceType results(table.search(query_vec,distance_typeDistanceType.L2).limit(5).to_list())五、Auto Embedding让向量生成自动化Auto Embedding 是 PyTiDB 最省心的功能。配置好嵌入函数后插入文档时无需手动生成向量PyTiDB 自动调用嵌入 API 并存储。5.1 配置 OpenAI Auto EmbeddingfrompytidbimportTiDBClientfrompytidb.embeddingsimportEmbeddingFunction# 定义嵌入函数调用 OpenAI APIembed_fnEmbeddingFunction(model_nametext-embedding-3-small,api_keysk-your-openai-key,dimensions1536)dbTiDBClient.connect(...)# 创建带 Auto Embedding 的表tabledb.get_table(auto_embed_docs,embedding_functionembed_fn# 绑定嵌入函数)# 插入时不需要提供 text_vecPyTiDB 自动生成table.bulk_insert([Document(textTiDB 支持水平扩展),Document(text向量搜索基于余弦相似度),Document(textHTAP 混合负载处理),],batch_size50)# batch_size 控制每批调用嵌入 API 的数量# 查询时也自动嵌入resultstable.search(数据库扩展方案).limit(3).to_list()# 内部自动将 数据库扩展方案 转为向量再做相似度搜索5.2 使用本地 BGE 模型零 API 成本frompytidb.embeddingsimportHuggingFaceEmbeddingFunction# 使用本地 BGE 模型首次运行自动下载约 400MBembed_fnHuggingFaceEmbeddingFunction(model_nameBAAI/bge-base-zh-v1.5,devicecpu,# 或 cuda 使用 GPUnormalizeTrue# BGE 模型需要归一化)tabledb.get_table(bge_docs,embedding_functionembed_fn)# 后续使用方式完全相同table.bulk_insert([Document(text纯中文知识库内容...)])resultstable.search(搜索中文语义).limit(5).to_list()六、对话记忆为 AI Agent 提供持久化上下文6.1 为什么需要基于向量的对话记忆普通的对话记忆是把历史消息全部放入 context window对话轮次一多Token 消耗爆炸。基于向量的记忆只检索最相关的历史消息是 Token 高效的解决方案。6.2 实现向量化对话记忆importjsonfromdatetimeimportdatetimefrompytidbimportTiDBClientfrompytidb.schemaimportDocumentfrompytidb.embeddingsimportEmbeddingFunctionclassVectorConversationMemory: 基于 TiDB 向量搜索的对话记忆 - 每条消息持久化存储不丢失 - 检索时只取最相关的历史消息 def__init__(self,db:TiDBClient,session_id:str):self.session_idsession_id embed_fnEmbeddingFunction(model_nametext-embedding-3-small,api_keysk-your-key)self.tabledb.get_table(conversation_memory,embedding_functionembed_fn)defadd_message(self,role:str,content:str):添加一条对话消息self.table.bulk_insert([Document(textcontent,meta{session_id:self.session_id,role:role,# user 或 assistanttimestamp:datetime.now().isoformat()})])defget_relevant_history(self,query:str,top_k:int5)-list: 检索与当前问题最相关的历史消息 results(self.table.search(query).filter({session_id:self.session_id})# 只查当前会话.limit(top_k).to_list())# 按时间排序让上下文有序returnsorted(results,keylambdax:x.get(meta,{}).get(timestamp,))defbuild_context_messages(self,current_query:str)-list: 构建 LLM 调用所需的消息列表带相关历史 historyself.get_relevant_history(current_query,top_k6)messages[]forhinhistory:metah.get(meta,{})messages.append({role:meta.get(role,user),content:h[text]})# 加上当前问题messages.append({role:user,content:current_query})returnmessages# 使用示例dbTiDBClient.connect(host...,...)memoryVectorConversationMemory(db,session_iduser_alice_20240410)# 模拟对话importopenai llmopenai.OpenAI()defchat_with_memory(question:str):# 构建带历史记忆的消息messagesmemory.build_context_messages(question)# 系统提示system_msg{role:system,content:你是一个专业的技术助手。}all_messages[system_msg]messages# 调用 LLMresponsellm.chat.completions.create(modelgpt-4o-mini,messagesall_messages,temperature0.1)answerresponse.choices[0].message.content# 记录对话历史memory.add_message(user,question)memory.add_message(assistant,answer)returnanswer# 多轮对话测试print(chat_with_memory(TiDB 向量搜索支持多少维度))print(chat_with_memory(那这个维度上限和主流嵌入模型比怎么样))# 需要记忆上一轮print(chat_with_memory(如何创建向量索引来优化查询性能))七、高级用法自定义 TableModel当内置的 Document 模型不够用时比如需要额外的业务字段可以自定义 TableModelfrompytidb.schemaimportTableModel,ColumnfromsqlalchemyimportBigInteger,String,Text,DateTime,Integerfromsqlalchemy.sqlimportfuncclassArticleTable(TableModel,tableTrue):自定义文章向量表__tablename__articlesid:intColumn(BigInteger,primary_keyTrue,autoincrementTrue)title:strColumn(String(500),comment文章标题)content:strColumn(Text,comment文章内容)author:strColumn(String(100),comment作者)category:strColumn(String(50),comment分类)view_count:intColumn(Integer,default0,comment浏览量)created_at:datetimeColumn(DateTime,server_defaultfunc.now())# 向量列继承 VectorMixin 自动处理# 通过 embedding_function 绑定时自动添加# 使用自定义模型tabledb.get_table(articles,schemaArticleTable,embedding_functionembed_fn)# 插入时使用 dicttable.bulk_insert([{title:TiDB 8.5 发版解读,content:TiDB 8.5 LTS 带来了...,author:PingCAP,category:数据库,view_count:0}])# 向量搜索自动嵌入 titlecontent 的组合文本resultstable.search(分布式数据库新特性).limit(5).to_list()八、错误处理与最佳实践8.1 连接重试与错误处理frompytidbimportTiDBClientfrompytidb.exceptionsimportTiDBConnectionErrorimporttimedefcreate_client_with_retry(max_retries:int3,**connect_kwargs):带重试的客户端创建forattemptinrange(max_retries):try:dbTiDBClient.connect(**connect_kwargs)returndbexceptTiDBConnectionErrorase:ifattemptmax_retries-1:raiseprint(f连接失败{2**attempt}秒后重试... ({e}))time.sleep(2**attempt)# 指数退避8.2 批量操作性能对比importtime# 测试逐条插入 vs 批量插入documents[Document(textf测试文档{i})foriinrange(1000)]# 方式1逐条慢不推荐starttime.time()fordocindocuments[:100]:table.insert(doc)print(f逐条插入100条{time.time()-start:.2f}s)# 方式2批量插入快推荐starttime.time()table.bulk_insert(documents[:1000],batch_size100)print(f批量插入1000条{time.time()-start:.2f}s)# 实测批量比逐条快 10-50 倍九、总结PyTiDB 是官方出品、专为 AI 应用设计的 Python SDK大幅简化了 TiDB 向量搜索的开发体验。核心能力总结TiDBClient负责连接管理支持 SSL/连接池生产可用三种搜索模式向量/全文/混合通过search_type参数切换接口统一Auto Embedding自动化嵌入生成支持 OpenAI、Hugging Face 等多种模型Document 模型 元数据过滤让向量搜索与业务过滤无缝结合对话记忆基于向量检索历史消息Token 效率远高于全量 context相比直接用 pymysql 拼 SQLPyTiDB 让代码量减少约 60%同时降低了出错概率。无论是快速原型还是生产级 AI 应用PyTiDB 都是 TiDB 技术栈下的首选 Python SDK。相关资源PyTiDB 官方文档PyTiDB GitHubPyTiDB PyPI