使用Rockset作为Langchain的文档加载器-优快云博客

使用Rockset作为Langchain的文档加载器

1. 技术背景介绍

Rockset是一个实时分析数据库，能够在没有运营负担的情况下对海量半结构化数据进行查询。数据在被摄取后的不到一秒内即可查询，同时，针对这些数据的分析查询通常在毫秒级完成。Rockset对计算进行了优化，使其适用于在100TB以下（或通过汇总超过100TB）的高并发应用。

2. 核心原理解析

使用Rockset与Langchain集成，可以通过SQL查询从Rockset集合中加载文档。核心在于构建一个RocksetLoader对象，该对象通过与Rockset的Python客户端进行交互执行查询并加载数据。

3. 代码实现演示

环境设置

首先，你需要确保已经拥有一个Rockset账号并获取了API Key。还需要知道你的API区域。本示例假设你使用的是位于Oregon的us-west-2区域。

# 安装必要的库
%pip install --upgrade --quiet rockset

然后设置环境变量ROCKSET_API_KEY。

加载文档

通过以下代码示例初始化一个RocksetLoader，并执行SQL查询以加载文档：

from langchain_community.document_loaders import RocksetLoader
from rockset import Regions, RocksetClient, models

# 初始化Rockset客户和Loader
loader = RocksetLoader(
    RocksetClient(Regions.usw2a1, "your-api-key"),  # 替换为你的API密钥
    models.QueryRequestSql(query="SELECT * FROM langchain_demo LIMIT 3"),  # 执行SQL查询
    ["text"],  # 内容列
    metadata_keys=["id", "date"]  # 元数据列
)

# 以惰性方式加载文档
docs_iterator = loader.lazy_load()

# 或一次性加载所有文档
docs_list = loader.load()

# 打印加载的文档
for doc in docs_list:
    print(doc.page_content)
    print(doc.metadata)

这里，SQL查询SELECT * FROM langchain_demo LIMIT 3 被执行，集合中的text列被用作页面内容，而记录的id和date列被用作元数据。

使用多个列作为内容

你亦可以选择使用多个列作为内容：

loader = RocksetLoader(
    RocksetClient(Regions.usw2a1, "your-api-key"),
    models.QueryRequestSql(query="SELECT * FROM langchain_demo LIMIT 1 WHERE id=38"),
    ["sentence1", "sentence2"],  # 两个内容列
)

假设sentence1字段内容为"This is the first sentence."，sentence2字段内容为"This is the second sentence."，结果文档的page_content将会是：

This is the first sentence.
This is the second sentence.

你可以自定义连接多个内容列的方式，例如将内容用空格连接：

RocksetLoader(
    RocksetClient(Regions.usw2a1, "your-api-key"),
    models.QueryRequestSql(query="SELECT * FROM langchain_demo LIMIT 1 WHERE id=38"),
    ["sentence1", "sentence2"],
    content_columns_joiner=lambda docs: " ".join([doc[1] for doc in docs])
)

结果将是：This is the first sentence. This is the second sentence.

如果你想包含列名，可以像这样设置：

RocksetLoader(
    RocksetClient(Regions.usw2a1, "your-api-key"),
    models.QueryRequestSql(query="SELECT * FROM langchain_demo LIMIT 1 WHERE id=38"),
    ["sentence1", "sentence2"],
    content_columns_joiner=lambda docs: "\n".join([f"{doc[0]}: {doc[1]}" for doc in docs])
)

结果将是：

sentence1: This is the first sentence.
sentence2: This is the second sentence.

4. 应用场景分析

Rockset与Langchain的集成非常适合用于需要实时查询和处理海量数据的应用场景。例如：

实时数据分析和报表
大规模数据挖掘和机器学习训练数据准备
高并发的数据查询服务

5. 实践建议

确保SQL查询高效，尽量避免扫描大量数据
正确配置元数据列，以便后续处理更方便
在使用多个内容列时，合理设置连接方式以保证内容可读性

如果遇到问题欢迎在评论区交流。