作为一个实时搜索和分析数据库，Rockset 使用索引来提供可扩展且高性能的个性化、产品搜索、语义搜索、聊天机器人应用等。由于 Rockset 专为实时场景而构建，你可以在不断更新的流式数据上构建这些响应迅速的应用程序。通过将 Rockset 与 LlamaIndex 集成，你可以轻松地基于自己的实时数据使用 LLM，构建可用于生产的向量搜索应用。

我们将演示如何在 LlamaIndex 中将 Rockset 用作向量存储。

教程

在本例中，我们将使用 OpenAI 的 text-embedding-ada-002 模型生成嵌入，并使用 Rockset 作为向量存储来保存嵌入。我们将从文件中摄取文本，并针对内容提问。

设置环境

1. 在 Rockset 控制台中创建一个集合，并将 Write API 作为数据源。将集合命名为 llamaindex_demo。使用 VECTOR_ENFORCE 配置以下摄取转换，以定义嵌入字段并利用性能和存储优化：

SELECT
    _input.* EXCEPT(_meta),
    VECTOR_ENFORCE(
        _input.embedding,
        1536,
        'float'
    ) as embedding
FROM _input

2. 在 Rockset 控制台中创建 API 密钥，并设置 ROCKSET_API_KEY 环境变量。在此处找到你的 API 服务器，并设置 ROCKSET_API_SERVER 环境变量。设置 OPENAI_API_KEY 环境变量。

3. 安装依赖项。

终端窗口

pip3 install llama_index rockset

4. LlamaIndex 允许你从多种来源摄取数据。在本例中，我们将读取名为 constitution.txt 的文本文件，它是美国宪法的转录文本，可在此处找到。

数据摄取

使用 LlamaIndex 的 SimpleDirectoryReader 类将文本文件转换为 Document 对象列表。

%pip install llama-index-llms-openai
%pip install llama-index-vector-stores-rocksetdb

from llama_index.core import SimpleDirectoryReader

docs = SimpleDirectoryReader(
    input_files=["{path to}/consitution.txt"]
).load_data()

实例化 LLM 和服务上下文。

from llama_index.core import Settings
from llama_index.llms.openai import OpenAI

Settings.llm = OpenAI(temperature=0.8, model="gpt-3.5-turbo")

实例化向量存储和存储上下文。

from llama_index.core import StorageContext
from llama_index.vector_stores.rocksetdb import RocksetVectorStore

vector_store = RocksetVectorStore(collection="llamaindex_demo")
storage_context = StorageContext.from_defaults(vector_store=vector_store)

将文档添加到 llamaindex_demo 集合，并创建索引。

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(
    docs,
    storage_context=storage_context,
)

查询

针对文档提出问题并生成回答。

response = index.as_query_engine().query("What is the duty of the president?")

print(str(response))

运行程序。

$ python3 main.py
The duty of the president is to faithfully execute the Office of President of the United States, preserve, protect and defend the Constitution of the United States, serve as the Commander in Chief of the Army and Navy, grant reprieves and pardons for offenses against the United States (except in cases of impeachment), make treaties and appoint ambassadors and other public ministers, take care that the laws be faithfully executed, and commission all the officers of the United States.

元数据过滤

元数据过滤允许你检索符合特定过滤条件的相关文档。

1. 将节点添加到向量存储并创建索引。

from llama_index.vector_stores.rocksetdb import RocksetVectorStore
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.vector_stores.types import NodeWithEmbedding
from llama_index.core.schema import TextNode

nodes = [
    NodeWithEmbedding(
        node=TextNode(
            text="Apples are blue",
            metadata={"type": "fruit"},
        ),
        embedding=[],
    )
]
index = VectorStoreIndex(
    nodes,
    storage_context=StorageContext.from_defaults(
        vector_store=RocksetVectorStore(collection="llamaindex_demo")
    ),
)

2. 定义元数据过滤器。

from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="type", value="fruit")]
)

3. 检索满足过滤条件的相关文档。

retriever = index.as_retriever(filters=filters)
retriever.retrieve("What colors are apples?")

从现有集合创建索引

你可以使用现有集合中的数据创建索引。

from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.rocksetdb import RocksetVectorStore

vector_store = RocksetVectorStore(collection="llamaindex_demo")

index = VectorStoreIndex.from_vector_store(vector_store)

从新集合创建索引

你还可以创建一个新的 Rockset 集合作为向量存储。

from llama_index.vector_stores.rocksetdb import RocksetVectorStore

vector_store = RocksetVectorStore.with_new_collection(
    collection="llamaindex_demo",  # name of new collection
    dimensions=1536,  # specifies length of vectors in ingest tranformation (optional)
    # other RocksetVectorStore args
)

index = VectorStoreIndex(
    nodes,
    storage_context=StorageContext.from_defaults(vector_store=vector_store),
)

配置

• collection：要查询的集合名称（必填）。

RocksetVectorStore(collection="my_collection")

• workspace：包含该集合的工作区名称。默认为 "commons"。

RocksetVectorStore(worksapce="my_workspace")

• api_key：用于验证 Rockset 请求的 API 密钥。如果传入了 client，则忽略此参数。默认为 ROCKSET_API_KEY 环境变量。

RocksetVectorStore(api_key="<my key>")

• api_server：用于 Rockset 请求的 API 服务器。如果传入了 client，则忽略此参数。默认为 ROCKSET_API_KEY 环境变量；如果未设置 ROCKSET_API_SERVER，则默认为 "https://api.use1a1.rockset.com"。

from rockset import Regions
RocksetVectorStore(api_server=Regions.euc1a1)

• client：用于执行 Rockset 请求的 Rockset 客户端对象。如果未指定，则会在内部使用 api_key 参数（或 ROCKSET_API_SERVER 环境变量）和 api_server 参数（或 ROCKSET_API_SERVER 环境变量）构造客户端对象。

from rockset import RocksetClient
RocksetVectorStore(client=RocksetClient(api_key="<my key>"))

• embedding_col：包含嵌入的数据库字段名称。默认为 "embedding"。

RocksetVectorStore(embedding_col="my_embedding")

• metadata_col：包含节点数据的数据库字段名称。默认为 "metadata"。

RocksetVectorStore(metadata_col="node")

• distance_func：用于衡量向量关系的度量。默认为余弦相似度。

RocksetVectorStore(distance_func=RocksetVectorStore.DistanceFunc.DOT_PRODUCT)