Hybrid Search

通过使用全文检索，你可以基于精确关键词检索文档。通过使用向量检索，你可以基于语义相似度检索文档。那么，能否将这两种检索方式结合起来，以提升检索质量并覆盖更多场景？答案是可以，这种方式被称为混合检索（hybrid search），并且在 AI 应用中被广泛使用。

在 TiDB 中，混合检索的一般流程如下：

使用 TiDB 进行 全文检索 和 向量检索。
使用 reranker（重排序器）对两种检索结果进行融合。

本教程演示如何在 TiDB 中使用 pytidb Python SDK 实现混合检索，该 SDK 内置了 embedding 和 reranking 支持。使用 pytidb 并非强制要求 —— 你也可以直接通过 SQL 进行检索，并使用自定义的 reranking 模型。

前置条件

混合检索依赖于全文检索和向量检索。全文检索目前仍处于早期阶段，我们正在持续向更多用户开放。目前，全文检索仅在以下产品选项和区域可用：

TiDB Cloud Serverless：Frankfurt (eu-central-1) 和 Singapore (ap-southeast-1)

要完成本教程，请确保你拥有位于支持区域的 TiDB Cloud Serverless 集群。如果还没有，请参考创建 TiDB Cloud Serverless 集群进行创建。

快速开始

步骤 1. 安装 pytidb Python SDK

pip install "pytidb[models]"

# （可选）如果你不想使用内置的 embedding 函数和 reranker：
# pip install pytidb

# （可选）如需将查询结果转换为 pandas DataFrame：
# pip install pandas

步骤 2. 连接 TiDB

from pytidb import TiDBClient

db = TiDBClient.connect(
    host="HOST_HERE",
    port=4000,
    username="USERNAME_HERE",
    password="PASSWORD_HERE",
    database="DATABASE_HERE",
)

你可以在 TiDB Cloud 控制台获取这些连接参数：

进入 Clusters 页面，点击目标集群名称进入集群概览页。

点击右上角的 Connect，弹出连接对话框，显示连接参数。

例如，连接参数如下所示：

HOST:     gateway01.us-east-1.prod.shared.aws.tidbcloud.com
PORT:     4000
USERNAME: 4EfqPF23YKBxaQb.root
PASSWORD: abcd1234
DATABASE: test
CA:       /etc/ssl/cert.pem

则连接 TiDB Cloud Serverless 集群的 Python 代码如下：

db = TiDBClient.connect(
    host="gateway01.us-east-1.prod.shared.aws.tidbcloud.com",
    port=4000,
    username="4EfqPF23YKBxaQb.root",
    password="abcd1234",
    database="test",
)

注意，上述示例仅用于演示。你需要使用自己的参数，并妥善保管。

步骤 3. 创建数据表

以创建名为 chunks 的表为例，包含以下字段：

id (int)：chunk 的 ID。
text (text)：chunk 的文本内容。
text_vec (vector)：文本的向量表示，由 pytidb 中的 embedding 模型自动生成。
user_id (int)：创建该 chunk 的用户 ID。

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)

步骤 4. 插入数据

table.bulk_insert(
    [
        Chunk(id=2, text="bar", user_id=2),   # 👈 text 字段会自动嵌入为向量
        Chunk(id=3, text="baz", user_id=3),   # 并存储到 "text_vec" 字段
        Chunk(id=4, text="qux", user_id=4),   # 中。
    ]
)

步骤 5. 执行混合检索

本例中，使用 jina-reranker 模型对检索结果进行重排序。

from pytidb.rerankers import Reranker

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

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

完整示例请参考 pytidb hybrid search demo。

反馈与帮助

全文检索目前仍处于早期阶段，开放范围有限。如果你希望在尚未开放的区域体验全文检索，或有任何反馈和帮助需求，欢迎联系我们：