构建基于嵌入向量的AI应用: 使用Chroma和SelfQueryRetriever

本文链接：https://blog.csdn.net/antja_/article/details/149271043

在AI应用程序中，处理和检索跨越大量文档的复杂查询是一个常见的任务。Chroma是一个功能强大的向量数据库，专注于通过嵌入向量建模来构建AI应用程序。本文将详细介绍如何使用Chroma创建一个向量存储，并演示如何用SelfQueryRetriever对其进行查询。

技术背景介绍

Chroma利用嵌入向量实现高效的文档检索，使得基于内容的查询更精准。SelfQueryRetriever是一个能够根据自定义元数据和内容描述自动检索相关文档的工具。结合这两者，我们可以实现复杂的查询机制。

核心原理解析

Chroma通过将文档转换为向量存储创建一个易于查询的数据库。SelfQueryRetriever使用描述性语言模型（LLM）来解析查询问题，并根据元数据和描述进行检索。

代码实现演示

下面是一个完整的代码示例，展示如何使用Chroma创建一个向量存储，并使用SelfQueryRetriever进行查询。

创建Chroma向量存储

首先，我们需要创建Chroma向量存储，并导入初始数据。该数据集包含几部电影的摘要：

%pip install --upgrade --quiet lark
%pip install --upgrade --quiet langchain-chroma

import os
import getpass
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings

# 设定OpenAI API密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

# 初始化OpenAI Embeddings
embeddings = OpenAIEmbeddings()

# 创建文档集
docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={"year": 1993, "rating": 7.7, "genre": "science fiction"},
    ),
    Document(
        page_content="Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
        metadata={"year": 2010, "director": "Christopher Nolan", "rating": 8.2},
    ),
    # 其他文档...
]

# 创建Chroma向量存储
vectorstore = Chroma.from_documents(docs, embeddings)

创建SelfQueryRetriever

接下来，我们使用SelfQueryRetriever进行自定义查询：

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]"),
    AttributeInfo(name="year", description="The year the movie was released", type="integer"),
    # 其他属性信息...
]

# 描述文档内容
document_content_description = "Brief summary of a movie"
llm = OpenAI(temperature=0)

# 初始化检索器
retriever = SelfQueryRetriever.from_llm(
    llm, vectorstore, document_content_description, metadata_field_info, verbose=True
)

# 测试检索器
results = retriever.invoke("What are some movies about dinosaurs")
print(results)