Putting a dollar cap on one LightRAG indexing job

[domondi1]

Indexing makes a lot of LLM calls: entity and relation extraction for every chunk, gleaning passes, and summaries, several at a time with llm_model_max_async. On a big document set the cost of one ainsert is hard to predict, and there isn't a per-job limit built in. Here's one way to give each indexing job a hard dollar budget without changing LightRAG.

openai_complete_if_cache passes its extra kwargs straight to chat.completions.create, so llm_model_kwargs can carry a base_url and per-request extra_headers (the job's id and budget) to a small OpenAI-compatible gateway running in the same process:

import asyncio
import inferrail
from lightrag import LightRAG
from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed

base_url = inferrail.start()  # gateway on a background thread, uses OPENAI_API_KEY

def rag_for_job(job_id: str, budget_usd: str, working_dir: str) -> LightRAG:
    return LightRAG(
        working_dir=working_dir,
        llm_model_func=gpt_4o_mini_complete,
        llm_model_kwargs={
            "base_url": base_url,
            "max_tokens": 2000,
            "extra_headers": {"X-Inferrail-Attribute-Work-Id": job_id,
                              "X-Inferrail-Budget-Usd": budget_usd},
        },
        embedding_func=openai_embed,  # embeddings go straight to OpenAI
    )

async def main():
    rag = rag_for_job("index-2026-10-03", "2.00", "./rag_storage")
    await rag.initialize_storages()
    await rag.ainsert(open("docs.txt").read())

asyncio.run(main())

Each call reserves its worst-case cost before it's sent, so the concurrent extraction calls can't all spend the same remaining money, and a call that doesn't fit is refused with a 402 before it reaches OpenAI. LightRAG doesn't retry a 402 (only rate limits, timeouts and connection errors), so the job stops quickly instead of spending more. inferrail work index-2026-10-03 shows what the job cost.

Caveats: set max_tokens (without it each call reserves a 4,096-token default), leave OPENAI_API_BASE unset so embeddings keep going to OpenAI (the gateway only handles chat completions), and cached LLM responses don't count because they never make a call. The model needs a price in Inferrail; gpt-4o-mini has one built in.

Setup details: guide. I maintain Inferrail (open source, Apache-2.0), so take the suggestion with that in mind. Curious how others here keep indexing cost bounded on large corpora.