在LangChain中使用OpenAI Embeddings与Chroma向量存储限制每分钟令牌数

核心思路 #

不用直接依赖Chroma.from_documents一次性处理所有文档，而是通过分批追加文档或自定义速率控制的Embedding类来控制每分钟令牌消耗，避免触发OpenAI的速率限制。

方法1：分批处理文档并控制速率 #

Chroma的add_documents方法是追加而非覆盖数据，因此可以先初始化空向量库，再分批次添加文档，每批处理后根据令牌消耗等待对应时间。

代码实现 #

import time
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma

# 拆分文档为块
text_splitter = CharacterTextSplitter(chunk_size=1500, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

# 初始化Embeddings和空Chroma库（指定持久化目录避免数据丢失）
embeddings = OpenAIEmbeddings()
db = Chroma(embedding_function=embeddings, persist_directory="./chroma_db")

# 配置速率参数（适配OpenAI ada-002的令牌限制）
max_tokens_per_minute = 1000000
# 估算单文档块令牌数：1500字符≈1125令牌（按1.33字符/令牌的通用比例）
tokens_per_chunk = 1125
# 计算每批可处理的文档数量
batch_size = max_tokens_per_minute // tokens_per_chunk

# 分批处理并追加到向量库
for i in range(0, len(texts), batch_size):
    batch = texts[i:i+batch_size]
    db.add_documents(batch)
    # 等待1分钟+5秒缓冲，确保不触发速率限制
    time.sleep(65)

# 最终持久化数据
db.persist()

方法2：自定义带速率控制的Embedding包装器 #

封装OpenAIEmbeddings类，在生成Embedding时自动计算令牌消耗，超过限制时自动等待，无需手动分批。

代码实现 #

import time
import tiktoken
from langchain.embeddings import OpenAIEmbeddings

class RateLimitedOpenAIEmbeddings(OpenAIEmbeddings):
    def __init__(self, max_tokens_per_minute=1000000, **kwargs):
        super().__init__(**kwargs)
        self.max_tokens_per_minute = max_tokens_per_minute
        self.tokens_used = 0
        self.last_reset_time = time.time()

    def _count_tokens(self, texts):
        # 用tiktoken精确计算令牌数（需先安装tiktoken：pip install tiktoken）
        encoding = tiktoken.encoding_for_model(self.model)
        return sum(len(encoding.encode(text.page_content)) for text in texts)

    def embed_documents(self, texts):
        current_time = time.time()
        
        # 每分钟重置令牌计数
        if current_time - self.last_reset_time >= 60:
            self.tokens_used = 0
            self.last_reset_time = current_time

        batch_tokens = self._count_tokens(texts)
        # 若当前批次令牌超过剩余额度，等待到下一分钟
        if self.tokens_used + batch_tokens > self.max_tokens_per_minute:
            wait_time = 60 - (current_time - self.last_reset_time)
            time.sleep(wait_time + 5)  # 加5秒缓冲
            self.tokens_used = 0
            self.last_reset_time = time.time()

        self.tokens_used += batch_tokens
        return super().embed_documents([text.page_content for text in texts])

# 使用自定义速率控制的Embeddings
embeddings = RateLimitedOpenAIEmbeddings(max_tokens_per_minute=1000000)
db = Chroma.from_documents(texts, embeddings)

关键说明 #

1. 两种方法都不会覆盖已有数据：add_documents是追加模式，Chroma.from_documents仅在首次创建时初始化库。
2. 方法1适合快速实现，方法2更灵活，可复用在其他需要Embedding的场景中。
3. 用tiktoken计算令牌数比估算更准确，建议优先使用。

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