使用Chroma和SelfQueryRetriever构建智能向量数据库应用

技术背景介绍

在构建AI应用时，向量数据库（Vector Database）是一个相当关键的组件。它能够存储文本、图像等数据的向量表示，从而支持高效的相似度搜索。Chroma是一个用于构建AI应用的向量数据库，可以方便地存储和检索嵌入向量。

在本文中，我们将展示如何使用SelfQueryRetriever结合Chroma向量数据库，实现一个智能的文档检索系统。

核心原理解析

Chroma向量数据库

Chroma允许将文档以向量的形式存储，并基于向量的相似度进行高效检索。这对于需要快速查找相似内容的AI应用非常有用。

SelfQueryRetriever

SelfQueryRetriever是LangChain库中的一个可用对象，它能够基于预定义的文档元数据和内容描述进行智能查询，自动构建查询表达式并过滤检索结果。

代码实现演示

1. 安装必要依赖

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

2. 获取OpenAI API Key

import getpass
import os

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

3. 创建Chroma向量存储

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

# 初始化OpenAI嵌入
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},
    )
    # 更多文档...
]

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

4. 创建自查询检索器

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",
    ),
    AttributeInfo(
        name="director",
        description="The name of the movie director",
        type="string",
    ),
    AttributeInfo(
        name="rating", description="A 1-10 rating for the movie", type="float"
    ),
]

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
)

5. 测试检索功能

# 基于查询进行检索
result = retriever.invoke("What are some movies about dinosaurs")
print(result)

# 基于过滤进行检索
result = retriever.invoke("I want to watch a movie rated higher than 8.5")
print(result)

# 基于查询和过滤进行检索
result = retriever.invoke("Has Greta Gerwig directed any movies about women")
print(result)

应用场景分析

• 电影推荐系统：通过嵌入向量表示电影剧情和相关信息，快速推荐用户感兴趣的电影。
• 文档管理系统：在大量文档中，通过内容检索和元数据过滤，快速找到相关文档。
• 智能问答系统：基于用户问题及文档元数据，提供精准答案。

实践建议

1. 选择合适的嵌入模型：根据具体应用场景，选择适合的嵌入模型（如OpenAI嵌入）。
2. 合理设计元数据结构：确保文档元数据尽可能详细和规范，方便智能检索。
3. 优化查询机制：结合实际需求调整查询和过滤策略，提高检索效率和准确性。

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