如何不依赖LangChain，直接用OpenAI API在ChromaDB生成嵌入？

问题描述 #

之前我通过LangChain的OpenAIEmbeddings模型，用下面的代码就能把向量嵌入存储到ChromaDB，不需要手动输入文本：

db = Chroma(persist_directory=CHROMA_PATH, embedding_function=get_embedding())

现在我想脱离LangChain，直接调用OpenAI官方API生成嵌入，调用方式如下：

client = OpenAI()

response = client.embeddings.create(
    model=model,
    input=[text]
)

但这里的问题是，OpenAI API必须传入文本才能生成嵌入，而ChromaDB的embedding_function参数对应的函数看起来不需要显式传入文本。我试过遍历文档逐个生成嵌入并添加到元数据：

for chunk in chunks_with_ids:
    chunk_embedding = get_openai_embedding(chunk.page_content)
    chunk.metadata["embeddings"] = chunk_embedding

但我想知道有没有更集成的方式，能直接用OpenAI API生成嵌入并存储到ChromaDB，而不用手动遍历处理？

解决方案 #

方法1：自定义符合Chroma要求的Embedding函数 #

ChromaDB的embedding_function参数要求传入的函数能接收文本列表作为输入，返回对应的嵌入向量列表。你可以基于OpenAI API封装一个这样的函数，之后就能像之前用LangChain那样直接集成到Chroma中，Chroma会自动调用这个函数为文本生成嵌入并存储。

示例代码：

from openai import OpenAI
import chromadb

# 初始化OpenAI客户端
client = OpenAI(api_key="你的API密钥")
CHROMA_PATH = "./chroma_db"

# 自定义嵌入函数
def openai_embedding_function(texts):
    # 调用OpenAI Embeddings API
    response = client.embeddings.create(
        model="text-embedding-3-small",  # 根据需求选择模型
        input=texts
    )
    # 提取嵌入向量并返回
    return [data.embedding for data in response.data]

# 初始化Chroma，传入自定义嵌入函数
db = chromadb.PersistentClient(path=CHROMA_PATH).get_or_create_collection(
    name="your_collection_name",
    embedding_function=openai_embedding_function
)

# 添加文档时，Chroma会自动调用嵌入函数生成向量
db.add(
    documents=["文档1内容", "文档2内容"],
    metadatas=[{"source": "doc1"}, {"source": "doc2"}],
    ids=["id1", "id2"]
)

这种方式完全不需要手动处理嵌入生成，Chroma会在你调用add方法时自动触发嵌入函数，和之前用LangChain的体验一致。

方法2：批量生成嵌入后直接存入Chroma #

如果你需要更灵活的控制（比如自定义批量大小、处理异常），可以先批量调用OpenAI API生成所有嵌入，再直接传入Chroma的add方法（此时不需要指定embedding_function）：

示例代码：

from openai import OpenAI
import chromadb

client = OpenAI(api_key="你的API密钥")
CHROMA_PATH = "./chroma_db"

# 准备文档数据
documents = ["文档1内容", "文档2内容"]
metadatas = [{"source": "doc1"}, {"source": "doc2"}]
ids = ["id1", "id2"]

# 批量生成嵌入
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=documents
)
embeddings = [data.embedding for data in response.data]

# 初始化Chroma集合（无需指定embedding_function）
db = chromadb.PersistentClient(path=CHROMA_PATH).get_or_create_collection(
    name="your_collection_name"
)

# 直接传入预生成的嵌入
db.add(
    documents=documents,
    metadatas=metadatas,
    ids=ids,
    embeddings=embeddings
)

这种方式适合需要对嵌入生成过程做额外处理的场景，比如分批处理大数量文档、添加重试逻辑等。

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