在 Astra DB 中使用 SelfQueryRetriever 进行向量存储与查询-优快云博客

本文链接：https://blog.youkuaiyun.com/fGVBSAbe/article/details/145333044

# 在 Astra DB 中使用 SelfQueryRetriever 进行向量存储与查询

## 技术背景介绍

DataStax Astra DB 是一个无服务器的向量数据库，基于 Cassandra 构建，并通过易用的 JSON API 提供服务。在这篇文章中，我们将演示如何在 Astra DB 中创建向量存储，并使用 SelfQueryRetriever 进行查询。我们的目标是展示如何对一组电影摘要数据进行向量存储，并通过自然语言进行检索。

## 核心原理解析

SelfQueryRetriever 结合了自然语言处理与向量检索技术，允许用户通过简单的自然语言查询从大规模数据集中提取信息。这需要将数据转换为向量表示，并存储于支持向量操作的数据库中。Astra DB 提供了无缝的向量存储功能，支持高效的数据查询。

## 代码实现演示

首先，我们将创建一个 Astra DB 向量存储，初始化一组包含电影摘要的文档。接下来，我们使用 OpenAI Embeddings 来生成这些文档的向量表示。

```python
# 安装必要的依赖库
%pip install --upgrade --quiet lark astrapy langchain-openai

# 引入库并获取 OpenAI API Key
import os
from getpass import getpass
from langchain_openai.embeddings import OpenAIEmbeddings

os.environ["OPENAI_API_KEY"] = getpass("OpenAI API Key:")
embeddings = OpenAIEmbeddings()

# 配置 Astra DB API Endpoint 和 Token
ASTRA_DB_API_ENDPOINT = input("ASTRA_DB_API_ENDPOINT = ")
ASTRA_DB_APPLICATION_TOKEN = getpass("ASTRA_DB_APPLICATION_TOKEN = ")

from langchain_community.vectorstores import AstraDB
from langchain_core.documents import Document

# 定义电影摘要文档
docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={"year": 1993, "rating": 7.7, "genre": "science fiction"},
    ),
    # 其他文档省略以节省篇幅
]

# 创建向量存储
vectorstore = AstraDB.from_documents(
    docs,
    embeddings,
    collection_name="astra_self_query_demo",
    api_endpoint=ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
)

# 描述文档的元数据字段
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import OpenAI

metadata_field_info = [
    AttributeInfo(
        name="genre",
        description="The genre of the movie",
        type="string or list[string]",
    ),
    # 其他字段省略
]

document_content_description = "Brief summary of a movie"
llm = OpenAI(temperature=0)

# 创建 SelfQueryRetriever 实例
retriever = SelfQueryRetriever.from_llm(
    llm, vectorstore, document_content_description, metadata_field_info, verbose=True
)

# 进行查询
print(retriever.invoke("What are some movies about dinosaurs?"))
print(retriever.invoke("I want to watch a movie rated higher than 8.5"))