如何基于LangChain与ChromaDB实现搜索查询的分页功能

现有代码回顾 #

你当前使用LangChain查询ChromaDB的代码如下：

from langchain.vectorstores.chroma import Chroma
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings

@app.route('/search',methods=['POST'])
def searchPost():
    query = request.json['query']
    try:
        embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
        vectordb = Chroma(persist_directory=persist_directory, 
                    embedding_function=embedding_function)    

        output = vectordb.similarity_search(query=query)

        metadata_ids = [doc.metadata['id'] for doc in output]

        return metadata_ids
    except Exception as e:
        print('error',e)
        return 'error'

分页实现方案 #

要实现分页功能，核心是利用Chroma原生支持的offset参数，结合前端传入的页码和每页条数，在数据库层面完成分页查询，避免全量加载后再切片的低效方式。具体实现步骤如下：

1. 接收分页参数 #

从请求中获取page（当前页码，默认1）和page_size（每页条数，默认10）两个参数，同时做类型转换确保为整数。

2. 计算偏移量 #

根据页码和每页条数计算数据库查询的偏移量：offset = (page - 1) * page_size（Chroma的offset从0开始计数，因此页码从1开始时需做此转换）。

3. 带分页的相似搜索 #

调用similarity_search时，传入k=page_size控制每页返回条数，同时传入offset控制从第几条开始返回结果。

4. 返回结构化分页结果（可选） #

额外返回总匹配条数、总页数等信息，方便前端构建分页导航组件。

修改后的完整代码 #

from langchain.vectorstores.chroma import Chroma
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings

@app.route('/search',methods=['POST'])
def searchPost():
    query = request.json['query']
    # 获取分页参数，设置默认值并做类型转换
    page = int(request.json.get('page', 1))
    page_size = int(request.json.get('page_size', 10))
    
    try:
        embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
        vectordb = Chroma(persist_directory=persist_directory, 
                    embedding_function=embedding_function)    
        
        # 计算数据库查询偏移量
        offset = (page - 1) * page_size
        
        # 执行带分页的相似搜索
        output = vectordb.similarity_search(
            query=query,
            k=page_size,
            offset=offset
        )

        metadata_ids = [doc.metadata['id'] for doc in output]
        
        # 获取符合条件的总文档数（若有过滤条件，需同步传入where参数）
        total = vectordb._client.count(where={})
        
        # 返回结构化分页结果
        return {
            'data': metadata_ids,
            'total': total,
            'current_page': page,
            'page_size': page_size,
            'total_pages': (total + page_size - 1) // page_size
        }
    except Exception as e:
        print('error', e)
        return {'error': str(e)}

关键说明 #

• 参数合法性校验：建议额外添加校验逻辑，确保page和page_size为正整数，避免非法输入导致的查询错误。
• 过滤条件适配：如果查询需要结合filter参数，需将相同的过滤条件同步传递给similarity_search和count方法，保证总条数统计的准确性。
• 性能优势：通过Chroma的offset参数实现数据库层面分页，相比全量查询后切片的方式，能大幅减少数据传输量和内存占用，提升查询效率。

内容的提问来源于stack exchange，提问作者Venkatesh