From 3818fa64dda974299f2fd35aeb9d5f0c48725626 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:50:24 -0500 Subject: [PATCH 1/6] Add tests for config-defined default headers in config --- tests/guardrails/test__http.py | 39 +++++++++++++ tests/guardrails/test_iorails.py | 84 +++++++++++++++++++++++++++ tests/guardrails/test_model_engine.py | 79 +++++++++++++++++++++++++ 3 files changed, 202 insertions(+) diff --git a/tests/guardrails/test__http.py b/tests/guardrails/test__http.py index e0b623d852..8e6b14a62b 100644 --- a/tests/guardrails/test__http.py +++ b/tests/guardrails/test__http.py @@ -24,6 +24,7 @@ DEFAULT_TIMEOUT_CONNECT, DEFAULT_TIMEOUT_TOTAL, RETRYABLE_STATUS_CODES, + merge_headers_case_insensitive, safe_read_body, ) @@ -66,6 +67,44 @@ async def test_501_chars_truncated(self): assert len(result) == 500 +class TestMergeHeadersCaseInsensitive: + """Test case-insensitive header merging (override wins, base preserved).""" + + def test_overrides_added_to_base(self): + """Non-colliding override keys are added to the base headers.""" + base = {"Content-Type": "application/json"} + result = merge_headers_case_insensitive(base, {"X-Tenant": "acme"}) + assert result == {"Content-Type": "application/json", "X-Tenant": "acme"} + + def test_override_replaces_case_insensitive_match(self): + """An override replaces a base header that differs only by case, keeping one entry.""" + base = {"Content-Type": "application/json", "Authorization": "Bearer base"} + result = merge_headers_case_insensitive(base, {"authorization": "Bearer override"}) + auth_keys = [key for key in result if key.lower() == "authorization"] + assert auth_keys == ["authorization"] + assert result["authorization"] == "Bearer override" + + def test_none_overrides_returns_base_copy(self): + """None overrides yields an unchanged copy that does not alias the base.""" + base = {"Content-Type": "application/json"} + result = merge_headers_case_insensitive(base, None) + assert result == base + assert result is not base + + def test_empty_overrides_returns_base_copy(self): + """Empty overrides yields an unchanged copy that does not alias the base.""" + base = {"Content-Type": "application/json"} + result = merge_headers_case_insensitive(base, {}) + assert result == base + assert result is not base + + def test_base_not_mutated(self): + """Merging does not mutate the caller's base dict.""" + base = {"Authorization": "Bearer base"} + merge_headers_case_insensitive(base, {"authorization": "Bearer override"}) + assert base == {"Authorization": "Bearer base"} + + class TestSharedConstants: """Test values of shared HTTP constants.""" diff --git a/tests/guardrails/test_iorails.py b/tests/guardrails/test_iorails.py index 11503da8e6..59ce327ab7 100644 --- a/tests/guardrails/test_iorails.py +++ b/tests/guardrails/test_iorails.py @@ -68,6 +68,90 @@ def test_rails_manager_uses_engine_registry(self, iorails_sync): assert iorails_sync.rails_manager.engine_registry is iorails_sync.engine_registry +DEFAULT_HEADERS_CONFIG = { + "models": [ + { + "type": "main", + "engine": "nim", + "model": "meta/llama-3.3-70b-instruct", + "parameters": {"default_headers": {"X-Main-Route": "main-pool"}}, + }, + { + "type": "content_safety", + "engine": "nim", + "model": "nvidia/llama-3.1-nemoguard-8b-content-safety", + "parameters": {"default_headers": {"X-Safety-Route": "safety-pool"}}, + }, + ], +} + + +@pytest_asyncio.fixture +async def iorails_with_header_config(): + """Build an IORails whose main and content_safety models carry distinct + parameters.default_headers, for end-to-end per-model header assertions.""" + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + iorails = IORails(RailsConfig.from_content(config=DEFAULT_HEADERS_CONFIG)) + try: + yield iorails + finally: + await iorails.stop() + + +class TestConfigDefaultHeadersEndToEnd: + """End-to-end: IORails built from config routes each model's + parameters.default_headers onto that model's outbound request, and the + main LLM and a second LLM carry only their own headers.""" + + @staticmethod + def _mock_client(): + """Build a mock aiohttp client whose post() records its call args.""" + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={"choices": [{"message": {"content": "ok"}}]}) + + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + return mock_client + + async def _headers_for_model(self, iorails, model_type): + """Route a model_call for model_type through a mock client and return the sent headers.""" + engine = iorails.engine_registry._get_engine(model_type, ModelEngine) + engine._client = self._mock_client() + engine._running = True + await iorails.engine_registry.model_call(model_type, [{"role": "user", "content": "Hi"}]) + return engine._client.post.call_args[1]["headers"] + + @pytest.mark.asyncio + async def test_main_llm_carries_only_its_config_headers(self, iorails_with_header_config): + """The main model request carries its own default_header plus base headers, and not the second model's.""" + headers = await self._headers_for_model(iorails_with_header_config, "main") + assert headers["X-Main-Route"] == "main-pool" + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + assert "X-Safety-Route" not in headers + + @pytest.mark.asyncio + async def test_second_llm_carries_only_its_config_headers(self, iorails_with_header_config): + """The content_safety model request carries its own default_header plus base headers, and not the main model's.""" + headers = await self._headers_for_model(iorails_with_header_config, "content_safety") + assert headers["X-Safety-Route"] == "safety-pool" + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + assert "X-Main-Route" not in headers + + @pytest.mark.asyncio + async def test_main_and_second_llm_get_distinct_headers(self, iorails_with_header_config): + """The two models' outbound header sets differ, proving per-model routing.""" + main_headers = await self._headers_for_model(iorails_with_header_config, "main") + safety_headers = await self._headers_for_model(iorails_with_header_config, "content_safety") + assert main_headers != safety_headers + assert main_headers.get("X-Main-Route") == "main-pool" + assert safety_headers.get("X-Safety-Route") == "safety-pool" + + class TestGenerateAsync: """Test the generate_async input-check → LLM → output-check pipeline.""" diff --git a/tests/guardrails/test_model_engine.py b/tests/guardrails/test_model_engine.py index 389195b09f..7cf095f27c 100644 --- a/tests/guardrails/test_model_engine.py +++ b/tests/guardrails/test_model_engine.py @@ -686,6 +686,85 @@ async def test_call_raises_if_not_started(self): await engine.call([{"role": "user", "content": "Hi"}]) +class TestModelEngineDefaultHeaders: + """Test that config-level parameters.default_headers reach outbound requests.""" + + @staticmethod + def _mock_client(): + """Build a mock aiohttp client whose post() records call args.""" + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={"choices": [{"message": {"content": "ok"}}]}) + + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + return mock_client + + @staticmethod + def _headers_from(mock_client): + """Extract the headers dict passed to the mocked post().""" + return mock_client.post.call_args[1]["headers"] + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_config_default_headers_merged_into_request(self): + """A configured default_header is sent alongside Content-Type and Authorization.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"X-Tenant": "acme"}})) + engine._client = self._mock_client() + engine._running = True + + await engine.call([{"role": "user", "content": "Hi"}]) + + headers = self._headers_from(engine._client) + assert headers["X-Tenant"] == "acme" + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_config_default_header_overrides_authorization_case_insensitive(self): + """A configured 'authorization' header replaces the api_key-derived Authorization.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"authorization": "Bearer custom"}})) + engine._client = self._mock_client() + engine._running = True + + await engine.call([{"role": "user", "content": "Hi"}]) + + headers = self._headers_from(engine._client) + auth_keys = [key for key in headers if key.lower() == "authorization"] + assert auth_keys == ["authorization"] + assert headers["authorization"] == "Bearer custom" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_config_default_headers_absent_from_body(self): + """default_headers configure transport, never the JSON request body.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"X-Tenant": "acme"}})) + engine._client = self._mock_client() + engine._running = True + + await engine.call([{"role": "user", "content": "Hi"}]) + + body = engine._client.post.call_args[1]["json"] + assert "default_headers" not in body + assert "X-Tenant" not in body + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_no_config_default_headers_leaves_base_headers(self): + """Without default_headers the request carries only the base headers.""" + engine = ModelEngine(_make_model()) + engine._client = self._mock_client() + engine._running = True + + await engine.call([{"role": "user", "content": "Hi"}]) + + headers = self._headers_from(engine._client) + assert headers == {"Content-Type": "application/json", "Authorization": "Bearer test-key"} + + class TestModelEngineStreamCall: """Test ModelEngine.stream_call() SSE streaming.""" From 7df961a8d5a6a67b9fdf880d3f3c1317c30816a4 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:05:44 -0500 Subject: [PATCH 2/6] Propagate config headers through to model calls --- nemoguardrails/guardrails/_http.py | 19 +++++++++++++++++++ nemoguardrails/guardrails/model_engine.py | 16 +++++++++++----- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/guardrails/_http.py b/nemoguardrails/guardrails/_http.py index 254a14327f..8eead19274 100644 --- a/nemoguardrails/guardrails/_http.py +++ b/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,23 @@ 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(): + existing = next((key for key in merged if key.lower() == name.lower()), None) + if existing is not None: + 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: diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index 2716960b70..e3c7722bc3 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -35,6 +35,7 @@ DEFAULT_MAX_ATTEMPTS, DEFAULT_TIMEOUT_CONNECT, DEFAULT_TIMEOUT_TOTAL, + merge_headers_case_insensitive, safe_read_body, ) from nemoguardrails.guardrails.base_engine import BaseEngine @@ -87,11 +88,11 @@ # be used in streaming or non-streaming mode. `stream_call` sets the # streaming default (`include_usage`), overridable via `llm_params`. "stream_options", - # client-only options — these configure the OpenAI-compatible client - # (constructor kwargs), not the chat-completion request body. IORails - # doesn't wire the shared client yet; reserve them so they're never - # forwarded as body fields, leaving proper client support to a future - # refactor. + # client-only options — these configure transport, not the + # chat-completion request body, so they are never forwarded as body + # fields. `default_headers` is read into `self.default_headers` and + # applied as request headers by `_prepare_request`; `default_query` + # is reserved for future client wiring. "default_headers", "default_query", } @@ -477,6 +478,10 @@ def __init__(self, model_config: Model) -> None: max_attempts=int(params.get("max_attempts") or DEFAULT_MAX_ATTEMPTS), ) + # Static per-model HTTP headers from `parameters.default_headers`, applied + # to every request by `_prepare_request` (LLMRails parity). + self.default_headers: dict[str, str] = dict(params.get("default_headers") or {}) + # Default `llm_params` used on inference are the subset of Model.parameters after # filtering out keys in _RESERVED_LLM_PARAMETERS. Exposed as a read-only # MappingProxyType view so callers can't mutate the shared per-engine defaults. @@ -550,6 +555,7 @@ def _prepare_request(self, messages: LLMMessages, **kwargs: Any) -> _RequestPara headers: dict[str, str] = {"Content-Type": "application/json"} if self.api_key: headers["Authorization"] = f"Bearer {self.api_key}" + headers = merge_headers_case_insensitive(headers, self.default_headers) body: dict[str, Any] = {"model": self.model_name, "messages": messages, **kwargs} return _RequestParams(client=client, url=url, headers=headers, body=body) From 926dc1eee1024f765761cf68bc33026f425e8000 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:13:59 -0500 Subject: [PATCH 3/6] Add end-to-end test to check headers are used correctly --- nemoguardrails/guardrails/model_engine.py | 4 +- tests/guardrails/test_iorails.py | 50 +++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index e3c7722bc3..3157e363a7 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -90,9 +90,7 @@ "stream_options", # client-only options — these configure transport, not the # chat-completion request body, so they are never forwarded as body - # fields. `default_headers` is read into `self.default_headers` and - # applied as request headers by `_prepare_request`; `default_query` - # is reserved for future client wiring. + # fields. "default_headers", "default_query", } diff --git a/tests/guardrails/test_iorails.py b/tests/guardrails/test_iorails.py index 59ce327ab7..1068cab5e7 100644 --- a/tests/guardrails/test_iorails.py +++ b/tests/guardrails/test_iorails.py @@ -20,6 +20,8 @@ import pytest import pytest_asyncio +from aiohttp import web +from aiohttp.test_utils import TestServer from nemoguardrails import Guardrails from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult @@ -152,6 +154,54 @@ async def test_main_and_second_llm_get_distinct_headers(self, iorails_with_heade assert safety_headers.get("X-Safety-Route") == "safety-pool" +class TestConfigDefaultHeadersOverHTTP: + """True end-to-end: run generate_async over a real loopback HTTP server and + assert on the headers the aiohttp client actually put on the wire.""" + + @pytest.mark.asyncio + async def test_config_default_headers_sent_on_the_wire(self): + """A configured default_header reaches the provider request, alongside the + api-key Authorization, and never leaks into the JSON body.""" + captured: dict = {} + + async def handler(request): + captured["headers"] = dict(request.headers) + captured["body"] = await request.json() + return web.json_response({"choices": [{"message": {"role": "assistant", "content": "ok"}}]}) + + app = web.Application() + app.router.add_post("/v1/chat/completions", handler) + server = TestServer(app) + await server.start_server() + try: + config = RailsConfig.from_content( + config={ + "models": [ + { + "type": "main", + "engine": "openai", + "model": "test-model", + "parameters": { + "base_url": str(server.make_url("/")), + "default_headers": {"X-Tenant": "acme"}, + }, + } + ] + } + ) + with patch.dict("os.environ", {"OPENAI_API_KEY": "test-key"}): + iorails = IORails(config) + async with iorails: + result = await iorails.generate_async(messages=[{"role": "user", "content": "hi"}]) + finally: + await server.close() + + assert result == {"role": "assistant", "content": "ok"} + assert captured["headers"]["X-Tenant"] == "acme" + assert captured["headers"]["Authorization"] == "Bearer test-key" + assert "X-Tenant" not in captured["body"] + + class TestGenerateAsync: """Test the generate_async input-check → LLM → output-check pipeline.""" From c1e04ba0584967b96114f258aba548976360fd68 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:53:56 -0500 Subject: [PATCH 4/6] Address review feedback --- nemoguardrails/guardrails/_http.py | 3 +-- nemoguardrails/guardrails/model_engine.py | 6 +++--- tests/guardrails/test__http.py | 8 ++++++++ tests/guardrails/test_model_engine.py | 17 +++++++++++++++++ 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/nemoguardrails/guardrails/_http.py b/nemoguardrails/guardrails/_http.py index 8eead19274..34d3ecb2c1 100644 --- a/nemoguardrails/guardrails/_http.py +++ b/nemoguardrails/guardrails/_http.py @@ -35,8 +35,7 @@ def merge_headers_case_insensitive(base: Mapping[str, str], overrides: Optional[ """ merged = dict(base) for name, value in (overrides or {}).items(): - existing = next((key for key in merged if key.lower() == name.lower()), None) - if existing is not None: + for existing in [key for key in merged if key.lower() == name.lower()]: del merged[existing] merged[name] = value return merged diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index 3157e363a7..4a76d74352 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -476,9 +476,9 @@ def __init__(self, model_config: Model) -> None: max_attempts=int(params.get("max_attempts") or DEFAULT_MAX_ATTEMPTS), ) - # Static per-model HTTP headers from `parameters.default_headers`, applied - # to every request by `_prepare_request` (LLMRails parity). - self.default_headers: dict[str, str] = dict(params.get("default_headers") or {}) + self.default_headers: Mapping[str, str] = MappingProxyType( + {str(key): str(value) for key, value in (params.get("default_headers") or {}).items()} + ) # Default `llm_params` used on inference are the subset of Model.parameters after # filtering out keys in _RESERVED_LLM_PARAMETERS. Exposed as a read-only diff --git a/tests/guardrails/test__http.py b/tests/guardrails/test__http.py index 8e6b14a62b..b94064a85d 100644 --- a/tests/guardrails/test__http.py +++ b/tests/guardrails/test__http.py @@ -104,6 +104,14 @@ def test_base_not_mutated(self): merge_headers_case_insensitive(base, {"authorization": "Bearer override"}) assert base == {"Authorization": "Bearer base"} + def test_removes_all_case_equivalent_base_keys(self): + """An override collapses every case variant of a name in the base into a single header.""" + base = {"Authorization": "base-a", "authorization": "base-b"} + result = merge_headers_case_insensitive(base, {"Authorization": "override"}) + auth_keys = [key for key in result if key.lower() == "authorization"] + assert auth_keys == ["Authorization"] + assert result["Authorization"] == "override" + class TestSharedConstants: """Test values of shared HTTP constants.""" diff --git a/tests/guardrails/test_model_engine.py b/tests/guardrails/test_model_engine.py index 7cf095f27c..1ad8b8bbd1 100644 --- a/tests/guardrails/test_model_engine.py +++ b/tests/guardrails/test_model_engine.py @@ -17,6 +17,7 @@ import asyncio import json +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import aiohttp @@ -764,6 +765,22 @@ async def test_no_config_default_headers_leaves_base_headers(self): headers = self._headers_from(engine._client) assert headers == {"Content-Type": "application/json", "Authorization": "Bearer test-key"} + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + def test_config_default_header_values_coerced_to_str(self): + """Non-string YAML header values (int, bool) are stored as strings.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"X-Count": 3, "X-Flag": True}})) + assert engine.default_headers["X-Count"] == "3" + assert engine.default_headers["X-Flag"] == "True" + assert all(isinstance(value, str) for value in engine.default_headers.values()) + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + def test_default_headers_mapping_is_immutable(self): + """default_headers is a read-only mapping so callers cannot mutate shared per-engine state.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"X-Tenant": "acme"}})) + headers: Any = engine.default_headers + with pytest.raises(TypeError): + headers["X-Injected"] = "nope" + class TestModelEngineStreamCall: """Test ModelEngine.stream_call() SSE streaming.""" From 979370dc15ea2588134b5de2324100611e8f3a67 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:58:38 -0500 Subject: [PATCH 5/6] Document default headers --- .../yaml-schema/model-configuration.mdx | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/docs/configure-rails/yaml-schema/model-configuration.mdx b/docs/configure-rails/yaml-schema/model-configuration.mdx index 1547108313..ec6306fed3 100644 --- a/docs/configure-rails/yaml-schema/model-configuration.mdx +++ b/docs/configure-rails/yaml-schema/model-configuration.mdx @@ -280,3 +280,40 @@ 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`. + +This behavior is the same for both the IORails and LLMRails engines. From fd22f2e5f2384dc4b77a7edde2d0ffb28fa16b35 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:49:23 -0500 Subject: [PATCH 6/6] Initial checkin of inference-time http headers --- .../yaml-schema/model-configuration.mdx | 37 +++- nemoguardrails/guardrails/api_engine.py | 5 +- nemoguardrails/guardrails/engine_registry.py | 19 +- nemoguardrails/guardrails/guardrails.py | 25 ++- nemoguardrails/guardrails/guardrails_types.py | 38 ++++ nemoguardrails/guardrails/iorails.py | 53 ++++- nemoguardrails/guardrails/model_engine.py | 59 +++++- tests/guardrails/test_api_engine.py | 67 +++++++ tests/guardrails/test_engine_registry.py | 101 +++++++++- tests/guardrails/test_guardrails.py | 148 +++++++++++++- tests/guardrails/test_guardrails_types.py | 87 ++++++++- tests/guardrails/test_iorails.py | 181 +++++++++++++++++- tests/guardrails/test_iorails_streaming.py | 95 ++++++++- tests/guardrails/test_model_engine.py | 162 ++++++++++++++++ 14 files changed, 1041 insertions(+), 36 deletions(-) diff --git a/docs/configure-rails/yaml-schema/model-configuration.mdx b/docs/configure-rails/yaml-schema/model-configuration.mdx index ec6306fed3..d18e4a3dcc 100644 --- a/docs/configure-rails/yaml-schema/model-configuration.mdx +++ b/docs/configure-rails/yaml-schema/model-configuration.mdx @@ -316,4 +316,39 @@ Configured headers are sent only as HTTP headers and never appear in the request 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`. -This behavior is the same for both the IORails and LLMRails engines. +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. + + +`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. + diff --git a/nemoguardrails/guardrails/api_engine.py b/nemoguardrails/guardrails/api_engine.py index 109851fea5..3aaac51bef 100644 --- a/nemoguardrails/guardrails/api_engine.py +++ b/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) diff --git a/nemoguardrails/guardrails/engine_registry.py b/nemoguardrails/guardrails/engine_registry.py index 1b7ff6f242..800c4fa1d5 100644 --- a/nemoguardrails/guardrails/engine_registry.py +++ b/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 diff --git a/nemoguardrails/guardrails/guardrails.py b/nemoguardrails/guardrails/guardrails.py index 6a438d8dc7..abaaf785cf 100644 --- a/nemoguardrails/guardrails/guardrails.py +++ b/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,11 +192,19 @@ 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( @@ -203,6 +212,7 @@ async def generate_async( 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. @@ -210,6 +220,13 @@ async def generate_async( """ 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,8 +243,8 @@ 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( @@ -235,9 +252,13 @@ async def _with_startup(iterator: AsyncIterator[str | dict]) -> AsyncIterator[st 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)) diff --git a/nemoguardrails/guardrails/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index e7ab281f59..edcc2e2f34 100644 --- a/nemoguardrails/guardrails/guardrails_types.py +++ b/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) diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index ac85b7877f..f3cc51b1e4 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -25,9 +25,9 @@ import logging import time import warnings -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Mapping from contextlib import nullcontext, suppress -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from nemoguardrails.actions.llm.utils import _extract_and_remove_think_tags from nemoguardrails.base_guardrails import BaseGuardrails @@ -41,6 +41,7 @@ RailDirection, TimedLLMResponse, get_request_id, + request_http_headers, serialize_prompt, truncate, ) @@ -743,6 +744,7 @@ def generate( prompt: Optional[str] = None, messages: Optional[LLMMessages] = None, options: Optional[Union[dict, GenerationOptions]] = None, + http_headers: Optional[Mapping[str, Any]] = None, **kwargs, ) -> Union[LLMMessage, GenerationResponse]: """Synchronous version of generate_async. @@ -751,6 +753,13 @@ def generate( the ``generate()`` call. For production use, use the asynchronous `generate_async()` and `stream_async()` methods for non-streaming and streaming requests respectively. + + Args: + prompt: Single user turn; ignored when ``messages`` is given. + messages: Conversation messages in OpenAI format. + options: GenerationOptions; when set, a GenerationResponse is returned. + http_headers: Per-request HTTP headers for every model and API call. + **kwargs: Forwarded to the pipeline; ``state`` is rejected. """ messages = self._convert_to_messages(prompt, messages) @@ -765,7 +774,9 @@ async def _run_sync_iorails(): """Spin up a short-lived IORails engine for one synchronous generate call.""" # Avoid counting this sync-API bridge as a separate user-created IORails instance. async with IORails(sync_config, _report_usage=False) as iorails_engine: - return await iorails_engine.generate_async(messages=messages, options=options, **kwargs) + return await iorails_engine.generate_async( + messages=messages, options=options, http_headers=http_headers, **kwargs + ) return asyncio.run(_run_sync_iorails()) @@ -774,6 +785,7 @@ async def generate_async( prompt: Optional[str] = None, messages: Optional[LLMMessages] = None, options: Optional[Union[dict, GenerationOptions]] = None, + http_headers: Optional[Mapping[str, Any]] = None, **kwargs, ) -> Union[LLMMessage, GenerationResponse]: """Public entry: submit the request to the internal work queue. @@ -790,13 +802,22 @@ async def generate_async( (OTEL HTTP semconv). A ``QueueFull`` rejection shows up in BOTH ``requests.errors{error.type=QueueFull}`` and ``nonstream.rejections`` — honest dual-signal reporting. + + Args: + prompt: Single user turn; ignored when ``messages`` is given. + messages: Conversation messages in OpenAI format. + options: GenerationOptions; when set, a GenerationResponse is returned. + http_headers: Per-request HTTP headers for every model and API call. + **kwargs: Forwarded to the pipeline; ``state`` is rejected. """ messages = self._convert_to_messages(prompt, messages) await self.start() metrics_ctx = request_metrics() if self._metrics_enabled else nullcontext() with metrics_ctx: try: - return await self._generate_async_queue.submit(self._run_generate, messages, options=options, **kwargs) + return await self._generate_async_queue.submit( + self._run_generate, messages, options=options, http_headers=http_headers, **kwargs + ) except asyncio.QueueFull: if self._metrics_enabled: record_nonstream_rejected() @@ -806,6 +827,7 @@ async def _run_generate( self, messages: LLMMessages, options: Optional[Union[dict, GenerationOptions]] = None, + http_headers: Optional[Mapping[str, Any]] = None, **kwargs, ) -> Union[LLMMessage, GenerationResponse]: """Runs inside a queue worker task. Wraps the pipeline in @@ -815,7 +837,7 @@ async def _run_generate( lifecycle scope by ``generate_async``, not here. """ tracer = self._tracer if self._tracing_enabled else None - with traced_request(tracer) as (request_span, req_id): + with request_http_headers(http_headers), traced_request(tracer) as (request_span, req_id): t0 = time.monotonic() try: result = await self._do_generate(messages, req_id, request_span, options=options, **kwargs) @@ -1137,6 +1159,10 @@ def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = No Mirrors ``generate``: spins up a short-lived IORails engine with tracing and metrics disabled and runs the check on it. For production use, prefer the asynchronous ``check_async``. + + Args: + messages: Conversation messages in OpenAI format. + rail_types: Rail types to run; None auto-detects from the message roles. """ if check_sync_call_from_async_loop(): raise RuntimeError( @@ -1167,6 +1193,10 @@ async def check_async(self, messages: LLMMessages, rail_types: Optional[list[Rai Submitted through the same admission queue as ``generate_async`` so the check path shares non-streaming concurrency limits, request metrics, and the per-request trace span. + + Args: + messages: Conversation messages in OpenAI format. + rail_types: Rail types to run; None auto-detects from the message roles. """ await self.start() metrics_ctx = request_metrics() if self._metrics_enabled else nullcontext() @@ -1275,6 +1305,7 @@ def stream_async( messages: LLMMessages, options: Optional[Union[dict, GenerationOptions]] = None, include_metadata: Optional[bool] = False, + http_headers: Optional[Mapping[str, Any]] = None, ) -> AsyncIterator[Union[str, dict]]: """Stream LLM response tokens with input/output rails applied. @@ -1286,10 +1317,9 @@ def stream_async( Args: messages: Conversation messages in OpenAI format. - options: Optional GenerationOptions (llm_params are forwarded to - the main LLM call). - include_metadata: When True, chunks are dicts with ``text`` and - ``metadata`` keys instead of plain strings. + options: GenerationOptions; ``llm_params`` are forwarded to the main LLM call. + include_metadata: When True, chunks are ``text``/``metadata`` dicts, not strings. + http_headers: Per-request HTTP headers for every model and API call. Returns: An async iterator of string chunks (or dicts). @@ -1482,7 +1512,10 @@ async def _wrapped_iterator(): # request span is the current OTEL context when create_task() # below snapshots contextvars — that's what makes rail / LLM # spans raised inside _generation_task attach as children. - with traced_request(tracer) as (request_span, req_id): + # request_http_headers is entered in the same scope for the + # same reason: _generation_task inherits the headers from the + # context snapshot taken at create_task() time. + with request_http_headers(http_headers), traced_request(tracer) as (request_span, req_id): t0 = time.monotonic() # Accumulate chunks the consumer actually receives. # Declared outside the try so the outer finally can diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index 4a76d74352..7a75f2b972 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -93,6 +93,11 @@ # fields. "default_headers", "default_query", + # inference-time header argument name — reserved so a model configured + # with `parameters.http_headers` (the per-request argument mistaken for + # the config key `default_headers`) can't collide with the explicit + # http_headers keyword the engine registry passes into chat_completion(). + "http_headers", } ) @@ -545,8 +550,21 @@ def _ensure_running(self) -> None: model_name=self.model_name, ) - def _prepare_request(self, messages: LLMMessages, **kwargs: Any) -> _RequestParams: - """Build the client, URL, headers, and body common to every request.""" + def _prepare_request( + self, + messages: LLMMessages, + *, + http_headers: Optional[Mapping[str, str]] = None, + **kwargs: Any, + ) -> _RequestParams: + """Build the client, URL, headers, and body common to every request. + + Headers are layered so the more specific source wins, matched + case-insensitively: the derived ``Content-Type``/``Authorization`` base, + then the model's configured ``default_headers``, then the per-request + *http_headers*. ``http_headers`` is a transport argument and never + reaches ``body``. + """ client = cast(RetryClient, self._client) url = self.base_url + _CHAT_COMPLETIONS_ENDPOINT @@ -554,6 +572,7 @@ def _prepare_request(self, messages: LLMMessages, **kwargs: Any) -> _RequestPara if self.api_key: headers["Authorization"] = f"Bearer {self.api_key}" headers = merge_headers_case_insensitive(headers, self.default_headers) + headers = merge_headers_case_insensitive(headers, http_headers) body: dict[str, Any] = {"model": self.model_name, "messages": messages, **kwargs} return _RequestParams(client=client, url=url, headers=headers, body=body) @@ -588,6 +607,8 @@ def _wrap_exception(self, exc: Exception, req_id: str, t0: float, label: str = " async def call( self, messages: LLMMessages, + *, + http_headers: Optional[Mapping[str, str]] = None, **kwargs: Any, ) -> dict: """Make a POST request to the /v1/chat/completions endpoint. @@ -597,6 +618,9 @@ async def call( Args: messages: List of message dicts in OpenAI format. + http_headers: Per-request HTTP headers layered over the model's + configured ``default_headers``; sent as headers only, never + added to the request body. **kwargs: Additional parameters for the request body (temperature, max_tokens, etc.) Returns: @@ -606,7 +630,7 @@ async def call( ModelEngineError: If the request fails after all retries. """ self._ensure_running() - req = self._prepare_request(messages, **kwargs) + req = self._prepare_request(messages, http_headers=http_headers, **kwargs) req_id = get_request_id() log.info("[%s] HTTP POST %s model='%s'", req_id, req.url, self.model_name) @@ -636,6 +660,8 @@ async def call( async def stream_call( self, messages: LLMMessages, + *, + http_headers: Optional[Mapping[str, str]] = None, **kwargs: Any, ) -> AsyncGenerator[LLMResponseChunk, None]: """Make a streaming POST request to the /v1/chat/completions endpoint. @@ -667,6 +693,9 @@ async def stream_call( Args: messages: List of message dicts in OpenAI format. + http_headers: Per-request HTTP headers layered over the model's + configured ``default_headers``; sent as headers only, never + added to the request body. **kwargs: Additional parameters for the request body (temperature, max_tokens, etc.) Yields: @@ -680,7 +709,7 @@ async def stream_call( # Request usage on the terminal stream chunk (LLMRails parity); a caller can override # or disable it via stream_options in llm_params. kwargs.setdefault("stream_options", {"include_usage": True}) - req = self._prepare_request(messages, stream=True, **kwargs) + req = self._prepare_request(messages, stream=True, http_headers=http_headers, **kwargs) # For streaming, disable the total timeout (response body streams # for the full generation duration) and use sock_read to detect stalls @@ -797,7 +826,13 @@ async def stream_call( except Exception as exc: raise self._wrap_exception(exc, req_id, t0, label="Stream request") from exc - async def chat_completion(self, messages: LLMMessages, **kwargs: Any) -> LLMResponse: + async def chat_completion( + self, + messages: LLMMessages, + *, + http_headers: Optional[Mapping[str, str]] = None, + **kwargs: Any, + ) -> LLMResponse: """Generate a chat completion and return a structured ``LLMResponse``. Calls the /v1/chat/completions endpoint and parses the OpenAI-format @@ -805,10 +840,14 @@ async def chat_completion(self, messages: LLMMessages, **kwargs: Any) -> LLMResp provider exposes ``reasoning_content``), usage, finish reason, and request id. + ``http_headers`` carries the request's inference-time headers; see + ``_prepare_request`` for how they layer over the configured + ``default_headers``. + Raises: ModelEngineError: If the request fails or the response format is unexpected. """ - response = await self.call(messages, **kwargs) + response = await self.call(messages, http_headers=http_headers, **kwargs) try: return _parse_chat_completion(response) except ValueError as exc: @@ -818,7 +857,11 @@ async def chat_completion(self, messages: LLMMessages, **kwargs: Any) -> LLMResp ) from exc async def stream_chat_completion( - self, messages: LLMMessages, **kwargs: Any + self, + messages: LLMMessages, + *, + http_headers: Optional[Mapping[str, str]] = None, + **kwargs: Any, ) -> AsyncGenerator[LLMResponseChunk, None]: """Stream a chat completion and yield ``LLMResponseChunk`` objects. @@ -829,7 +872,7 @@ async def stream_chat_completion( Raises: ModelEngineError: If the request fails after all retries. """ - async for chunk in self.stream_call(messages, **kwargs): + async for chunk in self.stream_call(messages, http_headers=http_headers, **kwargs): yield chunk def parse_tools(self, llm_params: Optional[dict]) -> Toolset: diff --git a/tests/guardrails/test_api_engine.py b/tests/guardrails/test_api_engine.py index 882c1834c1..0c4cc1c845 100644 --- a/tests/guardrails/test_api_engine.py +++ b/tests/guardrails/test_api_engine.py @@ -232,6 +232,21 @@ async def test_context_manager_calls_start_and_stop(self): class TestAPIEngineCall: """Test APIEngine.call() HTTP request construction and error handling.""" + @staticmethod + def _start_with_recording_client(engine): + """Start *engine* against a mock client and return the client, whose post() records call args.""" + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={"jailbreak": False, "score": -0.87}) + + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + engine._client = mock_client + engine._running = True + return mock_client + @pytest.mark.asyncio async def test_successful_call(self): """Successful call returns parsed JSON and posts to correct URL with headers.""" @@ -284,6 +299,58 @@ async def test_call_without_api_key_omits_auth_header(self): headers = call_kwargs[1]["headers"] assert "Authorization" not in headers + @pytest.mark.asyncio + async def test_call_applies_inference_headers(self): + """Per-request http_headers are sent alongside the derived base headers.""" + engine = APIEngine(base_url="https://api.example.com", endpoint="/v1/classify", api_key="test-key") + mock_client = self._start_with_recording_client(engine) + + await engine.call({"input": "Hi"}, http_headers={"X-Tenant": "acme"}) + + headers = mock_client.post.call_args[1]["headers"] + assert headers["X-Tenant"] == "acme" + assert headers["Content-Type"] == "application/json" + assert headers["Accept"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + + @pytest.mark.asyncio + async def test_inference_header_overrides_base_header_case_insensitive(self): + """A per-request header replaces a base header of the same name, regardless of case.""" + engine = APIEngine(base_url="https://api.example.com", endpoint="/v1/classify", api_key="test-key") + mock_client = self._start_with_recording_client(engine) + + await engine.call({"input": "Hi"}, http_headers={"authorization": "Bearer per-request"}) + + headers = mock_client.post.call_args[1]["headers"] + auth_keys = [key for key in headers if key.lower() == "authorization"] + assert auth_keys == ["authorization"] + assert headers["authorization"] == "Bearer per-request" + + @pytest.mark.asyncio + async def test_inference_headers_absent_from_body(self): + """http_headers configures transport and never becomes a request-body field.""" + engine = APIEngine(base_url="https://api.example.com", endpoint="/v1/classify", api_key="test-key") + mock_client = self._start_with_recording_client(engine) + + await engine.call({"input": "Hi"}, http_headers={"X-Tenant": "acme"}) + + assert mock_client.post.call_args[1]["json"] == {"input": "Hi"} + + @pytest.mark.asyncio + async def test_no_inference_headers_leaves_base_headers(self): + """Omitting http_headers leaves only the derived base headers.""" + engine = APIEngine(base_url="https://api.example.com", endpoint="/v1/classify", api_key="test-key") + mock_client = self._start_with_recording_client(engine) + + await engine.call({"input": "Hi"}) + + headers = mock_client.post.call_args[1]["headers"] + assert headers == { + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer test-key", + } + @pytest.mark.asyncio async def test_call_http_error_raises_api_engine_error(self): """HTTP 4xx/5xx raises APIEngineError with status and endpoint.""" diff --git a/tests/guardrails/test_engine_registry.py b/tests/guardrails/test_engine_registry.py index 9c840fe49c..00b3ef6041 100644 --- a/tests/guardrails/test_engine_registry.py +++ b/tests/guardrails/test_engine_registry.py @@ -29,6 +29,7 @@ from nemoguardrails.guardrails import telemetry from nemoguardrails.guardrails.api_engine import APIEngine from nemoguardrails.guardrails.engine_registry import EngineRegistry +from nemoguardrails.guardrails.guardrails_types import request_http_headers from nemoguardrails.guardrails.model_engine import ModelEngine from nemoguardrails.guardrails.tool_schema import Toolset from nemoguardrails.rails.llm.config import RailsConfig @@ -330,7 +331,7 @@ async def test_generate_from_correct_engine(self, manager): result = await manager.model_call("main", messages) assert result is expected - engine.chat_completion.assert_called_once_with(messages) + engine.chat_completion.assert_called_once_with(messages, http_headers=None) @pytest.mark.asyncio async def test_passes_kwargs_to_engine(self, manager): @@ -600,7 +601,8 @@ async def test_llm_params_take_precedence_over_config_parameters(self, span_expo @pytest.mark.asyncio async def test_model_call_excludes_non_body_keys(self, span_exporter): """Transport (base_url/timeout) and streaming-control (stream) keys in - parameters never reach the request body; only the sampling param does.""" + parameters never reach the request body; only the sampling param does, + alongside the explicit transport-only http_headers argument.""" tracer, _ = span_exporter registry = _registry_with_main_params( {"base_url": "https://custom.example.com", "timeout": 5, "stream": True, "temperature": 0.5}, @@ -612,7 +614,7 @@ async def test_model_call_excludes_non_body_keys(self, span_exporter): await registry.model_call("main", [{"role": "user", "content": "hi"}]) body = engine.chat_completion.call_args[1] - assert body == {"temperature": 0.5} + assert body == {"temperature": 0.5, "http_headers": None} @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) @pytest.mark.asyncio @@ -893,7 +895,7 @@ async def test_calls_correct_api_engine(self, manager): result = await manager.api_call("jailbreak_detection", {"input": "hello"}) assert result == mock_response - api_engine.call.assert_called_once_with({"input": "hello"}) + api_engine.call.assert_called_once_with({"input": "hello"}, http_headers=None) @pytest.mark.asyncio async def test_passes_kwargs_to_api_engine(self, manager): @@ -903,7 +905,7 @@ async def test_passes_kwargs_to_api_engine(self, manager): await manager.api_call("jailbreak_detection", {"input": "test"}, extra_param="value") - api_engine.call.assert_called_once_with({"input": "test"}, extra_param="value") + api_engine.call.assert_called_once_with({"input": "test"}, http_headers=None, extra_param="value") @pytest.mark.asyncio async def test_raises_key_error_for_unknown_api_name(self, manager): @@ -912,6 +914,95 @@ async def test_raises_key_error_for_unknown_api_name(self, manager): await manager.api_call("nonexistent", {"input": "test"}) +class TestEngineRegistryInferenceHeaders: + """The registry is the single choke point that reads the request-scoped + inference-time HTTP headers and forwards them to the engines as an explicit + transport argument, never as a request-body field.""" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_model_call_forwards_request_headers(self, manager): + """model_call passes the bound headers to chat_completion as http_headers.""" + engine = manager._get_engine("main", ModelEngine) + engine.chat_completion = AsyncMock(return_value=LLMResponse(content="ok")) + + with request_http_headers({"X-Tenant": "acme"}): + await manager.model_call("main", [{"role": "user", "content": "hi"}]) + + assert engine.chat_completion.call_args[1]["http_headers"] == {"X-Tenant": "acme"} + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_model_call_forwards_none_without_a_bound_scope(self, manager): + """Outside a request scope the engine is told there are no per-request headers.""" + engine = manager._get_engine("main", ModelEngine) + engine.chat_completion = AsyncMock(return_value=LLMResponse(content="ok")) + + await manager.model_call("main", [{"role": "user", "content": "hi"}]) + + assert engine.chat_completion.call_args[1]["http_headers"] is None + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_model_call_headers_reach_every_configured_model(self, manager): + """The same bound headers broadcast to a rail model, not just the main one.""" + main_engine = manager._get_engine("main", ModelEngine) + safety_engine = manager._get_engine("content_safety", ModelEngine) + main_engine.chat_completion = AsyncMock(return_value=LLMResponse(content="ok")) + safety_engine.chat_completion = AsyncMock(return_value=LLMResponse(content="safe")) + + with request_http_headers({"X-Tenant": "acme"}): + await manager.model_call("main", [{"role": "user", "content": "hi"}]) + await manager.model_call("content_safety", [{"role": "user", "content": "hi"}]) + + assert main_engine.chat_completion.call_args[1]["http_headers"] == {"X-Tenant": "acme"} + assert safety_engine.chat_completion.call_args[1]["http_headers"] == {"X-Tenant": "acme"} + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_stream_model_call_forwards_request_headers(self, manager): + """stream_model_call passes the bound headers through as http_headers.""" + engine = manager._get_engine("main", ModelEngine) + + captured: dict = {} + + async def _capturing_stream(messages, **kwargs): # noqa: ARG001 (signature dictated by ModelEngine) + captured.update(kwargs) + yield LLMResponseChunk(delta_content="hi", finish_reason="stop") + + engine.stream_chat_completion = _capturing_stream + + with request_http_headers({"X-Tenant": "acme"}): + async for _ in manager.stream_model_call("main", [{"role": "user", "content": "hi"}]): + pass + + assert captured["http_headers"] == {"X-Tenant": "acme"} + + @pytest.mark.asyncio + async def test_api_call_forwards_request_headers(self, manager): + """api_call passes the bound headers to the API engine as http_headers.""" + api_engine = manager._get_engine("jailbreak_detection", APIEngine) + api_engine.call = AsyncMock(return_value={"jailbreak": False, "score": -0.95}) + + with request_http_headers({"X-Tenant": "acme"}): + await manager.api_call("jailbreak_detection", {"input": "hello"}) + + api_engine.call.assert_called_once_with({"input": "hello"}, http_headers={"X-Tenant": "acme"}) + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_request_headers_do_not_enter_the_request_body(self, manager): + """The bound headers stay out of the merged body params sent to the engine.""" + engine = manager._get_engine("main", ModelEngine) + engine.chat_completion = AsyncMock(return_value=LLMResponse(content="ok")) + + with request_http_headers({"X-Tenant": "acme"}): + await manager.model_call("main", [{"role": "user", "content": "hi"}], temperature=0.2) + + sent_kwargs = engine.chat_completion.call_args[1] + assert sent_kwargs == {"temperature": 0.2, "http_headers": {"X-Tenant": "acme"}} + + class TestEngineRegistryApiEngineStartErrors: """Test start() error handling for API engines.""" diff --git a/tests/guardrails/test_guardrails.py b/tests/guardrails/test_guardrails.py index 2ef5d93455..ee06a54832 100644 --- a/tests/guardrails/test_guardrails.py +++ b/tests/guardrails/test_guardrails.py @@ -178,12 +178,17 @@ async def mock_stream(): with pytest.raises(NotImplementedError, match="IORails doesn't support update_llm()"): guardrails.update_llm(mock_new_llm) - guardrails.rails_engine.generate.assert_called_once_with(prompt=None, messages=messages, options=None) - guardrails.rails_engine.generate_async.assert_called_once_with(prompt=None, messages=messages, options=None) + guardrails.rails_engine.generate.assert_called_once_with( + prompt=None, messages=messages, options=None, http_headers=None + ) + guardrails.rails_engine.generate_async.assert_called_once_with( + prompt=None, messages=messages, options=None, http_headers=None + ) guardrails.rails_engine.stream_async.assert_called_once_with( messages=messages, options=None, include_metadata=False, + http_headers=None, ) @pytest.mark.asyncio @@ -1230,6 +1235,7 @@ async def mock_stream(): messages=[{"role": "user", "content": "hi"}], options=None, include_metadata=False, + http_headers=None, ) @pytest.mark.asyncio @@ -1260,6 +1266,7 @@ async def mock_stream(): messages=[{"role": "user", "content": "hi"}], options=opts, include_metadata=True, + http_headers=None, ) @pytest.mark.asyncio @@ -1294,6 +1301,7 @@ async def mock_stream(): messages=[{"role": "user", "content": "hi"}], options=None, include_metadata=False, + http_headers=None, ) @pytest.mark.asyncio @@ -1318,6 +1326,7 @@ async def mock_stream(): messages=[{"role": "user", "content": "hello"}], options=None, include_metadata=False, + http_headers=None, ) @@ -2010,3 +2019,138 @@ async def test_options_forwarded_unchanged(self, _mock_init, _content_safety_rai assert guardrails.rails_engine.generate.call_args.kwargs["options"] is options assert guardrails.rails_engine.generate_async.call_args.kwargs["options"] is options + + +class TestHttpHeadersForwarding: + """The facade forwards per-request ``http_headers`` to IORails and rejects + them on the LLMRails fallback, which has no inference-time header support.""" + + @pytest.mark.asyncio + @patch.object(IORails, "stop", new_callable=AsyncMock) + @patch.object(IORails, "start", new_callable=AsyncMock) + @patch.object(IORails, "__init__", return_value=None) + async def test_generate_forwards_headers_to_iorails( + self, mock_init, mock_start, mock_stop, _content_safety_rails_config + ): + """The sync generate passes http_headers straight through to IORails.""" + headers = {"X-Tenant": "acme"} + messages = [{"role": "user", "content": "hi"}] + + async with Guardrails(config=_content_safety_rails_config, use_iorails=True) as guardrails: + engine = _iorails_engine(guardrails) + engine.generate = MagicMock(return_value="ok") + + guardrails.generate(messages=messages, http_headers=headers) + + assert engine.generate.call_args.kwargs["http_headers"] is headers + + @pytest.mark.asyncio + @patch.object(IORails, "stop", new_callable=AsyncMock) + @patch.object(IORails, "start", new_callable=AsyncMock) + @patch.object(IORails, "__init__", return_value=None) + async def test_generate_async_forwards_headers_to_iorails( + self, mock_init, mock_start, mock_stop, _content_safety_rails_config + ): + """The async generate passes http_headers straight through to IORails.""" + headers = {"X-Tenant": "acme"} + messages = [{"role": "user", "content": "hi"}] + + async with Guardrails(config=_content_safety_rails_config, use_iorails=True) as guardrails: + engine = _iorails_engine(guardrails) + engine.generate_async = AsyncMock(return_value="ok") + + await guardrails.generate_async(messages=messages, http_headers=headers) + + assert engine.generate_async.call_args.kwargs["http_headers"] is headers + + @pytest.mark.asyncio + @patch.object(IORails, "stop", new_callable=AsyncMock) + @patch.object(IORails, "start", new_callable=AsyncMock) + @patch.object(IORails, "__init__", return_value=None) + async def test_stream_async_forwards_headers_to_iorails( + self, mock_init, mock_start, mock_stop, _content_safety_rails_config + ): + """http_headers is a supported stream_async kwarg and reaches IORails.""" + headers = {"X-Tenant": "acme"} + + async def mock_stream(): + yield "ok" + + async with Guardrails(config=_content_safety_rails_config, use_iorails=True) as guardrails: + engine = _iorails_engine(guardrails) + engine.stream_async = MagicMock(return_value=mock_stream()) + + with patch("nemoguardrails.guardrails.guardrails.log") as mock_log: + async for _ in guardrails.stream_async( + messages=[{"role": "user", "content": "hi"}], http_headers=headers + ): + pass + + mock_log.warning.assert_not_called() + + assert engine.stream_async.call_args.kwargs["http_headers"] is headers + + @pytest.mark.asyncio + @patch.object(LLMRails, "__init__", return_value=None) + async def test_generate_rejects_headers_on_llmrails(self, _mock_init, _content_safety_rails_config): + """The sync generate raises rather than silently dropping headers on the LLMRails fallback.""" + async with Guardrails(config=_content_safety_rails_config, use_iorails=False) as guardrails: + guardrails.rails_engine.generate = MagicMock(return_value="ok") + + with pytest.raises( + NotImplementedError, match=r"LLMRails doesn't support inference-time HTTP headers in generate\(\)" + ): + guardrails.generate(messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"}) + + guardrails.rails_engine.generate.assert_not_called() + + @pytest.mark.asyncio + @patch.object(LLMRails, "__init__", return_value=None) + async def test_generate_async_rejects_headers_on_llmrails(self, _mock_init, _content_safety_rails_config): + """The async generate raises rather than silently dropping headers on the LLMRails fallback.""" + async with Guardrails(config=_content_safety_rails_config, use_iorails=False) as guardrails: + guardrails.rails_engine.generate_async = AsyncMock(return_value="ok") + + with pytest.raises( + NotImplementedError, + match=r"LLMRails doesn't support inference-time HTTP headers in generate_async\(\)", + ): + await guardrails.generate_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + + guardrails.rails_engine.generate_async.assert_not_called() + + @pytest.mark.asyncio + @patch.object(LLMRails, "__init__", return_value=None) + async def test_stream_async_rejects_headers_on_llmrails(self, _mock_init, _content_safety_rails_config): + """Streaming raises on the LLMRails fallback before the engine is reached.""" + + async def mock_stream(): + yield "ok" + + async with Guardrails(config=_content_safety_rails_config, use_iorails=False) as guardrails: + guardrails.rails_engine.stream_async = MagicMock(return_value=mock_stream()) + + with pytest.raises( + NotImplementedError, + match=r"LLMRails doesn't support inference-time HTTP headers in stream_async\(\)", + ): + guardrails.stream_async(messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"}) + + guardrails.rails_engine.stream_async.assert_not_called() + + @pytest.mark.asyncio + @patch.object(LLMRails, "__init__", return_value=None) + async def test_llmrails_unaffected_without_headers(self, _mock_init, _content_safety_rails_config): + """Omitting http_headers leaves the LLMRails path exactly as before.""" + messages = [{"role": "user", "content": "hi"}] + + async with Guardrails(config=_content_safety_rails_config, use_iorails=False) as guardrails: + guardrails.rails_engine.generate_async = AsyncMock(return_value="ok") + + assert await guardrails.generate_async(messages=messages) == "ok" + + guardrails.rails_engine.generate_async.assert_awaited_once_with( + prompt=None, messages=messages, options=None + ) diff --git a/tests/guardrails/test_guardrails_types.py b/tests/guardrails/test_guardrails_types.py index c5d6c48096..b3b561f45f 100644 --- a/tests/guardrails/test_guardrails_types.py +++ b/tests/guardrails/test_guardrails_types.py @@ -15,7 +15,19 @@ """Unit tests for guardrails_types module.""" -from nemoguardrails.guardrails.guardrails_types import LLMMessage, LLMMessages, RailResult, truncate +import contextvars +from unittest.mock import MagicMock, patch + +import pytest + +from nemoguardrails.guardrails.guardrails_types import ( + LLMMessage, + LLMMessages, + RailResult, + get_http_headers, + request_http_headers, + truncate, +) class TestRailResult: @@ -123,3 +135,76 @@ def test_llm_messages_is_list(self): ] assert isinstance(msgs, list) assert all(isinstance(m, dict) for m in msgs) + + +class TestRequestHttpHeaders: + """Tests for the request-scoped inference-time HTTP header ContextVar.""" + + def test_defaults_to_none_outside_a_request(self): + """Without an active request scope there are no inference-time headers.""" + assert get_http_headers() is None + + def test_headers_visible_inside_scope(self): + """Headers bound by the context manager are readable for the duration of the scope.""" + with request_http_headers({"X-Tenant": "acme"}): + assert get_http_headers() == {"X-Tenant": "acme"} + + def test_headers_cleared_after_scope(self): + """Leaving the scope restores the previous (absent) value so headers don't leak between requests.""" + with request_http_headers({"X-Tenant": "acme"}): + pass + assert get_http_headers() is None + + def test_headers_cleared_when_scope_raises(self): + """A failing request still resets the ContextVar.""" + with pytest.raises(RuntimeError): + with request_http_headers({"X-Tenant": "acme"}): + raise RuntimeError("request failed") + assert get_http_headers() is None + + def test_none_binds_none(self): + """Passing None binds None rather than an empty dict.""" + with request_http_headers(None): + assert get_http_headers() is None + + def test_empty_mapping_binds_empty_dict(self): + """An empty mapping is preserved as an empty dict, distinct from None.""" + with request_http_headers({}): + assert get_http_headers() == {} + + def test_values_coerced_to_str(self): + """Non-string names and values are coerced, matching config default_headers handling.""" + with request_http_headers({"X-Count": 3, "X-Flag": True}): + headers = get_http_headers() + assert headers == {"X-Count": "3", "X-Flag": "True"} + + def test_bound_headers_do_not_alias_the_caller_mapping(self): + """The bound headers are a copy, so later caller mutation cannot change the request's headers.""" + caller_headers = {"X-Tenant": "acme"} + with request_http_headers(caller_headers): + caller_headers["X-Tenant"] = "other" + assert get_http_headers() == {"X-Tenant": "acme"} + + def test_nested_scope_restores_outer_headers(self): + """An inner scope shadows the outer headers and restores them on exit.""" + with request_http_headers({"X-Tenant": "outer"}): + with request_http_headers({"X-Tenant": "inner"}): + assert get_http_headers() == {"X-Tenant": "inner"} + assert get_http_headers() == {"X-Tenant": "outer"} + + def test_reset_from_a_different_context_is_tolerated(self): + """Exiting the scope from another context, as async-generator teardown does, does not raise.""" + scope = request_http_headers({"X-Tenant": "acme"}) + contextvars.Context().run(scope.__enter__) + + scope.__exit__(None, None, None) + + def test_unexpected_reset_error_is_reraised(self): + """A reset failure other than the cross-context one surfaces instead of being swallowed.""" + fake_var = MagicMock() + fake_var.reset.side_effect = ValueError("boom") + + with patch("nemoguardrails.guardrails.guardrails_types._http_headers_var", fake_var): + with pytest.raises(ValueError, match="boom"): + with request_http_headers({"X-Tenant": "acme"}): + pass diff --git a/tests/guardrails/test_iorails.py b/tests/guardrails/test_iorails.py index 1068cab5e7..895f0fa68e 100644 --- a/tests/guardrails/test_iorails.py +++ b/tests/guardrails/test_iorails.py @@ -16,6 +16,7 @@ """Unit tests for iorails module.""" import asyncio +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -24,7 +25,7 @@ from aiohttp.test_utils import TestServer from nemoguardrails import Guardrails -from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult +from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult, get_http_headers from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails from nemoguardrails.guardrails.model_engine import ModelEngine from nemoguardrails.rails.llm.config import RailsConfig @@ -202,6 +203,184 @@ async def handler(request): assert "X-Tenant" not in captured["body"] +_MAIN_MODEL = "meta/llama-3.3-70b-instruct" +_SAFETY_MODEL = "nvidia/llama-3.1-nemoguard-8b-content-safety" +_SAFE_VERDICT = json.dumps({"User Safety": "safe", "Response Safety": "safe"}) + + +class TestInferenceHttpHeaders: + """generate_async binds its http_headers for the whole request so every + downstream model call sees them, and clears them when the request ends.""" + + @pytest.mark.asyncio + async def test_headers_bound_at_the_main_model_call(self, iorails): + """Headers passed to generate_async are readable at the main LLM call.""" + seen: dict = {} + + async def _record_headers(model_type, messages, **kwargs): + seen[model_type] = get_http_headers() + return LLMResponse(content="Hello") + + iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.engine_registry.model_call = _record_headers + + await iorails.generate_async(messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"}) + + assert seen["main"] == {"X-Tenant": "acme"} + + @pytest.mark.asyncio + async def test_headers_cleared_after_the_request(self, iorails): + """The request scope is torn down so headers never leak into the next request.""" + iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.engine_registry.model_call = AsyncMock(return_value=LLMResponse(content="Hello")) + + await iorails.generate_async(messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"}) + + assert get_http_headers() is None + + @pytest.mark.asyncio + async def test_no_headers_leaves_the_scope_unbound(self, iorails): + """Omitting http_headers leaves nothing bound, so only config headers apply.""" + seen: dict = {} + + async def _record_headers(model_type, messages, **kwargs): + seen[model_type] = get_http_headers() + return LLMResponse(content="Hello") + + iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.engine_registry.model_call = _record_headers + + await iorails.generate_async(messages=[{"role": "user", "content": "hi"}]) + + assert seen["main"] is None + + @pytest.mark.asyncio + async def test_headers_bound_when_the_request_fails(self, iorails): + """A request that raises still clears the bound headers.""" + iorails.rails_manager.is_input_safe = AsyncMock(side_effect=RuntimeError("rail exploded")) + + with pytest.raises(RuntimeError, match="rail exploded"): + await iorails.generate_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + + assert get_http_headers() is None + + +class TestInferenceHttpHeadersOverHTTP: + """True end-to-end: run generate_async against a loopback HTTP server and + assert on the headers the main LLM and the content-safety rail model + actually put on the wire.""" + + @staticmethod + def _build_app(captured: dict): + """Serve chat completions, recording the headers and body sent per model.""" + + async def handler(request): + body = await request.json() + captured[body["model"]] = {"headers": dict(request.headers), "body": body} + content = _SAFE_VERDICT if body["model"] == _SAFETY_MODEL else "ok" + return web.json_response({"choices": [{"message": {"role": "assistant", "content": content}}]}) + + app = web.Application() + app.router.add_post("/v1/chat/completions", handler) + return app + + @staticmethod + def _build_config(base_url: str): + """A content-safety config whose two models point at the loopback server + and carry distinct default_headers, one of which the request will override.""" + return RailsConfig.from_content( + config={ + **CONTENT_SAFETY_CONFIG, + "models": [ + { + "type": "main", + "engine": "nim", + "model": _MAIN_MODEL, + "parameters": { + "base_url": base_url, + "default_headers": {"X-Tenant": "from-config", "X-Main-Route": "main-pool"}, + }, + }, + { + "type": "content_safety", + "engine": "nim", + "model": _SAFETY_MODEL, + "parameters": { + "base_url": base_url, + "default_headers": {"X-Safety-Route": "safety-pool"}, + }, + }, + ], + } + ) + + async def _run_request(self, captured: dict): + """Run one generate_async with inference-time headers against the loopback server.""" + server = TestServer(self._build_app(captured)) + await server.start_server() + try: + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + iorails = IORails(self._build_config(str(server.make_url("/")))) + async with iorails: + return await iorails.generate_async( + messages=[{"role": "user", "content": "hi"}], + http_headers={"X-Tenant": "from-request", "X-Trace": "abc123"}, + ) + finally: + await server.close() + + @pytest.mark.asyncio + async def test_headers_reach_the_main_llm_on_the_wire(self): + """The main LLM request carries the inference-time headers on top of its own config headers.""" + captured: dict = {} + + result = await self._run_request(captured) + + assert result == {"role": "assistant", "content": "ok"} + main_headers = captured[_MAIN_MODEL]["headers"] + assert main_headers["X-Trace"] == "abc123" + assert main_headers["X-Main-Route"] == "main-pool" + assert main_headers["Authorization"] == "Bearer test-key" + + @pytest.mark.asyncio + async def test_headers_broadcast_to_the_rail_model_on_the_wire(self): + """The content-safety rail model request carries the same inference-time headers.""" + captured: dict = {} + + await self._run_request(captured) + + safety_headers = captured[_SAFETY_MODEL]["headers"] + assert safety_headers["X-Trace"] == "abc123" + assert safety_headers["X-Tenant"] == "from-request" + assert safety_headers["X-Safety-Route"] == "safety-pool" + assert "X-Main-Route" not in safety_headers + + @pytest.mark.asyncio + async def test_inference_header_overrides_the_config_header_on_the_wire(self): + """Where both layers set the same name, the per-request value is what is sent.""" + captured: dict = {} + + await self._run_request(captured) + + assert captured[_MAIN_MODEL]["headers"]["X-Tenant"] == "from-request" + + @pytest.mark.asyncio + async def test_headers_never_reach_the_request_body(self): + """Inference-time headers are transport-only and absent from every JSON body sent.""" + captured: dict = {} + + await self._run_request(captured) + + for model_name, sent in captured.items(): + assert "http_headers" not in sent["body"], model_name + assert "X-Trace" not in sent["body"], model_name + + class TestGenerateAsync: """Test the generate_async input-check → LLM → output-check pipeline.""" diff --git a/tests/guardrails/test_iorails_streaming.py b/tests/guardrails/test_iorails_streaming.py index 5d6e783953..7f75918882 100644 --- a/tests/guardrails/test_iorails_streaming.py +++ b/tests/guardrails/test_iorails_streaming.py @@ -24,7 +24,7 @@ import pytest_asyncio from nemoguardrails.exceptions import StreamingNotSupportedError -from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.guardrails_types import RailResult, get_http_headers from nemoguardrails.guardrails.iorails import ( REFUSAL_MESSAGE, STREAM_MAX_CONCURRENCY, @@ -1082,3 +1082,96 @@ async def _stream(model_type, messages, **kwargs): terminal = empty_frames[0]["metadata"] assert terminal["usage"] == {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3} assert "response_metadata" in terminal and "usage_metadata" in terminal + + +class TestStreamAsyncInferenceHttpHeaders: + """stream_async binds its http_headers for the whole stream, so the main LLM + call and the rails running in the background generation task all see them, + and the binding is cleared once the stream ends.""" + + @pytest.mark.asyncio + async def test_headers_bound_at_the_main_stream_call(self, iorails_input_only): + """Headers passed to stream_async are readable at the main streaming LLM call.""" + seen: dict = {} + + async def _recording_stream(model_type, messages, **kwargs): + seen["stream"] = get_http_headers() + yield LLMResponseChunk(delta_content="Hi") + + _wire_mocks(iorails_input_only, stream=_recording_stream) + + await _collect( + iorails_input_only.stream_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + ) + + assert seen["stream"] == {"X-Tenant": "acme"} + + @pytest.mark.asyncio + async def test_headers_visible_to_rails_in_the_generation_task(self, iorails_input_only): + """The background generation task inherits the headers from its context snapshot.""" + seen: dict = {} + + async def _recording_input_rail(messages, *, enabled=True): + seen["input_rail"] = get_http_headers() + return RailResult(is_safe=True) + + _wire_mocks(iorails_input_only) + iorails_input_only.rails_manager.is_input_safe = _recording_input_rail + + await _collect( + iorails_input_only.stream_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + ) + + assert seen["input_rail"] == {"X-Tenant": "acme"} + + @pytest.mark.asyncio + async def test_headers_visible_to_streaming_output_rails(self, iorails_stream_first): + """Output rails checking buffered chunks see the same headers as the main call.""" + seen: dict = {} + + async def _recording_output_rail(messages, bot_response, *, enabled=True): + seen["output_rail"] = get_http_headers() + return RailResult(is_safe=True) + + _wire_mocks(iorails_stream_first) + iorails_stream_first.rails_manager.is_output_safe = _recording_output_rail + + await _collect( + iorails_stream_first.stream_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + ) + + assert seen["output_rail"] == {"X-Tenant": "acme"} + + @pytest.mark.asyncio + async def test_headers_cleared_after_the_stream(self, iorails_input_only): + """Draining the stream tears the scope down so headers can't leak into the next request.""" + _wire_mocks(iorails_input_only) + + await _collect( + iorails_input_only.stream_async( + messages=[{"role": "user", "content": "hi"}], http_headers={"X-Tenant": "acme"} + ) + ) + + assert get_http_headers() is None + + @pytest.mark.asyncio + async def test_no_headers_leaves_the_scope_unbound(self, iorails_input_only): + """Omitting http_headers leaves nothing bound, so only config headers apply.""" + seen: dict = {} + + async def _recording_stream(model_type, messages, **kwargs): + seen["stream"] = get_http_headers() + yield LLMResponseChunk(delta_content="Hi") + + _wire_mocks(iorails_input_only, stream=_recording_stream) + + await _collect(iorails_input_only.stream_async(messages=[{"role": "user", "content": "hi"}])) + + assert seen["stream"] is None diff --git a/tests/guardrails/test_model_engine.py b/tests/guardrails/test_model_engine.py index 1ad8b8bbd1..a05ca615c8 100644 --- a/tests/guardrails/test_model_engine.py +++ b/tests/guardrails/test_model_engine.py @@ -782,6 +782,168 @@ def test_default_headers_mapping_is_immutable(self): headers["X-Injected"] = "nope" +class TestModelEngineInferenceHeaders: + """Test that per-request http_headers layer over the base and config headers.""" + + @staticmethod + def _mock_client(): + """Build a mock aiohttp client whose post() records call args.""" + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={"choices": [{"message": {"content": "ok"}}]}) + + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + return mock_client + + @staticmethod + def _headers_from(mock_client): + """Extract the headers dict passed to the mocked post().""" + return mock_client.post.call_args[1]["headers"] + + @staticmethod + def _body_from(mock_client): + """Extract the JSON body passed to the mocked post().""" + return mock_client.post.call_args[1]["json"] + + def _started_engine(self, parameters=None): + """Build a started ModelEngine with a recording mock client.""" + engine = ModelEngine(_make_model(parameters=parameters)) + engine._client = self._mock_client() + engine._running = True + return engine + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_inference_headers_merged_into_request(self): + """A per-request header is sent alongside the base headers.""" + engine = self._started_engine() + + await engine.call([{"role": "user", "content": "Hi"}], http_headers={"X-Request-Tag": "batch-7"}) + + headers = self._headers_from(engine._client) + assert headers["X-Request-Tag"] == "batch-7" + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_config_headers_kept_when_names_do_not_collide(self): + """Config and per-request headers with different names both reach the request.""" + engine = self._started_engine(parameters={"default_headers": {"X-Tenant": "acme"}}) + + await engine.call([{"role": "user", "content": "Hi"}], http_headers={"X-Request-Tag": "batch-7"}) + + headers = self._headers_from(engine._client) + assert headers["X-Tenant"] == "acme" + assert headers["X-Request-Tag"] == "batch-7" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_inference_header_overrides_config_header_case_insensitive(self): + """On a name collision the per-request header wins over the configured one, regardless of case.""" + engine = self._started_engine(parameters={"default_headers": {"X-Tenant": "from-config"}}) + + await engine.call([{"role": "user", "content": "Hi"}], http_headers={"x-tenant": "from-request"}) + + headers = self._headers_from(engine._client) + tenant_keys = [key for key in headers if key.lower() == "x-tenant"] + assert tenant_keys == ["x-tenant"] + assert headers["x-tenant"] == "from-request" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_inference_header_overrides_authorization(self): + """A per-request Authorization header replaces the api-key-derived one.""" + engine = self._started_engine() + + await engine.call([{"role": "user", "content": "Hi"}], http_headers={"Authorization": "Bearer per-request"}) + + headers = self._headers_from(engine._client) + assert headers["Authorization"] == "Bearer per-request" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_full_precedence_base_then_config_then_request(self): + """Base < config < per-request: each layer overrides the one before it.""" + engine = self._started_engine( + parameters={"default_headers": {"Authorization": "Bearer from-config", "X-Tenant": "acme"}} + ) + + await engine.call( + [{"role": "user", "content": "Hi"}], + http_headers={"Authorization": "Bearer from-request"}, + ) + + headers = self._headers_from(engine._client) + assert headers["Authorization"] == "Bearer from-request" + assert headers["X-Tenant"] == "acme" + assert headers["Content-Type"] == "application/json" + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_inference_headers_absent_from_body(self): + """http_headers is a transport argument and never becomes a request-body field.""" + engine = self._started_engine() + + await engine.call([{"role": "user", "content": "Hi"}], http_headers={"X-Request-Tag": "batch-7"}) + + body = self._body_from(engine._client) + assert "http_headers" not in body + assert "X-Request-Tag" not in body + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_no_inference_headers_leaves_config_headers_untouched(self): + """Omitting http_headers leaves the configured headers exactly as PR-1 behavior.""" + engine = self._started_engine(parameters={"default_headers": {"X-Tenant": "acme"}}) + + await engine.call([{"role": "user", "content": "Hi"}]) + + headers = self._headers_from(engine._client) + assert headers == { + "Content-Type": "application/json", + "Authorization": "Bearer test-key", + "X-Tenant": "acme", + } + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + def test_http_headers_excluded_from_body_param_defaults(self): + """A model configured with parameters.http_headers never forwards it as a body field.""" + engine = ModelEngine(_make_model(parameters={"http_headers": {"X-Tenant": "acme"}, "temperature": 0.5})) + assert "http_headers" not in engine.body_param_defaults + assert engine.body_param_defaults["temperature"] == 0.5 + + @patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}) + @pytest.mark.asyncio + async def test_stream_call_applies_inference_headers(self): + """Streaming requests carry the per-request headers layered over the configured ones.""" + engine = ModelEngine(_make_model(parameters={"default_headers": {"X-Tenant": "acme"}})) + sse_lines = [ + b'data: {"choices":[{"delta":{"content":"Hi"}}]}\n\n', + b"data: [DONE]\n\n", + ] + mock_client = MagicMock() + mock_client.post = MagicMock(return_value=_mock_streaming_response(sse_lines)) + engine._client = mock_client + engine._running = True + + chunks = [ + chunk + async for chunk in engine.stream_call( + [{"role": "user", "content": "Hi"}], http_headers={"X-Request-Tag": "batch-7"} + ) + ] + + assert [chunk.delta_content for chunk in chunks] == ["Hi"] + headers = self._headers_from(mock_client) + assert headers["X-Tenant"] == "acme" + assert headers["X-Request-Tag"] == "batch-7" + assert "http_headers" not in self._body_from(mock_client) + + class TestModelEngineStreamCall: """Test ModelEngine.stream_call() SSE streaming."""