Per-run dollar budgets for Haystack pipelines and agents (no custom component)

[domondi1]

One Haystack pipeline often serves many runs, and an agent run can make several model calls. This shows a way to give each run its own dollar budget with the built-in OpenAIChatGenerator, enforced outside the pipeline.

Point the generator at a local OpenAI-compatible gateway (Inferrail) and pass the run's id and budget per pipeline.run:

llm = OpenAIChatGenerator(
    api_key=Secret.from_env_var("INFERRAIL_GATEWAY_TOKEN"),
    model="default",  # an Inferrail route name
    api_base_url="http://127.0.0.1:8000/v1",
    generation_kwargs={"max_tokens": 300},
)
# ...
pipeline.run({
    "llm": {"generation_kwargs": {"max_tokens": 300, "extra_headers": {
        "X-Inferrail-Attribute-Work-Id": "support-ticket-4812",
        "X-Inferrail-Budget-Usd": "0.05",
    }}},
    # ...other components' inputs
})

The first call of a run creates its budget. Calls that no longer fit get HTTP 402 before the provider (the generator raises openai.APIStatusError with status_code == 402; inside pipeline.run Haystack wraps it in PipelineRuntimeError), and inferrail work support-ticket-4812 prints the run's cost afterwards. Tested with Haystack 3.2.0 and Inferrail 0.4.8 (one pipeline, two runs; the second run's call was refused at its budget).

Limits: each call reserves an estimate (prompt plus max_tokens), so set max_tokens; only calls through the gateway are counted. Walkthrough: recipe. Disclosure: I maintain Inferrail (Apache-2.0).