LlamaIndex VectorStoreIndex.from_documents() — embedding batches fail mid-ingestion

from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

# 1) Embedding — set embed_batch_size to control provider load
embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",
    embed_batch_size=32,          # smaller = safer on rate limits
    num_workers=4,                # parallelism (per-provider judgment)
)

# 2) Vector store with explicit collection
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection("my_docs")
vector_store = ChromaVectorStore(chroma_collection=collection)

# 3) Docstore — tracks what has already been ingested
docstore = SimpleDocumentStore.from_persist_path("./docstore.json") \
    if os.path.exists("./docstore.json") else SimpleDocumentStore()

# 4) Pipeline: node parser -> embedding -> vector store
pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=512, chunk_overlap=64),
        embed_model,
    ],
    docstore=docstore,           # skips already-processed docs on rerun
    vector_store=vector_store,
)

# 5) Run — safe to interrupt and resume
nodes = pipeline.run(
    documents=your_documents,
    show_progress=True,
    num_workers=4,               # docstore parallelism
)
print(f"Processed {len(nodes)} new nodes")

# 6) Persist docstore for next run
docstore.persist("./docstore.json")

# 7) Build index from the populated vector store
index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store,
    embed_model=embed_model,
)

from typing import List
import tiktoken
from llama_index.core.schema import BaseNode, TransformComponent
from llama_index.core.node_parser import SentenceSplitter

# --- Guard: verify no chunk exceeds embedding token limit ---
class TokenLimitGuard(TransformComponent):
    """Split any node whose text exceeds the embedding model's limit."""
    max_tokens: int = 8000        # keep 192 tokens of headroom under 8192
    tokenizer_name: str = "cl100k_base"

    def __call__(self, nodes: List[BaseNode], **kwargs) -> List[BaseNode]:
        enc = tiktoken.get_encoding(self.tokenizer_name)
        out = []
        for node in nodes:
            tokens = enc.encode(node.text)
            if len(tokens) <= self.max_tokens:
                out.append(node)
                continue
            # Split into halves recursively (rare — node parser should prevent this)
            print(f"⚠ Node {node.node_id[:8]}... has {len(tokens)} tokens, splitting")
            mid = len(tokens) // 2
            first_text = enc.decode(tokens[:mid])
            second_text = enc.decode(tokens[mid:])
            # Preserve metadata
            first = node.copy()
            first.text = first_text
            first.node_id = f"{node.node_id}_a"
            second = node.copy()
            second.text = second_text
            second.node_id = f"{node.node_id}_b"
            out.extend([first, second])
        return out

# --- Configure embedding with tight rate control ---
from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",
    embed_batch_size=20,             # smaller batches on strict TPM
    max_retries=6,                    # retry on rate limits (SDK-level)
    timeout=60,
    reuse_client=True,
    additional_kwargs={"encoding_format": "float"},
)

# Wire into the pipeline
pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=512, chunk_overlap=64),
        TokenLimitGuard(max_tokens=8000),   # ← guards against overflow
        embed_model,
    ],
    docstore=docstore,
    vector_store=vector_store,
)

# --- Chroma: 5461 hard limit; use 2000 for safety ---
from llama_index.vector_stores.chroma import ChromaVectorStore
chroma_store = ChromaVectorStore(chroma_collection=collection)

# --- Pinecone: 2MB request cap; 100-200 for dense vectors is safe ---
from llama_index.vector_stores.pinecone import PineconeVectorStore
pinecone_store = PineconeVectorStore(
    pinecone_index=idx,
    insert_kwargs={"batch_size": 100},   # per-batch upsert size
)

# --- Qdrant: batch_size configurable in insert ---
from llama_index.vector_stores.qdrant import QdrantVectorStore
qdrant_store = QdrantVectorStore(
    client=qdrant_client,
    collection_name="my_docs",
    batch_size=64,
)

# --- pgvector via LlamaIndex ---
from llama_index.vector_stores.postgres import PGVectorStore
pg_store = PGVectorStore.from_params(
    database="mydb",
    host="localhost",
    port=5432,
    user="user",
    password="pass",
    table_name="documents",
    embed_dim=1536,
    hybrid_search=False,
    perform_setup=True,
)

# --- On the index side ---
from llama_index.core import VectorStoreIndex, StorageContext

storage_context = StorageContext.from_defaults(vector_store=chroma_store)

# insert_batch_size on the index builder controls how many nodes go in per
# vector_store.add() call
index = VectorStoreIndex(
    nodes=nodes,
    storage_context=storage_context,
    embed_model=embed_model,
    show_progress=True,
    insert_batch_size=2000,       # ← store-appropriate
)