feat(iorails): Inference-time http headers

docs/configure-rails/yaml-schema/model-configuration.mdx

@@ -280,3 +280,75 @@ models:
 ```
 
 Common parameters vary by provider. For built-in engines, see the OpenAI-compatible client options. For LangChain engines, refer to the corresponding LangChain provider documentation.
+
+### Custom HTTP Headers
+
+Set `parameters.default_headers` to attach custom HTTP headers to every request a model sends to its provider endpoint.
+This is useful for multi-tenant routing, request attribution, or correlation IDs.
+Headers are configured per model: each model sends only the headers declared in its own `parameters` block, so a header you want on several models must be repeated under each one.
+
+```yaml
+models:
+  - type: main
+    engine: nim
+    model: meta/llama-3.3-70b-instruct
+    parameters:
+      default_headers:
+        X-Tenant-Id: acme-corp
+        X-Routing-Pool: main-llm
+        X-Correlation-Source: guardrails-app
+
+  - type: content_safety
+    engine: nim
+    model: nvidia/llama-3.1-nemoguard-8b-content-safety
+    parameters:
+      default_headers:
+        X-Tenant-Id: acme-corp
+        X-Routing-Pool: safety-pool
+        X-Team: trust-and-safety
+```
+
+The configured headers apply on top of the request's base headers, which are `Content-Type` and the `Authorization` bearer token derived from the API key.
+Header names are matched case-insensitively, and a configured header overrides a base header of the same name.
+For example, setting `Authorization` under `default_headers` replaces the token built from the API key, which lets you use a custom authentication scheme.
+Configured headers are sent only as HTTP headers and never appear in the request body.
+
+Header values must be strings; a bare numeric or boolean value in YAML, such as `X-Retry: 3`, is converted to its string form.
+Keep secrets such as API keys in environment variables through `api_key_env_var` rather than hardcoding them in `default_headers` in a checked-in `config.yml`.
+
+Configured `default_headers` behave the same way on both the IORails and LLMRails engines.
+
+### Per-Request HTTP Headers
+
+Headers declared in `config.yml` are fixed when the configuration loads, so they cannot carry anything that varies from one request to the next.
+For values such as an end-user's tenant token, a per-call billing tag, or a trace ID from the surrounding application, pass an `http_headers` argument at inference time.
+
+```python
+response = await guardrails.generate_async(
+    messages=[{"role": "user", "content": "Hello"}],
+    http_headers={
+        "X-Tenant-Id": "acme-corp",
+        "X-Correlation-Id": request_id,
+    },
+)
+```
+
+The argument is accepted by `generate()`, `generate_async()`, and `stream_async()`.
+
+Unlike `default_headers`, which apply only to the model they are declared under, `http_headers` apply to *every* provider request the call triggers: the main LLM call and each rail's model or API call.
+This is what makes them suitable for correlation IDs and tenant attribution, where you want the same value on every outbound request in a single logical operation.
+
+Header names are matched case-insensitively across three layers, with the more specific layer winning:
+
+```
+Content-Type, Authorization  <  parameters.default_headers  <  http_headers
+```
+
+So a name set in both `default_headers` and `http_headers` is sent with the per-request value, while names set in only one layer are sent unchanged.
+As with `default_headers`, values are converted to strings and are sent only as HTTP headers, never in the request body.
+
+<Note>
+`http_headers` is supported by the IORails engine only.
+Passing it when your configuration routes to LLMRails raises `NotImplementedError` rather than dropping the headers, so a tenant or authorization header cannot go missing unnoticed.
+Use `parameters.default_headers` for values that both engines must send.
+</Note>

nemoguardrails/guardrails/_http.py

@@ -15,6 +15,8 @@
 
 """Shared aiohttp helpers for IORails engine HTTP clients."""
 
+from typing import Mapping, Optional
+
 import aiohttp
 
 DEFAULT_MAX_ATTEMPTS = 3
@@ -23,6 +25,22 @@
 RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})
 
 
+def merge_headers_case_insensitive(base: Mapping[str, str], overrides: Optional[Mapping[str, str]]) -> dict[str, str]:
+    """Merge ``overrides`` onto a copy of ``base``, matching header names case-insensitively.
+
+    HTTP header names are case-insensitive, so an override replaces any base
+    header with the same name regardless of case, adopting the override's
+    casing (mirrors ``BaseClient._build_headers``). ``base`` is not mutated;
+    ``None`` or empty ``overrides`` yields a plain copy of ``base``.
+    """
+    merged = dict(base)
+    for name, value in (overrides or {}).items():
+        for existing in [key for key in merged if key.lower() == name.lower()]:
+            del merged[existing]
+        merged[name] = value
+    return merged
+
+
 async def safe_read_body(response: aiohttp.ClientResponse, max_chars: int = 500) -> str:
     """Read response body for error messages, truncating if too large."""
     try:

nemoguardrails/guardrails/api_engine.py

@@ -19,6 +19,7 @@
 
 import logging
 import time
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Optional, cast
 
 import aiohttp
@@ -28,6 +29,7 @@
     DEFAULT_MAX_ATTEMPTS,
     DEFAULT_TIMEOUT_CONNECT,
     DEFAULT_TIMEOUT_TOTAL,
+    merge_headers_case_insensitive,
     safe_read_body,
 )
 from nemoguardrails.guardrails.base_engine import BaseEngine
@@ -91,7 +93,7 @@ def from_jailbreak_config(cls, jailbreak_config: JailbreakDetectionConfig) -> AP
             api_key=jailbreak_config.get_api_key(),
         )
 
-    async def call(self, body: dict[str, Any], **kwargs) -> dict:
+    async def call(self, body: dict[str, Any], *, http_headers: Optional[Mapping[str, str]] = None, **kwargs) -> dict:
         """POST the JSON body to the configured endpoint and return the parsed response."""
         if not self._running:
             raise APIEngineError("APIEngine has not been started. Call start() first.", endpoint=self.url)
@@ -105,6 +107,7 @@ async def call(self, body: dict[str, Any], **kwargs) -> dict:
         }
         if self.api_key:
             headers["Authorization"] = f"Bearer {self.api_key}"
+        headers = merge_headers_case_insensitive(headers, http_headers)
 
         req_id = get_request_id()
         log.info("[%s] HTTP POST %s", req_id, url)

nemoguardrails/guardrails/engine_registry.py

@@ -27,7 +27,7 @@
 
 from nemoguardrails.guardrails.api_engine import APIEngine
 from nemoguardrails.guardrails.base_engine import BaseEngine
-from nemoguardrails.guardrails.guardrails_types import get_request_id, truncate
+from nemoguardrails.guardrails.guardrails_types import get_http_headers, get_request_id, truncate
 from nemoguardrails.guardrails.model_engine import ModelEngine
 from nemoguardrails.guardrails.telemetry import (
     api_call_span,
@@ -188,6 +188,11 @@ async def model_call(self, model_type: str, messages: list[dict], **kwargs: Any)
         (one observation each for ``input`` and ``output`` token types,
         only when ``LLMResponse.usage`` is populated).
 
+        The request's inference-time HTTP headers are read here from the
+        request-scoped ContextVar and passed to the engine as an explicit
+        ``http_headers`` argument — deliberately not folded into
+        ``merged_params``, which becomes the request body.
+
         Raises:
             KeyError: If no engine is registered with the given name.
             TypeError: If the named engine is not a ModelEngine.
@@ -219,7 +224,7 @@ async def model_call(self, model_type: str, messages: list[dict], **kwargs: Any)
             # they land on the span even if the call raises.
             set_llm_request_attributes(span, merged_params)
             with duration_ctx:
-                result = await engine.chat_completion(messages, **merged_params)
+                result = await engine.chat_completion(messages, http_headers=get_http_headers(), **merged_params)
             # Set response/usage and content attrs inside the span context so
             # the helpers see the live LLM CLIENT span and the attributes land
             # before it closes.  Both are skipped on exception, which never
@@ -320,7 +325,9 @@ async def stream_model_call(
                 # — it's never read in that branch.
                 t0 = time.monotonic() if self._metrics_enabled else 0.0
                 last_chunk_time: Optional[float] = None
-                async for chunk in engine.stream_chat_completion(messages, **merged_params):
+                async for chunk in engine.stream_chat_completion(
+                    messages, http_headers=get_http_headers(), **merged_params
+                ):
                     if self._metrics_enabled:
                         # Per OTEL semconv, "first chunk" / "output chunk"
                         # mean content-bearing chunks — gate on
@@ -420,6 +427,10 @@ def extract_tool_exchanges(self, model_type: str, messages: list[dict]) -> list[
     async def api_call(self, api_name: str, message: dict[str, Any], **kwargs: Any) -> dict[str, Any]:
         """Route an API request to the named API engine.
 
+        Like ``model_call``, the request's inference-time HTTP headers are read
+        from the request-scoped ContextVar and passed as an explicit
+        ``http_headers`` argument rather than merged into the request body.
+
         Raises:
             KeyError: If no engine is registered with the given name.
             TypeError: If the named engine is not an APIEngine.
@@ -429,7 +440,7 @@ async def api_call(self, api_name: str, message: dict[str, Any], **kwargs: Any)
 
         with api_call_span(self._tracer, api_name):
             api_engine = self._get_engine(api_name, APIEngine)
-            response = await api_engine.call(message, **kwargs)
+            response = await api_engine.call(message, http_headers=get_http_headers(), **kwargs)
 
         log.debug("[%s] API engine '%s' response: %s", req_id, api_name, truncate(response))
         return response

nemoguardrails/guardrails/guardrails.py

@@ -23,6 +23,7 @@
 
 import logging
 import warnings
+from collections.abc import Mapping
 from typing import Any, AsyncIterator, Callable, List, Optional, Tuple, Type, Union, cast
 
 from typing_extensions import Self
@@ -191,25 +192,41 @@ def generate(
         prompt: str | None = None,
         messages: LLMMessages | None = None,
         options: Optional[Union[dict, GenerationOptions]] = None,
+        http_headers: Optional[Mapping[str, Any]] = None,
         **kwargs,
     ) -> Union[str, dict, GenerationResponse, Tuple[dict, dict]]:
         """Generate an LLM response synchronously with guardrails applied.
         Supported in both IORails and LLMRails
         """
+        if isinstance(self.rails_engine, IORails):
+            return self.rails_engine.generate(
+                prompt=prompt, messages=messages, options=options, http_headers=http_headers, **kwargs
+            )
+
+        if http_headers is not None:
+            raise NotImplementedError("LLMRails doesn't support inference-time HTTP headers in generate()")
         return self.rails_engine.generate(prompt=prompt, messages=messages, options=options, **kwargs)
 
     async def generate_async(
         self,
         prompt: str | None = None,
         messages: LLMMessages | None = None,
         options: Optional[Union[dict, GenerationOptions]] = None,
+        http_headers: Optional[Mapping[str, Any]] = None,
         **kwargs,
     ) -> str | dict | GenerationResponse | tuple[dict, dict]:
         """Generate an LLM response asynchronously with guardrails applied.
         Supported by both LLMRails and IORails
         """
         await self._ensure_started()
 
+        if isinstance(self.rails_engine, IORails):
+            return await self.rails_engine.generate_async(
+                prompt=prompt, messages=messages, options=options, http_headers=http_headers, **kwargs
+            )
+
+        if http_headers is not None:
+            raise NotImplementedError("LLMRails doesn't support inference-time HTTP headers in generate_async()")
         return await self.rails_engine.generate_async(prompt=prompt, messages=messages, options=options, **kwargs)
 
     def stream_async(
@@ -226,18 +243,22 @@ async def _with_startup(iterator: AsyncIterator[str | dict]) -> AsyncIterator[st
                 yield chunk
 
         if isinstance(self.rails_engine, IORails):
-            # IORails.stream_async() only accepts messages, options, include_metadata
-            unsupported = set(kwargs) - {"options", "include_metadata"}
+            # IORails.stream_async() only accepts messages, options, include_metadata, http_headers
+            unsupported = set(kwargs) - {"options", "include_metadata", "http_headers"}
             if unsupported:
                 log.warning("IORails stream_async: ignoring unsupported kwargs: %s", unsupported)
             return _with_startup(
                 self.rails_engine.stream_async(
                     messages=stream_messages,
                     options=kwargs.get("options"),
                     include_metadata=kwargs.get("include_metadata", False),
+                    http_headers=kwargs.get("http_headers"),
                 )
             )
 
+        if kwargs.pop("http_headers", None) is not None:
+            raise NotImplementedError("LLMRails doesn't support inference-time HTTP headers in stream_async()")
+
         llmrails = cast(LLMRails, self.rails_engine)
         return _with_startup(llmrails.stream_async(messages=stream_messages, **kwargs))
 

nemoguardrails/guardrails/guardrails_types.py

@@ -15,6 +15,8 @@
 
 
 import secrets
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
 from contextvars import ContextVar, Token
 from dataclasses import dataclass, field
 from enum import Enum
@@ -136,6 +138,42 @@ def reset_request_id(token: Token[str]) -> None:
     _request_id_var.reset(token)
 
 
+_http_headers_var: ContextVar[Optional[dict[str, str]]] = ContextVar("http_headers", default=None)
+
+
+def get_http_headers() -> Optional[dict[str, str]]:
+    """Return the inference-time HTTP headers bound to the current request, or None."""
+    return _http_headers_var.get()
+
+
+@contextmanager
+def request_http_headers(headers: Optional[Mapping[str, Any]]) -> Iterator[None]:
+    """Bind *headers* as the current request's inference-time HTTP headers.
+
+    Set once at the IORails request boundary and read once at the
+    ``EngineRegistry`` choke point, so every model and API call the request
+    makes carries the same headers without threading an argument through each
+    rail. Names and values are coerced to ``str`` so a non-string value behaves
+    the same as one configured under ``parameters.default_headers``.
+
+    ``ContextVar.reset()`` raises ``ValueError("... was created in a different
+    Context")`` when the enclosing async generator is closed from an outer
+    task's context. That one error is expected on streaming teardown and is
+    tolerated here, the same way the request-ID ContextVar is cleaned up; any
+    other ``ValueError`` indicates a bug and is re-raised.
+    """
+    coerced = None if headers is None else {str(name): str(value) for name, value in headers.items()}
+    token = _http_headers_var.set(coerced)
+    try:
+        yield
+    finally:
+        try:
+            _http_headers_var.reset(token)
+        except ValueError as exc:
+            if "different Context" not in str(exc):
+                raise
+
+
 def truncate(text: object, max_len: int | None = None) -> str:
     """Return ``str(text)`` truncated to *max_len* characters (default: LOG_CONTENT_TRUNCATE_LENGTH)."""
     s = str(text)