fix(knowledge): 优化 ChromaDB 插入逻辑,支持分批处理
This commit is contained in:
parent
e592c75662
commit
1ce9afa0d3
@ -200,8 +200,23 @@ class ChromaKB(KnowledgeBase):
|
||||
metadatas = [chunk["metadata"] for chunk in chunks]
|
||||
ids = [chunk["id"] for chunk in chunks]
|
||||
|
||||
# 插入到 ChromaDB
|
||||
collection.add(documents=documents, metadatas=metadatas, ids=ids)
|
||||
# 插入到 ChromaDB - 分批处理以避免超出 OpenAI 批次大小限制
|
||||
batch_size = 64 # OpenAI 的最大批次大小限制
|
||||
total_batches = (len(chunks) + batch_size - 1) // batch_size
|
||||
|
||||
for i in range(0, len(chunks), batch_size):
|
||||
batch_documents = documents[i:i + batch_size]
|
||||
batch_metadatas = metadatas[i:i + batch_size]
|
||||
batch_ids = ids[i:i + batch_size]
|
||||
|
||||
collection.add(
|
||||
documents=batch_documents,
|
||||
metadatas=batch_metadatas,
|
||||
ids=batch_ids
|
||||
)
|
||||
|
||||
batch_num = i // batch_size + 1
|
||||
logger.info(f"Processed batch {batch_num}/{total_batches} for {filename}")
|
||||
|
||||
logger.info(f"Inserted {content_type} {item} into ChromaDB. Done.")
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user