向量搜索与 AI 集成

一、向量搜索概述

1.1 什么是向量搜索

向量搜索是 AI 时代的核心技术之一。它的基本思路是：将文本、图片、音频等非结构化数据转换为高维向量（数值数组），然后通过计算向量之间的距离来衡量相似度。

文本 → 向量化（Embedding） → 向量存储 → 相似度搜索

"如何学习 TiDB"
     │
     v  (通过 Embedding 模型)
[0.12, -0.34, 0.56, 0.78, -0.23, ...]  (1536 维向量)

搜索 "分布式数据库教程"
     │
     v  (同样转为向量)
[0.15, -0.31, 0.52, 0.75, -0.20, ...]

计算两个向量的距离 → 距离越近，语义越相似

1.2 为什么数据库需要向量搜索

在 AI 应用中，向量搜索有独特需求：

TiDB 的向量搜索让用户可以在同一套系统中同时处理结构化业务数据和向量语义搜索。

二、TiDB 向量数据类型

2.1 VECTOR 类型

TiDB 原生支持 VECTOR 数据类型：

-- 创建包含向量列的表
CREATE TABLE documents (
    id BIGINT PRIMARY KEY AUTO_RANDOM,
    title VARCHAR(255) NOT NULL,
    content TEXT,
    embedding VECTOR NOT NULL,  -- 向量列
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

2.2 插入向量数据

-- 直接插入向量
INSERT INTO documents (title, content, embedding) VALUES (
    'TiDB 架构简介',
    'TiDB 是一款分布式 NewSQL 数据库...',
    '[0.12, -0.34, 0.56, 0.78, -0.23, 0.45, 0.67, -0.11]'
);

-- 批量插入
INSERT INTO documents (title, content, embedding) VALUES
    ('分布式事务', 'Percolator 模型...', '[0.23, -0.12, 0.89, 0.34, ...]'),
    ('HTAP 原理', 'TiFlash 列存引擎...', '[0.45, -0.56, 0.12, 0.78, ...]'),
    ('SQL 优化', '执行计划与索引...', '[0.67, -0.23, 0.34, 0.56, ...]');

2.3 向量函数

TiDB 提供了丰富的向量操作函数：

-- 创建向量（从字符串）
SELECT VEC_COSINE_DISTANCE(
    '[0.1, 0.2, 0.3]',
    '[0.1, 0.2, 0.3]'
) AS distance;
-- 输出: 0 (完全相同)

-- 常用的距离度量
-- 余弦距离 (Cosine Distance): 衡量方向相似度，范围 [0, 2]
SELECT VEC_COSINE_DISTANCE(emb1, emb2) FROM ...;

-- L2 距离 (欧氏距离): 衡量绝对距离
SELECT VEC_L2_DISTANCE(emb1, emb2) FROM ...;

-- 注意: TiDB v8.5.0 暂不支持 VEC_INNER_PRODUCT
-- 如需内积相似度，可使用 1 - VEC_COSINE_DISTANCE 近似

三、向量搜索实践

3.1 相似度搜索

-- 创建示例表和数据
CREATE TABLE knowledge_base (
    id BIGINT PRIMARY KEY AUTO_RANDOM,
    question TEXT NOT NULL,
    answer TEXT NOT NULL,
    embedding VECTOR NOT NULL
);

-- 插入数据（向量由 Embedding 模型生成）
INSERT INTO knowledge_base (question, answer, embedding) VALUES
    ('如何部署 TiDB 集群？', '使用 TiUP 工具...', '[0.15, -0.28, 0.42, ...]'),
    ('TiDB 支持事务吗？', 'TiDB 支持分布式 ACID 事务...', '[0.33, -0.12, 0.56, ...]'),
    ('TiFlash 是什么？', 'TiFlash 是 TiDB 的列存引擎...', '[0.21, -0.45, 0.33, ...]'),
    ('如何进行数据迁移？', '可以使用 DM 工具...', '[0.18, -0.32, 0.61, ...]'),
    ('TiDB 的备份方案有哪些？', '支持 BR 物理备份和 Dumpling 逻辑备份...', '[0.25, -0.18, 0.47, ...]');

-- 搜索与查询最相似的文档
-- 假设查询 "TiDB 集群部署方法" 的向量为:
SET @query_vec = '[0.16, -0.30, 0.40, 0.55, -0.22, ...]';

SELECT
    id,
    question,
    VEC_COSINE_DISTANCE(embedding, @query_vec) AS similarity
FROM knowledge_base
ORDER BY similarity ASC  -- 距离越小越相似
LIMIT 3;

3.2 向量索引

对向量列建立索引可以大幅提升搜索速度：

-- 注意: TiDB v8.5.0 Community (Playground) 暂不支持 SQL 方式创建向量索引
-- ALTER TABLE ... ADD VECTOR INDEX 语法在 Community 版本中不可用
-- 向量索引 (HNSW) 需要 TiFlash 组件支持

-- 在 Playground 环境中，向量搜索仍然可以工作，只是没有索引加速
-- 对于小规模数据（万级以下），全表扫描的速度仍然可以接受

-- 生产环境建议:
-- 1. 部署 TiFlash 组件以启用向量索引
-- 2. 或对大规模数据使用专用向量数据库（如 Milvus）配合 TiDB 使用

HNSW 参数说明（生产环境启用 TiFlash 后使用）：

-- 创建带参数的向量索引（需要 TiFlash 支持）
ALTER TABLE knowledge_base
    ADD VECTOR INDEX idx_embedding (embedding)
    ALGORITHM HNSW
    WITH (
        m = 16,              -- 每个节点的最大连接数，越大搜索越精确但存储越多
        ef_construction = 200 -- 构建索引时的搜索范围，越大构建越慢但索引质量越好
    );

3.3 混合搜索（向量 + 结构化过滤）

这是 TiDB 向量搜索的最大优势 — 可以在同一个查询中结合向量和结构化条件：

-- 搜索特定分类下的相似文档
SELECT
    id,
    question,
    category,
    VEC_COSINE_DISTANCE(embedding, @query_vec) AS similarity
FROM knowledge_base
WHERE category = '部署运维'
ORDER BY similarity ASC
LIMIT 5;

-- 结合时间范围
SELECT
    id,
    question,
    VEC_COSINE_DISTANCE(embedding, @query_vec) AS similarity
FROM knowledge_base
WHERE created_at > '2024-01-01'
ORDER BY similarity ASC
LIMIT 5;

-- 多条件组合
SELECT
    id,
    question,
    author,
    VEC_COSINE_DISTANCE(embedding, @query_vec) AS similarity
FROM knowledge_base
WHERE status = 'published'
    AND category IN ('架构', '运维')
    AND view_count > 100
ORDER BY similarity ASC
LIMIT 10;

四、完整实践：构建 AI 知识库问答系统

4.1 架构设计

用户提问
  │
  v
┌──────────────┐
│  应用层       │
│  (Python)    │
└──────┬───────┘
       │
       │ 1. 将用户问题转为向量 (调用 Embedding API)
       v
┌──────────────┐
│   TiDB       │
│  向量搜索     │  2. 搜索最相关的文档
└──────┬───────┘
       │
       │ 3. 返回相关文档 + 原始问题
       v
┌──────────────┐
│  LLM (GPT)   │  4. 生成回答 (RAG)
└──────┬───────┘
       │
       v
   最终回答

4.2 Python 实现

import mysql.connector
import requests
import json
from typing import List, Tuple

class TiDBVectorSearch:
    """TiDB 向量搜索封装"""

    def __init__(self, host: str, port: int, user: str, password: str, database: str):
        self.conn = mysql.connector.connect(
            host=host, port=port, user=user, password=password, database=database
        )
        self.cursor = self.conn.cursor(dictionary=True)

    def get_embedding(self, text: str) -> List[float]:
        """调用 Embedding API 获取向量"""
        # 这里以 OpenAI 兼容的 API 为例，可以替换为其他服务
        response = requests.post(
            "https://api.openai.com/v1/embeddings",
            headers={"Authorization": f"Bearer {API_KEY}"},
            json={"input": text, "model": "text-embedding-ada-002"}
        )
        return response.json()["data"][0]["embedding"]

    def insert_document(self, title: str, content: str, embedding: List[float]):
        """插入文档"""
        vec_str = json.dumps(embedding)
        self.cursor.execute(
            "INSERT INTO knowledge_base (title, content, embedding) VALUES (%s, %s, %s)",
            (title, content, vec_str)
        )
        self.conn.commit()

    def search(self, query: str, top_k: int = 5) -> List[dict]:
        """搜索相似文档"""
        query_vec = json.dumps(self.get_embedding(query))
        self.cursor.execute(f"""
            SELECT id, title, content,
                   VEC_COSINE_DISTANCE(embedding, %s) AS similarity
            FROM knowledge_base
            ORDER BY similarity ASC
            LIMIT %s
        """, (query_vec, top_k))
        return self.cursor.fetchall()

    def close(self):
        self.cursor.close()
        self.conn.close()


# 使用示例
searcher = TiDBVectorSearch(
    host="127.0.0.1", port=4000,
    user="root", password="", database="ai_demo"
)

# 插入文档
docs = [
    ("TiDB 架构", "TiDB 采用计算存储分离架构..."),
    ("事务处理", "TiDB 支持乐观和悲观两种事务模式..."),
    ("HTAP 分析", "TiFlash 提供列存引擎用于分析查询..."),
    ("数据迁移", "DM 工具支持从 MySQL 全量+增量迁移..."),
    ("备份恢复", "BR 工具提供物理备份恢复能力..."),
]

for title, content in docs:
    embedding = searcher.get_embedding(content)
    searcher.insert_document(title, content, embedding)

# 搜索
results = searcher.search("TiDB 是怎么处理分布式事务的？", top_k=3)
for r in results:
    print(f"标题: {r['title']}")
    print(f"相似度: {r['similarity']:.4f}")
    print(f"内容: {r['content'][:100]}...")
    print()

searcher.close()

4.3 结合 LLM 的 RAG 问答

def rag_answer(query: str, searcher: TiDBVectorSearch) -> str:
    """基于检索增强生成 (RAG) 回答"""
    # 1. 搜索相关文档
    docs = searcher.search(query, top_k=3)

    # 2. 构建 Prompt
    context = "\n\n".join([f"- {d['title']}: {d['content']}" for d in docs])

    prompt = f"""基于以下参考信息回答问题。如果参考信息中没有相关内容，请说明无法回答。

参考信息:
{context}

问题: {query}

回答:"""

    # 3. 调用 LLM 生成回答
    response = requests.post(
        "https://api.openai.com/v1/chat/completions",
        headers={"Authorization": f"Bearer {API_KEY}"},
        json={
            "model": "gpt-4",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.7
        }
    )

    return response.json()["choices"][0]["message"]["content"]

# 使用
answer = rag_answer("TiDB 的事务机制是怎样的？", searcher)
print(answer)

五、与 AI 框架集成

5.1 集成 LangChain

# 注意: TiDB 向量存储集成可能需要安装额外的包
# pip install langchain langchain-openai
# 或使用 MySQL 兼容的向量存储方案

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.text_splitter import CharacterTextSplitter

# 使用 MySQL 兼容的向量存储
# TiDB 兼容 MySQL 协议，可以使用 MySQLVectorStore 或自定义实现
# 也可以直接使用 TiDB 的 SQL 接口进行向量操作

# 示例: 使用自定义 TiDB 向量存储
class TiDBVectorStore:
    """TiDB 向量存储实现 — 利用 TiDB 原生 VECTOR 类型"""

    def __init__(self, connection_string: str, embedding, table_name: str):
        self.table_name = table_name
        self.embedding = embedding
        # 使用 mysql-connector-python 连接
        import mysql.connector
        from urllib.parse import urlparse
        parsed = urlparse(connection_string)
        self.conn = mysql.connector.connect(
            host=parsed.hostname,
            port=parsed.port or 4000,
            user=parsed.username,
            password=parsed.password or "",
            database=parsed.path.lstrip("/")
        )

    def add_texts(self, texts: list):
        """添加文本并生成向量"""
        cursor = self.conn.cursor()
        for text in texts:
            vec = self.embedding.embed_query(text)
            vec_str = str(vec)
            cursor.execute(
                f"INSERT INTO {self.table_name} (content, embedding) VALUES (%s, %s)",
                (text, vec_str)
            )
        self.conn.commit()

    def similarity_search(self, query: str, k: int = 4):
        """相似度搜索"""
        vec = self.embedding.embed_query(query)
        vec_str = str(vec)
        cursor = self.conn.cursor(dictionary=True)
        cursor.execute(f"""
            SELECT content,
                   VEC_COSINE_DISTANCE(embedding, %s) AS distance
            FROM {self.table_name}
            ORDER BY distance ASC
            LIMIT %s
        """, (vec_str, k))
        return cursor.fetchall()

# 初始化 Embedding
embeddings = OpenAIEmbeddings()

# 连接 TiDB 向量存储
vector_store = TiDBVectorStore(
    connection_string="mysql://root@127.0.0.1:4000/ai_demo",
    embedding=embeddings,
    table_name="knowledge_base"
)

# 创建检索器
retriever = vector_store  # 简化示例，实际需实现 as_retriever 接口

# 创建 RAG 链
qa_chain = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model="gpt-4", temperature=0),
    chain_type="stuff",
    retriever=retriever
)

# 问答
result = qa_chain.run("TiDB 如何实现分布式事务？")
print(result)

5.2 集成 LlamaIndex

# 注意: 需要安装 llama-index 及相关集成包
# pip install llama-index llama-index-llms-openai llama-index-embeddings-openai

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# TiDB 向量存储实现 — 利用 TiDB 原生 VECTOR 类型
class TiDBVectorStore:
    """自定义 TiDB 向量存储，适配 LlamaIndex"""

    def __init__(self, connection_string: str, table_name: str):
        self.table_name = table_name
        import mysql.connector
        from urllib.parse import urlparse
        parsed = urlparse(connection_string)
        self.conn = mysql.connector.connect(
            host=parsed.hostname,
            port=parsed.port or 4000,
            user=parsed.username,
            password=parsed.password or "",
            database=parsed.path.lstrip("/")
        )

    def add(self, nodes):
        """添加节点到向量存储"""
        cursor = self.conn.cursor()
        for node in nodes:
            # 这里简化处理，实际需要调用 Embedding 模型
            vec_str = str(node.embedding) if hasattr(node, 'embedding') else "[]"
            cursor.execute(
                f"INSERT INTO {self.table_name} (content, embedding) VALUES (%s, %s)",
                (node.text, vec_str)
            )
        self.conn.commit()

    def query(self, query_str: str, similarity_top_k: int = 2):
        """查询相似文档"""
        # 简化示例，实际需要调用 Embedding 模型
        cursor = self.conn.cursor(dictionary=True)
        cursor.execute(f"""
            SELECT content,
                   VEC_COSINE_DISTANCE(embedding, '[]') AS distance
            FROM {self.table_name}
            ORDER BY distance ASC
            LIMIT %s
        """, (similarity_top_k,))
        return cursor.fetchall()

# 连接 TiDB
tidb_vector_store = TiDBVectorStore(
    connection_string="mysql://root@127.0.0.1:4000/ai_demo",
    table_name="knowledge_base"
)

# 创建索引（使用自定义存储）
# 注意: 完整集成需要实现 LlamaIndex 的 VectorStore 接口
index = VectorStoreIndex.from_documents(
    documents=[],  # 实际使用时传入文档
    embed_model=OpenAIEmbedding()
)

# 查询引擎
query_engine = index.as_query_engine(
    llm=OpenAI(model="gpt-4"),
    similarity_top_k=3
)

# 问答
response = query_engine.query("TiDB 的备份恢复方案有哪些？")
print(response)

六、向量搜索性能优化

6.1 索引构建建议

-- 大数据量时，构建索引前先插入数据，再创建索引
-- 而不是在插入过程中维护索引

-- 1. 批量插入数据
INSERT INTO knowledge_base ... ;  -- 大量数据

-- 2. 插入完成后创建索引
ALTER TABLE knowledge_base
    ADD VECTOR INDEX idx_embedding (embedding);

6.2 搜索参数调优

-- 查询时设置 HNSW 搜索参数
SET @query_vec = '[...]';

SELECT /*+ HNSW_EF_SEARCH(100) */
    id, question,
    VEC_COSINE_DISTANCE(embedding, @query_vec) AS similarity
FROM knowledge_base
ORDER BY similarity ASC
LIMIT 10;

6.3 维度选择

维度越低：存储越小，搜索越快，但精度可能下降。

七、与其他向量数据库对比

TiDB 的独特优势：向量搜索与结构化查询在同一个 SQL 中完成，无需维护两套系统。