Building a haystack custom generator n4n.ai integration gives your pipelines a single OpenAI-compatible surface for 240+ models without rewriting components. This guide shows how to implement a custom Haystack 2.x generator, point it at the gateway, and verify token metering works end to end.
Step 1: Install dependencies
Haystack 2.x ships as haystack-ai. The OpenAI Python client is the easiest way to speak to any OpenAI-compatible HTTP surface, including the gateway.
pip install haystack-ai openai python-dotenvPin versions in production. Haystack moves fast; lock to a known-good release and test before bumping.
Step 2: Define the custom generator component
Haystack components are plain Python classes decorated with @component. A generator needs a run method that accepts a prompt and returns replies. Wrapping the OpenAI client gives you full control over headers, model naming, and error handling—something the built-in OpenAIGenerator does not expose cleanly when you need routing directives.
from haystack import component
from openai import OpenAI
import os
@component
class GatewayGenerator:
def __init__(
self,
model: str = "anthropic/claude-3-haiku",
temperature: float = 0.2,
max_tokens: int = 512,
):
self.client = OpenAI(
base_url=os.environ["GATEWAY_BASE_URL"],
api_key=os.environ["GATEWAY_API_KEY"],
# Forward routing or cache hints if your gateway expects them
default_headers={"x-gateway-route": "auto"},
)
self.model = model
self.temperature = temperature
self.max_tokens = max_tokens
@component.output_types(replies=list[str], meta=dict)
def run(self, prompt: str) -> dict:
resp = self.client.chat.completions.create(
model=self.model,
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
max_tokens=self.max_tokens,
)
reply = resp.choices[0].message.content or ""
return {
"replies": [reply],
"meta": {
"model": resp.model,
"usage": resp.usage.model_dump() if resp.usage else None,
},
}The meta output is not required by Haystack, but it is the only way you will see per-token usage without scraping logs. Keep it.
Step 3: Configure the gateway connection
n4n.ai exposes one OpenAI-compatible endpoint that fronts 240+ models and applies automatic fallback when a provider is rate-limited or degraded. Set GATEWAY_BASE_URL to that endpoint’s /v1 path and GATEWAY_API_KEY to your project token. Use a .env file locally; inject secrets via your orchestrator in CI.
export GATEWAY_BASE_URL="https://api.n4n.ai/v1"
export GATEWAY_API_KEY="sk-your-project-token"If you need to pin a provider or forward cache-control hints, the default_headers dict in the constructor is the right place. The gateway honors client routing directives and forwards provider cache-control hints, so a header like x-gateway-route or standard cache-control passes through to the upstream.
Model names follow the gateway’s namespace, not OpenAI’s. Expect strings like openai/gpt-4o-mini or meta/llama-3-70b. Hard-coding a model in a component is fine for a demo; read it from an env var or pipeline parameter in real code.
Step 4: Wire the generator into a pipeline
A generator alone is useful, but Haystack’s value is composition. Below is a minimal pipeline that takes a prompt, calls the custom generator, and prints the reply.
from haystack import Pipeline
from dotenv import load_dotenv
load_dotenv()
gen = GatewayGenerator(model="mistral/mixtral-8x7b-instruct")
pipe = Pipeline()
pipe.add_component("generator", gen)
result = pipe.run(
data={"generator": {"prompt": "Summarize LLM gateway fallback in one sentence."}}
)
print(result["generator"]["replies"][0])
print("usage:", result["generator"]["meta"]["usage"])For a RAG flow, drop a Retriever and PromptBuilder before the generator. The generator’s run signature stays identical; only the pipeline graph changes.
from haystack.components.builders import PromptBuilder
template = """
Context:
{% for doc in documents %}
{{ doc.content }}
{% endfor %}
Question: {{ question }}
Answer concisely.
"""
builder = PromptBuilder(template=template)
pipe = Pipeline()
pipe.add_component("builder", builder)
pipe.add_component("generator", GatewayGenerator())
pipe.connect("builder", "generator")Step 5: Verify success and inspect metering
Success means three things: the reply is non-empty, the model field reflects the gateway’s resolved backend, and the usage block contains prompt and completion tokens.
Run the script from Step 4. You should see a printed string and a dict similar to:
{
"prompt_tokens": 14,
"completion_tokens": 32,
"total_tokens": 46
}If total_tokens is null, the gateway did not return usage—check that you are calling a chat completion endpoint, not a legacy completion route. If the client raises AuthenticationError, your GATEWAY_API_KEY is wrong or missing. If it raises NotFoundError, the model string is not in the gateway’s catalog; list available models via the gateway’s /v1/models endpoint.
To confirm fallback actually triggers, force a degraded route in a staging environment or use a model that one upstream provider rate-limits. The gateway should return a response from a secondary provider without changing your code. That is the entire point of the haystack custom generator pattern: your pipeline talks to one stable interface.
Why not just use OpenAIGenerator
Haystack ships OpenAIGenerator and OpenAIChatGenerator. They accept api_base_url and work for simple cases. The moment you need custom headers, structured meta output, or retry logic that distinguishes provider errors from gateway errors, you are subclassing anyway. Writing the 30-line component above is less fragile than monkey-patching library internals.
A haystack custom generator also lets you swap the client entirely. If you later move to a non-OpenAI protocol, only the run method changes; the pipeline and the rest of your codebase stay put.
Operational notes
- Timeouts: Set
timeout=30on the OpenAI client. Gateways can hang when all providers are degraded. - Streaming: The OpenAI client supports
stream=True. If you add streaming, yield tokens inrunvia a generator and document the output type change. - Concurrency: The client is thread-safe. Create one instance per process; don’t rebuild the generator per request in a web server.
- Model allowlist: Restrict
modelto an internal enum. The gateway’s 240+ model catalog is great for exploration, dangerous for production if a typo silently routes to a costly backend.
The haystack custom generator approach is boring in the best way: a thin adapter, explicit I/O, and no hidden behavior. That is what you want between your pipeline and a multi-provider inference gateway.