From 49ae76f5d382fe7d5f7368977aa7b5db5f3a5df6 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:37:35 -0500 Subject: [PATCH 01/11] Initial checkin of GenerationResponse code (non-streaming only, no token usage) --- nemoguardrails/guardrails/iorails.py | 139 +++++- tests/guardrails/test_iorails.py | 7 +- .../test_iorails_generation_response.py | 405 ++++++++++++++++++ tests/guardrails/test_tool_rails_iorails.py | 18 +- 4 files changed, 544 insertions(+), 25 deletions(-) create mode 100644 tests/guardrails/test_iorails_generation_response.py diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 4b8df06120..9a92eb02e7 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -63,10 +63,16 @@ from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop from nemoguardrails.rails.llm.buffer import get_buffer_strategy from nemoguardrails.rails.llm.config import RailsConfig, _get_flow_name -from nemoguardrails.rails.llm.options import GenerationOptions, RailsResult, RailStatus, RailType +from nemoguardrails.rails.llm.options import ( + GenerationOptions, + GenerationResponse, + RailsResult, + RailStatus, + RailType, +) from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler from nemoguardrails.tracing.constants import GuardrailsAttributes -from nemoguardrails.types import LLMModel, LLMResponse, ToolCall +from nemoguardrails.types import LLMModel, LLMResponse, ToolCall, UsageInfo if TYPE_CHECKING: from opentelemetry.trace import Span @@ -170,6 +176,105 @@ def _build_assistant_message(content: str, tool_calls: Optional[list[ToolCall]]) } +def _usage_to_dict(usage: UsageInfo) -> dict: + """Serialize ``UsageInfo`` token counts to a dict, dropping ``None`` values. + + ``input_tokens``/``output_tokens``/``total_tokens`` default to 0 and are always + kept; ``reasoning_tokens``/``cached_tokens`` default to ``None`` and are omitted + unless the provider reported them. + """ + fields = { + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + "total_tokens": usage.total_tokens, + "reasoning_tokens": usage.reasoning_tokens, + "cached_tokens": usage.cached_tokens, + } + return {key: value for key, value in fields.items() if value is not None} + + +def _build_llm_metadata(response: LLMResponse) -> Optional[dict]: + """Assemble ``llm_metadata`` from the main call's ``provider_metadata`` plus usage. + + Mirrors LLMRails' ``llm_metadata`` (the provider metadata blob) and adds a + structured ``usage`` sub-key with token counts that LLMRails does not expose + outside its ``log``. Returns ``None`` when neither is present. + """ + metadata = dict(response.provider_metadata or {}) + if response.usage: + usage = _usage_to_dict(response.usage) + if usage: + metadata["usage"] = usage + return metadata or None + + +def _build_generation_response( + response_text: str, reasoning_content: Optional[str], response: LLMResponse +) -> GenerationResponse: + """Build the structured ``GenerationResponse`` returned when ``options`` are supplied. + + Reasoning goes to the ``reasoning_content`` field with clean message content (no + inline ```` prefix — that is the bare-return shape). Tool calls use the + canonical ``ToolCall.to_dict()`` shape (dict arguments), matching LLMRails' + ``GenerationResponse.tool_calls`` rather than the OpenAI-wire JSON-string shape + used for the bare message. ``llm_output`` stays ``None`` to match LLMRails, whose + ``raw_response`` source is never populated. + """ + result = GenerationResponse(response=[{"role": "assistant", "content": response_text}]) + if reasoning_content: + result.reasoning_content = reasoning_content + if response.tool_calls: + result.tool_calls = [tool_call.to_dict() for tool_call in response.tool_calls] + llm_metadata = _build_llm_metadata(response) + if llm_metadata: + result.llm_metadata = llm_metadata + return result + + +def _finalize_refusal(structured: bool) -> Union[LLMMessage, GenerationResponse]: + """Shape the refusal message for the active return contract (structured vs bare).""" + message: LLMMessage = {"role": "assistant", "content": REFUSAL_MESSAGE} + if structured: + return GenerationResponse(response=[message]) + return message + + +def _response_content_for_capture(result: Union[LLMMessage, GenerationResponse]) -> Optional[str]: + """Extract the assistant content for content capture from either return shape.""" + if isinstance(result, GenerationResponse): + response = result.response + if isinstance(response, list) and response: + last = response[-1] + content = last.get("content") if isinstance(last, dict) else None + return content if isinstance(content, str) else None + return response if isinstance(response, str) else None + return result.get("content") + + +def _raise_on_unsupported_options(options: Optional[GenerationOptions], state: object) -> None: + """Raise for ``GenerationOptions``/``state`` features IORails cannot honor. + + ``state`` and ``output_vars``/``output_data`` require Colang runtime state that + IORails does not have and raise ``ValueError``; ``log`` is deferred to a follow-up + PR and raises ``NotImplementedError``. These raise rather than silently returning + empty data, mirroring how LLMRails rejects options it cannot fulfill. + """ + if state is not None: + raise ValueError("state is not supported by IORails; it is a stateless input/output rails engine") + if options is None: + return + if options.output_vars: + raise ValueError("output_vars/output_data is not supported by IORails; it has no Colang context to return") + log_options = options.log + if log_options and ( + log_options.activated_rails + or log_options.llm_calls + or log_options.internal_events + or log_options.colang_history + ): + raise NotImplementedError("GenerationResponse `log` is not supported by IORails") + + def _coerce_generation_options(options: Optional[Union[dict, GenerationOptions]]) -> Optional[GenerationOptions]: """Normalize the request ``options`` argument into a ``GenerationOptions`` or None.""" if isinstance(options, GenerationOptions): @@ -464,7 +569,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): """Context manager (used for testing rather than long-lived instance)""" await self.stop() - def generate(self, messages: LLMMessages, **kwargs) -> LLMMessage: + def generate(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: """Synchronous version of generate_async. Telemetry is disabled for the ephemeral IORails object used for @@ -488,7 +593,7 @@ async def _run_sync_iorails(): return asyncio.run(_run_sync_iorails()) - async def generate_async(self, messages: LLMMessages, **kwargs) -> LLMMessage: + async def generate_async(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: """Public entry: submit the request to the internal work queue. The queue enforces non-streaming concurrency limits @@ -514,7 +619,7 @@ async def generate_async(self, messages: LLMMessages, **kwargs) -> LLMMessage: record_nonstream_rejected() raise - async def _run_generate(self, messages: LLMMessages, **kwargs) -> LLMMessage: + async def _run_generate(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: """Runs inside a queue worker task. Wraps the pipeline in ``traced_request`` so each request gets its own span + request ID, then delegates to ``_do_generate`` for the actual input rails → @@ -533,7 +638,7 @@ async def _run_generate(self, messages: LLMMessages, **kwargs) -> LLMMessage: # Capture content once here at the traced_request boundary so any # future early-return added to _do_generate is covered automatically. if self._content_capture_enabled: - set_request_content(request_span, messages, result.get("content")) + set_request_content(request_span, messages, _response_content_for_capture(result)) elapsed_ms = (time.monotonic() - t0) * 1000 log.info("[%s] generate_async completed time=%.1fms", req_id, elapsed_ms) return result @@ -560,12 +665,16 @@ def _guardrails_violation_payload(message: str, param: str) -> str: async def _do_generate( self, messages: LLMMessages, req_id: str, request_span: Optional["Span"] = None, **kwargs - ) -> LLMMessage: + ) -> Union[LLMMessage, GenerationResponse]: """Core pipeline: tool-result rails -> input rails -> LLM call -> tool-call + output rails.""" log.info("[%s] generate_async called", req_id) log.debug("[%s] generate_async messages=%s", req_id, truncate(messages)) options = _coerce_generation_options(kwargs.get("options")) + _raise_on_unsupported_options(options, kwargs.get("state")) + # When options are supplied we return a structured GenerationResponse + # (mirroring LLMRails' `if gen_options:` branch); otherwise a bare LLMMessage. + has_generation_response = options is not None # Pass llm_params (including tool definitions) unchanged to the LLM call. llm_kwargs = options.llm_params if (options and options.llm_params) else {} input_enabled = options.rails.input if options else True @@ -581,7 +690,7 @@ async def _do_generate( log.info("[%s] Tool result blocked: %s", req_id, tool_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.INPUT) - return {"role": "assistant", "content": REFUSAL_MESSAGE} + return _finalize_refusal(has_generation_response) if self._speculative_generation: response = await self._do_generate_speculative( @@ -591,7 +700,7 @@ async def _do_generate( response = await self._do_generate_sequential(messages, req_id, llm_kwargs, input_enabled=input_enabled) if response is None: - return {"role": "assistant", "content": REFUSAL_MESSAGE} + return _finalize_refusal(has_generation_response) # Log raw content before reasoning extraction and think-token removal log.debug("[%s] Raw LLM response: %s", req_id, truncate(response.content)) @@ -612,7 +721,7 @@ async def _do_generate( log.info("[%s] Tool call blocked: %s", req_id, tool_call.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return {"role": "assistant", "content": REFUSAL_MESSAGE} + return _finalize_refusal(has_generation_response) # Output rails check the final answer, not reasoning traces. # Reasoning is re-attached as tags only below so reasoning intentionally bypasses output @@ -627,10 +736,14 @@ async def _do_generate( log.info("[%s] Output blocked: %s", req_id, output_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return {"role": "assistant", "content": REFUSAL_MESSAGE} + return _finalize_refusal(has_generation_response) + + if has_generation_response: + return _build_generation_response(response_text, reasoning_content, response) - # TODO: Support returning GenerationResponse `reasoning_content` to match LLMRails - # For now, embed the reasoning on the content with think-tags + # Bare return path: reasoning is delivered inline as a prefix + # (LLMRails legacy shape). The structured path above instead puts it in the + # reasoning_content field and keeps the message content clean. if reasoning_content: response_text = f"{reasoning_content}\n" + response_text diff --git a/tests/guardrails/test_iorails.py b/tests/guardrails/test_iorails.py index ce287a47b0..14071362cc 100644 --- a/tests/guardrails/test_iorails.py +++ b/tests/guardrails/test_iorails.py @@ -26,7 +26,7 @@ from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails from nemoguardrails.guardrails.model_engine import ModelEngine from nemoguardrails.rails.llm.config import RailsConfig -from nemoguardrails.rails.llm.options import GenerationOptions +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse from nemoguardrails.types import LLMResponse, LLMResponseChunk, ToolCall, ToolCallFunction from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG, NEMOGUARDS_CONFIG @@ -90,7 +90,7 @@ async def test_safe_input_and_output(self, iorails): @pytest.mark.asyncio async def test_safe_input_and_output_with_generation_options(self, iorails): - """Returns LLM response when both input and output rails pass.""" + """Passing options returns a GenerationResponse wrapping the LLM response.""" messages = [{"role": "user", "content": "hi"}] llm_response = "Hello from LLM" @@ -103,7 +103,8 @@ async def test_safe_input_and_output_with_generation_options(self, iorails): result = await iorails.generate_async(messages, options=options) - assert result == {"role": "assistant", "content": llm_response} + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": llm_response}] iorails.rails_manager.is_input_safe.assert_called_once_with(messages, enabled=True) iorails.engine_registry.model_call.assert_called_once_with("main", messages, **llm_params) iorails.rails_manager.is_output_safe.assert_called_once_with(messages, llm_response, enabled=True) diff --git a/tests/guardrails/test_iorails_generation_response.py b/tests/guardrails/test_iorails_generation_response.py new file mode 100644 index 0000000000..4c7c7dbe76 --- /dev/null +++ b/tests/guardrails/test_iorails_generation_response.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Structured GenerationResponse return from IORails, matching LLMRails. + +When ``options`` is supplied, ``generate_async``/``generate`` return a +``GenerationResponse`` instead of a bare ``LLMMessage`` dict, mirroring LLMRails' +``if gen_options:`` branch. When ``options`` is absent the bare dict is returned +unchanged. The structured path populates ``response``, ``reasoning_content``, +``tool_calls`` (as ``ToolCall.to_dict()`` with dict arguments), and +``llm_metadata`` (main-call ``provider_metadata`` plus a ``usage`` sub-key); +``llm_output`` is always ``None`` (parity with LLMRails' unwired ``raw_response``); +``log`` is deferred; and ``output_vars``/``state``/``log`` request options raise. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio + +from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse +from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction, UsageInfo +from tests.guardrails.async_helpers import started_iorails +from tests.guardrails.test_data import NEMOGUARDS_CONFIG + + +@pytest_asyncio.fixture +async def iorails(): + """Started IORails instance with worker-queue teardown after each test.""" + async with started_iorails(NEMOGUARDS_CONFIG) as iorails: + yield iorails + + +@pytest.fixture +def iorails_sync(): + """Unstarted IORails instance for driving the synchronous ``generate`` path.""" + return IORails(RailsConfig.from_content(config=NEMOGUARDS_CONFIG)) + + +def _stub_safe_rails(iorails: IORails) -> None: + """Default-safe input, output, and tool-call rails so tests focus on the LLM response.""" + 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.rails_manager.are_tool_calls_safe = AsyncMock(return_value=RailResult(is_safe=True)) + + +def _stub_model(iorails: IORails, response: LLMResponse) -> None: + """Make the main-model call return a fixed structured LLMResponse.""" + iorails.engine_registry.model_call = AsyncMock(return_value=response) + + +_USER = [{"role": "user", "content": "hi"}] + + +class TestStructuredResponseTrigger: + """``options`` presence decides GenerationResponse vs. bare dict.""" + + @pytest.mark.asyncio + async def test_options_dict_returns_generation_response(self, iorails): + """A dict ``options`` argument switches the return type to GenerationResponse.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + result = await iorails.generate_async(_USER, options={"llm_params": {"temperature": 0.5}}) + + assert isinstance(result, GenerationResponse) + + @pytest.mark.asyncio + async def test_generation_options_returns_generation_response(self, iorails): + """A GenerationOptions instance also switches the return type.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + result = await iorails.generate_async(_USER, options=GenerationOptions()) + + assert isinstance(result, GenerationResponse) + + @pytest.mark.asyncio + async def test_no_generation_options_returns_messages(self, iorails): + """Without ``options`` the return stays the bare assistant-message dict.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + result = await iorails.generate_async(_USER) + + assert not isinstance(result, GenerationResponse) + assert result == {"role": "assistant", "content": "Hello"} + + +class TestResponseField: + """The ``response`` field wraps the assistant message in a one-element list.""" + + @pytest.mark.asyncio + async def test_plain_content_wrapped_in_list(self, iorails): + """``response`` is ``[{"role":"assistant","content": text}]``.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello there")) + + result = await iorails.generate_async(_USER, options={}) + + assert result.response == [{"role": "assistant", "content": "Hello there"}] + + +class TestReasoningContent: + """Reasoning goes to ``reasoning_content`` with clean content (no inline ).""" + + @pytest.mark.asyncio + async def test_native_reasoning_in_field_content_clean(self, iorails): + """Provider ``reasoning`` populates ``reasoning_content``; ``response`` content has no prefix.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello", reasoning="thinking step")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.reasoning_content == "thinking step" + assert result.response == [{"role": "assistant", "content": "Hello"}] + iorails.rails_manager.is_output_safe.assert_called_once_with(_USER, "Hello", enabled=True) + + @pytest.mark.asyncio + async def test_inline_think_tags_extracted_to_field(self, iorails): + """Inline tags are stripped into ``reasoning_content`` and never reach output rails.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="thinking stepHello")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.reasoning_content == "thinking step" + assert result.response == [{"role": "assistant", "content": "Hello"}] + iorails.rails_manager.is_output_safe.assert_called_once_with(_USER, "Hello", enabled=True) + + @pytest.mark.asyncio + async def test_no_reasoning_field_is_none(self, iorails): + """Absent reasoning leaves ``reasoning_content`` as None.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="plain answer")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.reasoning_content is None + assert result.response == [{"role": "assistant", "content": "plain answer"}] + + +class TestToolCalls: + """``tool_calls`` use the LLMRails ``ToolCall.to_dict()`` shape (dict arguments).""" + + @pytest.mark.asyncio + async def test_tool_calls_serialized_with_dict_arguments(self, iorails): + """``tool_calls`` is a list of ``to_dict()`` entries whose ``arguments`` stay a dict, not a JSON string.""" + _stub_safe_rails(iorails) + tool_call = ToolCall( + id="call_1", + type="function", + function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), + ) + _stub_model(iorails, LLMResponse(content="", tool_calls=[tool_call])) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.tool_calls == [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "SF"}}, + } + ] + + @pytest.mark.asyncio + async def test_response_message_has_no_tool_calls_key(self, iorails): + """In the structured path tool calls live only in the top-level field, not on the message.""" + _stub_safe_rails(iorails) + tool_call = ToolCall( + id="call_1", + type="function", + function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), + ) + _stub_model(iorails, LLMResponse(content="", tool_calls=[tool_call])) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": ""}] + assert "tool_calls" not in result.response[0] + + @pytest.mark.asyncio + async def test_no_tool_calls_field_is_none(self, iorails): + """A text-only response leaves ``tool_calls`` as None.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.tool_calls is None + + +class TestLLMMetadata: + """``llm_metadata`` carries main-call ``provider_metadata`` plus a ``usage`` sub-key.""" + + @pytest.mark.asyncio + async def test_provider_metadata_and_usage_merged(self, iorails): + """provider_metadata is surfaced and structured usage is added under ``usage``.""" + _stub_safe_rails(iorails) + _stub_model( + iorails, + LLMResponse( + content="Hello", + provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}, + usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15), + ), + ) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.llm_metadata == { + "response_headers": {"nvcf-status": "fulfilled"}, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + @pytest.mark.asyncio + async def test_usage_only_when_no_provider_metadata(self, iorails): + """With usage but no provider_metadata, ``llm_metadata`` holds just the ``usage`` sub-key.""" + _stub_safe_rails(iorails) + _stub_model( + iorails, + LLMResponse(content="Hello", usage=UsageInfo(input_tokens=3, output_tokens=4, total_tokens=7)), + ) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.llm_metadata == {"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}} + + @pytest.mark.asyncio + async def test_usage_includes_reasoning_and_cached_tokens_when_present(self, iorails): + """Non-None ``reasoning_tokens``/``cached_tokens`` appear in the ``usage`` sub-key.""" + _stub_safe_rails(iorails) + _stub_model( + iorails, + LLMResponse( + content="Hello", + usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15, reasoning_tokens=3, cached_tokens=2), + ), + ) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.llm_metadata == { + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "reasoning_tokens": 3, + "cached_tokens": 2, + } + } + + @pytest.mark.asyncio + async def test_no_metadata_or_usage_is_none(self, iorails): + """No provider_metadata and no usage leaves ``llm_metadata`` as None.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.llm_metadata is None + + +class TestLLMOutput: + """``llm_output`` is always None, matching LLMRails' unwired ``raw_response``.""" + + @pytest.mark.asyncio + async def test_llm_output_none_even_when_requested(self, iorails): + """``options={"llm_output": True}`` is accepted but the field stays None.""" + _stub_safe_rails(iorails) + _stub_model( + iorails, + LLMResponse(content="Hello", provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}), + ) + + result = await iorails.generate_async(_USER, options={"llm_output": True}) + + assert isinstance(result, GenerationResponse) + assert result.llm_output is None + + +class TestUnsupportedOptionGuards: + """Colang-coupled options IORails cannot honor raise ValueError.""" + + @pytest.mark.asyncio + async def test_output_vars_true_raises(self, iorails): + """``output_vars=True`` requests Colang context IORails has no access to.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + with pytest.raises(ValueError): + await iorails.generate_async(_USER, options={"output_vars": True}) + + @pytest.mark.asyncio + async def test_output_vars_list_raises(self, iorails): + """A list of ``output_vars`` also raises.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + with pytest.raises(ValueError): + await iorails.generate_async(_USER, options={"output_vars": ["relevant_chunks"]}) + + @pytest.mark.asyncio + async def test_log_request_flag_raises(self, iorails): + """Requesting any ``log`` detail raises NotImplementedError while ``log`` is deferred.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + with pytest.raises(NotImplementedError): + await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) + + @pytest.mark.asyncio + async def test_state_argument_raises(self, iorails): + """A ``state`` argument is unsupported for the stateless IORails engine.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello")) + + with pytest.raises(ValueError): + await iorails.generate_async(_USER, options={}, state={"conversation": []}) + + +class TestBlockedStructuredResponse: + """A blocked request returns the refusal in ``response`` with the other fields empty.""" + + @pytest.mark.asyncio + async def test_input_block_returns_refusal_response(self, iorails): + """Input-rail block yields a GenerationResponse whose response is the refusal message.""" + iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=False, reason="unsafe")) + iorails.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + _stub_model(iorails, LLMResponse(content="unused")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] + assert result.tool_calls is None + assert result.reasoning_content is None + assert result.llm_metadata is None + + @pytest.mark.asyncio + async def test_output_block_returns_refusal_response(self, iorails): + """Output-rail block yields a GenerationResponse whose response is the refusal message.""" + 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=False, reason="unsafe")) + _stub_model(iorails, LLMResponse(content="bad answer")) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] + + +class TestBarePathUnchanged: + """The optionless bare-dict path keeps its existing behavior.""" + + @pytest.mark.asyncio + async def test_bare_path_inlines_reasoning_prefix(self, iorails): + """Without ``options`` reasoning is still delivered inline as a prefix.""" + _stub_safe_rails(iorails) + _stub_model(iorails, LLMResponse(content="Hello", reasoning="thinking step")) + + result = await iorails.generate_async(_USER) + + assert result == {"role": "assistant", "content": "thinking step\nHello"} + + +class TestSyncGenerateStructured: + """The synchronous ``generate`` mirrors the async structured return.""" + + def test_sync_generate_with_options_returns_generation_response(self, iorails_sync): + """``generate(options=...)`` returns a GenerationResponse from the ephemeral engine.""" + iorails_sync.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails_sync.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails_sync.engine_registry.model_call = AsyncMock(return_value=LLMResponse(content="Hello")) + + with patch("nemoguardrails.guardrails.iorails.IORails", return_value=iorails_sync): + result = iorails_sync.generate(_USER, options={}) + + assert isinstance(result, GenerationResponse) diff --git a/tests/guardrails/test_tool_rails_iorails.py b/tests/guardrails/test_tool_rails_iorails.py index d8593c6c56..7a8b2c0791 100644 --- a/tests/guardrails/test_tool_rails_iorails.py +++ b/tests/guardrails/test_tool_rails_iorails.py @@ -309,20 +309,20 @@ class TestNonStreamingToolCalls: async def test_allowed_tool_call_passes(self, iorails): _inject_json_response(iorails, _tool_call_payload("get_weather", '{"city": "Paris"}')) result = await iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result["tool_calls"][0]["function"]["name"] == "get_weather" + assert result.tool_calls[0]["function"]["name"] == "get_weather" @pytest.mark.asyncio async def test_undeclared_tool_call_blocked(self, iorails): _inject_json_response(iorails, _tool_call_payload("rm_rf", "{}")) result = await iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result == {"role": "assistant", "content": REFUSAL_MESSAGE} + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] @pytest.mark.asyncio async def test_invalid_arguments_blocked(self, iorails): # Missing the required "city" argument violates the declared schema. _inject_json_response(iorails, _tool_call_payload("get_weather", "{}")) result = await iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result == {"role": "assistant", "content": REFUSAL_MESSAGE} + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] class TestSpeculativeToolCalls: @@ -337,13 +337,13 @@ async def speculative_iorails(self): async def test_undeclared_tool_call_blocked(self, speculative_iorails): _inject_json_response(speculative_iorails, _tool_call_payload("rm_rf", "{}")) result = await speculative_iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result == {"role": "assistant", "content": REFUSAL_MESSAGE} + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] @pytest.mark.asyncio async def test_allowed_tool_call_passes(self, speculative_iorails): _inject_json_response(speculative_iorails, _tool_call_payload("get_weather", '{"city": "Paris"}')) result = await speculative_iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result["tool_calls"][0]["function"]["name"] == "get_weather" + assert result.tool_calls[0]["function"]["name"] == "get_weather" @pytest.mark.asyncio async def test_input_toggle_forwarded_to_input_rails(self, speculative_iorails): @@ -474,7 +474,7 @@ async def test_tool_output_disabled_skips_tool_call_rail(self, iorails): _inject_json_response(iorails, _tool_call_payload("rm_rf", "{}")) options = {"llm_params": LLM_PARAMS, "rails": {"tool_output": False}} result = await iorails.generate_async(MESSAGES, options=options) - assert result["tool_calls"][0]["function"]["name"] == "rm_rf" + assert result.tool_calls[0]["function"]["name"] == "rm_rf" @pytest.mark.asyncio async def test_tool_input_disabled_skips_tool_result_rail(self, iorails): @@ -484,7 +484,7 @@ async def test_tool_input_disabled_skips_tool_result_rail(self, iorails): make_tool_conversation(result_call_id="call_999"), options={"rails": {"tool_input": False}}, ) - assert result == {"role": "assistant", "content": "ok"} + assert result.response == [{"role": "assistant", "content": "ok"}] @pytest.mark.asyncio async def test_input_toggle_forwarded_to_input_rails(self, iorails): @@ -531,7 +531,7 @@ async def test_duplicate_tool_definitions_block(self, iorails): _inject_json_response(iorails, _tool_call_payload("get_weather", '{"city": "Paris"}')) dup_params = {"tools": [WEATHER_TOOL, WEATHER_TOOL]} result = await iorails.generate_async(MESSAGES, options={"llm_params": dup_params}) - assert result == {"role": "assistant", "content": REFUSAL_MESSAGE} + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] class TestToolRailBlockMetrics: @@ -551,7 +551,7 @@ async def test_nonstream_tool_call_block_records_output_metric(self, iorails): _inject_json_response(iorails, _tool_call_payload("rm_rf", "{}")) with patch("nemoguardrails.guardrails.iorails.record_request_blocked") as record_blocked: result = await iorails.generate_async(MESSAGES, options={"llm_params": LLM_PARAMS}) - assert result == {"role": "assistant", "content": REFUSAL_MESSAGE} + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] record_blocked.assert_called_once_with(RailDirection.OUTPUT) @pytest.mark.asyncio From 054dca75df7a6db169909b3713279b25e5066881 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:12:33 -0500 Subject: [PATCH 02/11] Promote options to a named argument rather than smuggled in through kwargs --- nemoguardrails/base_guardrails.py | 9 +++-- nemoguardrails/guardrails/guardrails.py | 44 +++++++++++++++++----- nemoguardrails/guardrails/iorails.py | 49 ++++++++++++++++++------- tests/guardrails/test_guardrails.py | 19 +++++----- 4 files changed, 86 insertions(+), 35 deletions(-) diff --git a/nemoguardrails/base_guardrails.py b/nemoguardrails/base_guardrails.py index 359a03f16e..599e65649a 100644 --- a/nemoguardrails/base_guardrails.py +++ b/nemoguardrails/base_guardrails.py @@ -27,9 +27,10 @@ """ from abc import ABC, abstractmethod -from typing import Any, AsyncIterator +from typing import Any, AsyncIterator, Optional, Union from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.options import GenerationOptions class BaseGuardrails(ABC): @@ -44,12 +45,14 @@ class BaseGuardrails(ABC): config: RailsConfig @abstractmethod - def generate(self, *args: Any, **kwargs: Any) -> Any: + def generate(self, *args: Any, options: Optional[Union[dict, GenerationOptions]] = None, **kwargs: Any) -> Any: """Generate an LLM response synchronously with guardrails applied.""" ... @abstractmethod - async def generate_async(self, *args: Any, **kwargs: Any) -> Any: + async def generate_async( + self, *args: Any, options: Optional[Union[dict, GenerationOptions]] = None, **kwargs: Any + ) -> Any: """Generate an LLM response asynchronously with guardrails applied.""" ... diff --git a/nemoguardrails/guardrails/guardrails.py b/nemoguardrails/guardrails/guardrails.py index 9579f55333..3b31f024d3 100644 --- a/nemoguardrails/guardrails/guardrails.py +++ b/nemoguardrails/guardrails/guardrails.py @@ -38,7 +38,7 @@ from nemoguardrails.logging.explain import ExplainInfo from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.rails.llm.llmrails import LLMRails -from nemoguardrails.rails.llm.options import GenerationResponse, RailsResult, RailType +from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse, RailsResult, RailType from nemoguardrails.types import LLMModel log = logging.getLogger(__name__) @@ -205,35 +205,61 @@ async def _ensure_started(self) -> None: await self.startup() def generate( - self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> Union[str, dict, GenerationResponse, Tuple[dict, dict]]: """Generate an LLM response synchronously with guardrails applied. Supported in both IORails and LLMRails """ generate_messages = self._convert_to_messages(prompt, messages) - return self.rails_engine.generate(messages=generate_messages, **kwargs) + return self.rails_engine.generate(messages=generate_messages, options=options, **kwargs) @overload - async def generate_async(self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs) -> str: ... + async def generate_async( + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, + ) -> str: ... @overload async def generate_async( - self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> dict: ... @overload async def generate_async( - self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> GenerationResponse: ... @overload async def generate_async( - self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> tuple[dict, dict]: ... async def generate_async( - self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs + self, + prompt: str | None = None, + messages: LLMMessages | None = None, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> str | dict | GenerationResponse | tuple[dict, dict]: """Generate an LLM response asynchronously with guardrails applied. Supported by both LLMRails and IORails @@ -241,7 +267,7 @@ async def generate_async( await self._ensure_started() generate_messages = self._convert_to_messages(prompt, messages) - return await self.rails_engine.generate_async(messages=generate_messages, **kwargs) + return await self.rails_engine.generate_async(messages=generate_messages, options=options, **kwargs) def stream_async( self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 9a92eb02e7..15c0bc0c3b 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -569,7 +569,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): """Context manager (used for testing rather than long-lived instance)""" await self.stop() - def generate(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: + def generate( + self, + messages: LLMMessages, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, + ) -> Union[LLMMessage, GenerationResponse]: """Synchronous version of generate_async. Telemetry is disabled for the ephemeral IORails object used for @@ -589,11 +594,16 @@ 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, **kwargs) + return await iorails_engine.generate_async(messages, options=options, **kwargs) return asyncio.run(_run_sync_iorails()) - async def generate_async(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: + async def generate_async( + self, + messages: LLMMessages, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, + ) -> Union[LLMMessage, GenerationResponse]: """Public entry: submit the request to the internal work queue. The queue enforces non-streaming concurrency limits @@ -613,13 +623,18 @@ async def generate_async(self, messages: LLMMessages, **kwargs) -> Union[LLMMess 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, **kwargs) + return await self._generate_async_queue.submit(self._run_generate, messages, options=options, **kwargs) except asyncio.QueueFull: if self._metrics_enabled: record_nonstream_rejected() raise - async def _run_generate(self, messages: LLMMessages, **kwargs) -> Union[LLMMessage, GenerationResponse]: + async def _run_generate( + self, + messages: LLMMessages, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, + ) -> Union[LLMMessage, GenerationResponse]: """Runs inside a queue worker task. Wraps the pipeline in ``traced_request`` so each request gets its own span + request ID, then delegates to ``_do_generate`` for the actual input rails → @@ -630,7 +645,7 @@ async def _run_generate(self, messages: LLMMessages, **kwargs) -> Union[LLMMessa with traced_request(tracer) as (request_span, req_id): t0 = time.monotonic() try: - result = await self._do_generate(messages, req_id, request_span, **kwargs) + result = await self._do_generate(messages, req_id, request_span, options=options, **kwargs) except Exception: elapsed_ms = (time.monotonic() - t0) * 1000 log.error("[%s] generate_async failed time=%.1fms", req_id, elapsed_ms, exc_info=True) @@ -664,17 +679,23 @@ def _guardrails_violation_payload(message: str, param: str) -> str: ) async def _do_generate( - self, messages: LLMMessages, req_id: str, request_span: Optional["Span"] = None, **kwargs + self, + messages: LLMMessages, + req_id: str, + request_span: Optional["Span"] = None, + *, + options: Optional[Union[dict, GenerationOptions]] = None, + **kwargs, ) -> Union[LLMMessage, GenerationResponse]: """Core pipeline: tool-result rails -> input rails -> LLM call -> tool-call + output rails.""" log.info("[%s] generate_async called", req_id) log.debug("[%s] generate_async messages=%s", req_id, truncate(messages)) - options = _coerce_generation_options(kwargs.get("options")) + options = _coerce_generation_options(options) _raise_on_unsupported_options(options, kwargs.get("state")) # When options are supplied we return a structured GenerationResponse # (mirroring LLMRails' `if gen_options:` branch); otherwise a bare LLMMessage. - has_generation_response = options is not None + has_generation_options = options is not None # Pass llm_params (including tool definitions) unchanged to the LLM call. llm_kwargs = options.llm_params if (options and options.llm_params) else {} input_enabled = options.rails.input if options else True @@ -690,7 +711,7 @@ async def _do_generate( log.info("[%s] Tool result blocked: %s", req_id, tool_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.INPUT) - return _finalize_refusal(has_generation_response) + return _finalize_refusal(has_generation_options) if self._speculative_generation: response = await self._do_generate_speculative( @@ -700,7 +721,7 @@ async def _do_generate( response = await self._do_generate_sequential(messages, req_id, llm_kwargs, input_enabled=input_enabled) if response is None: - return _finalize_refusal(has_generation_response) + return _finalize_refusal(has_generation_options) # Log raw content before reasoning extraction and think-token removal log.debug("[%s] Raw LLM response: %s", req_id, truncate(response.content)) @@ -721,7 +742,7 @@ async def _do_generate( log.info("[%s] Tool call blocked: %s", req_id, tool_call.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return _finalize_refusal(has_generation_response) + return _finalize_refusal(has_generation_options) # Output rails check the final answer, not reasoning traces. # Reasoning is re-attached as tags only below so reasoning intentionally bypasses output @@ -736,9 +757,9 @@ async def _do_generate( log.info("[%s] Output blocked: %s", req_id, output_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return _finalize_refusal(has_generation_response) + return _finalize_refusal(has_generation_options) - if has_generation_response: + if has_generation_options: return _build_generation_response(response_text, reasoning_content, response) # Bare return path: reasoning is delivered inline as a prefix diff --git a/tests/guardrails/test_guardrails.py b/tests/guardrails/test_guardrails.py index d573fb0daa..40325b6b94 100644 --- a/tests/guardrails/test_guardrails.py +++ b/tests/guardrails/test_guardrails.py @@ -123,8 +123,8 @@ async def mock_stream(): guardrails.update_llm(mock_new_llm) # Verify all calls went to LLMRails - guardrails.rails_engine.generate.assert_called_once_with(messages=messages) - guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages) + guardrails.rails_engine.generate.assert_called_once_with(messages=messages, options=None) + guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages, options=None) guardrails.rails_engine.stream_async.assert_called_once_with(messages=messages) guardrails.rails_engine.explain.assert_called_once() guardrails.rails_engine.update_llm.assert_called_once_with(mock_new_llm) @@ -178,8 +178,8 @@ 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(messages=messages) - guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages) + guardrails.rails_engine.generate.assert_called_once_with(messages=messages, options=None) + guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages, options=None) guardrails.rails_engine.stream_async.assert_called_once_with( messages=messages, options=None, @@ -257,8 +257,8 @@ async def mock_stream(): guardrails.update_llm(mock_new_llm) # Verify all calls went to LLMRails - guardrails.rails_engine.generate.assert_called_once_with(messages=messages) - guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages) + guardrails.rails_engine.generate.assert_called_once_with(messages=messages, options=None) + guardrails.rails_engine.generate_async.assert_called_once_with(messages=messages, options=None) guardrails.rails_engine.stream_async.assert_called_once_with(messages=messages) guardrails.rails_engine.explain.assert_called_once() guardrails.rails_engine.update_llm.assert_called_once_with(mock_new_llm) @@ -649,7 +649,7 @@ async def test_generate_async_with_string_prompt(self, mock_llmrails_class, _nem # Verify generate_async was called with correct messages expected_messages = [{"role": "user", "content": "Hello async!"}] - mock_llmrails_instance.generate_async.assert_awaited_once_with(messages=expected_messages) + mock_llmrails_instance.generate_async.assert_awaited_once_with(messages=expected_messages, options=None) assert result == "Async response" @pytest.mark.asyncio @@ -668,7 +668,7 @@ async def test_generate_async_with_messages(self, mock_llmrails_class, _nemoguar ] result = await guardrails.generate_async(messages=messages) - mock_llmrails_instance.generate_async.assert_awaited_once_with(messages=messages) + mock_llmrails_instance.generate_async.assert_awaited_once_with(messages=messages, options=None) assert result == "Async conversation response" @pytest.mark.asyncio @@ -685,7 +685,7 @@ async def test_generate_async_with_kwargs(self, mock_llmrails_class, _nemoguards # Verify kwargs were passed through expected_messages = [{"role": "user", "content": "Test"}] mock_llmrails_instance.generate_async.assert_awaited_once_with( - messages=expected_messages, temperature=0.5, top_p=0.9 + messages=expected_messages, options=None, temperature=0.5, top_p=0.9 ) assert result == "Response" @@ -901,6 +901,7 @@ def test_generate_with_additional_parameters(self, mock_llmrails_class, _nemogua expected_messages = [{"role": "user", "content": "Test"}] mock_llmrails_instance.generate.assert_called_once_with( messages=expected_messages, + options=None, temperature=0.7, max_tokens=100, top_p=0.9, From b72afc0f1bea04223705cd7caaf116cd17ca39ea Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 16:57:10 -0500 Subject: [PATCH 03/11] Add functionality for LLMRails non-Colang GenerationResponse parity --- .../actions/content_safety_action.py | 13 +- nemoguardrails/guardrails/guardrails_types.py | 47 +++- nemoguardrails/guardrails/iorails.py | 258 ++++++++++++++---- nemoguardrails/guardrails/rail_action.py | 55 +++- nemoguardrails/guardrails/rails_manager.py | 43 ++- .../guardrails/test_iorails_generation_log.py | 222 +++++++++++++++ .../test_iorails_generation_log_capture.py | 138 ++++++++++ .../test_iorails_generation_response.py | 42 +-- 8 files changed, 715 insertions(+), 103 deletions(-) create mode 100644 tests/guardrails/test_iorails_generation_log.py create mode 100644 tests/guardrails/test_iorails_generation_log_capture.py diff --git a/nemoguardrails/guardrails/actions/content_safety_action.py b/nemoguardrails/guardrails/actions/content_safety_action.py index eb9a15e034..ce0229dc89 100644 --- a/nemoguardrails/guardrails/actions/content_safety_action.py +++ b/nemoguardrails/guardrails/actions/content_safety_action.py @@ -121,10 +121,13 @@ def _content_safety_to_rail_result(parsed: object) -> RailResult: """ if isinstance(parsed, (list, tuple)): if parsed and parsed[0] is True: - return RailResult(is_safe=True) + return RailResult(is_safe=True, return_value={"allowed": True, "policy_violations": []}) if parsed and parsed[0] is False: - if len(parsed) > 1: - categories = ", ".join(str(c) for c in parsed[1:]) - return RailResult(is_safe=False, reason=f"Safety categories: {categories}") - return RailResult(is_safe=False, reason="Unknown") + violations = [str(c) for c in parsed[1:]] + verdict = {"allowed": False, "policy_violations": violations} + if violations: + return RailResult( + is_safe=False, reason=f"Safety categories: {', '.join(violations)}", return_value=verdict + ) + return RailResult(is_safe=False, reason="Unknown", return_value=verdict) raise RuntimeError(f"Unexpected content safety parse result: {parsed}") diff --git a/nemoguardrails/guardrails/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index 546f66361c..f5e613a925 100644 --- a/nemoguardrails/guardrails/guardrails_types.py +++ b/nemoguardrails/guardrails/guardrails_types.py @@ -16,9 +16,11 @@ import secrets from contextvars import ContextVar, Token -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum -from typing import Any, TypeAlias +from typing import Any, Optional, TypeAlias + +from nemoguardrails.types import UsageInfo # LLMMessage can contain role/content, plus optional tool_calls / tool_call_id / name; content may be None LLMMessage: TypeAlias = dict[str, Any] @@ -32,13 +34,52 @@ class RailDirection(Enum): OUTPUT = "Output" +@dataclass(frozen=True, slots=True) +class RailCallRecord: + """One rail's execution record, carried on RailResult for GenerationLog synthesis. + + Captures what a single rail did — its verdict and the (at most one) model call it + made — as engine-neutral data. IORails maps a ``RailCallRecord`` to an + ``ActivatedRail`` (with a single synthetic ``ExecutedAction`` and ``LLMCallInfo``); + the raw ``usage``/timing is kept here so this module stays free of the pydantic + ``GenerationLog`` types. Tool rails that make no model call leave ``usage`` None. + """ + + flow: str + rail_type: str + is_safe: bool + action_name: Optional[str] = None + return_value: Any = None + task: Optional[str] = None + usage: Optional[UsageInfo] = None + llm_model_name: Optional[str] = None + llm_provider_name: Optional[str] = None + started_at: Optional[float] = None + finished_at: Optional[float] = None + duration: Optional[float] = None + + @dataclass(frozen=True, slots=True) class RailResult: - """Result of a rail safety check.""" + """Result of a rail safety check. + + ``records`` carries the per-rail execution records for every rail that ran in this + check (not just the blocking one), so IORails can synthesize a ``GenerationLog``. + It is empty unless log collection is active. ``return_value`` is the rail's + structured verdict (e.g. ``{"allowed": ..., "policy_violations": [...]}``) when the + action supplies one, used as the log's ``ExecutedAction.return_value``. + + ``records`` and ``return_value`` are log-capture metadata, not part of the safety + verdict, so they are excluded from equality and hashing (``compare=False``) — two + results with the same ``is_safe``/``reason``/``triggered_rail`` compare equal + regardless of captured log data. + """ is_safe: bool reason: str | None = None triggered_rail: str | None = None + records: tuple[RailCallRecord, ...] = field(default=(), compare=False) + return_value: Any = field(default=None, compare=False) # Default max character length for truncate(). Used to keep DEBUG log lines short. diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 15c0bc0c3b..21ef59b84a 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -37,6 +37,7 @@ from nemoguardrails.guardrails.guardrails_types import ( LLMMessage, LLMMessages, + RailCallRecord, RailDirection, get_request_id, truncate, @@ -60,19 +61,24 @@ traced_request, ) from nemoguardrails.llm.taskmanager import LLMTaskManager +from nemoguardrails.logging.explain import LLMCallInfo from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop from nemoguardrails.rails.llm.buffer import get_buffer_strategy from nemoguardrails.rails.llm.config import RailsConfig, _get_flow_name from nemoguardrails.rails.llm.options import ( + ActivatedRail, + ExecutedAction, + GenerationLog, GenerationOptions, GenerationResponse, + GenerationStats, RailsResult, RailStatus, RailType, ) from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler from nemoguardrails.tracing.constants import GuardrailsAttributes -from nemoguardrails.types import LLMModel, LLMResponse, ToolCall, UsageInfo +from nemoguardrails.types import LLMModel, LLMResponse, ToolCall if TYPE_CHECKING: from opentelemetry.trace import Span @@ -176,40 +182,137 @@ def _build_assistant_message(content: str, tool_calls: Optional[list[ToolCall]]) } -def _usage_to_dict(usage: UsageInfo) -> dict: - """Serialize ``UsageInfo`` token counts to a dict, dropping ``None`` values. +def _build_llm_metadata(response: LLMResponse) -> Optional[dict]: + """Return the main call's ``provider_metadata`` verbatim (LLMRails mirror). - ``input_tokens``/``output_tokens``/``total_tokens`` default to 0 and are always - kept; ``reasoning_tokens``/``cached_tokens`` default to ``None`` and are omitted - unless the provider reported them. + A pure passthrough — token usage lives in ``log`` (``log.llm_calls`` / ``log.stats``), + not here, matching LLMRails. Returns ``None`` when the main call carried no metadata. """ - fields = { - "input_tokens": usage.input_tokens, - "output_tokens": usage.output_tokens, - "total_tokens": usage.total_tokens, - "reasoning_tokens": usage.reasoning_tokens, - "cached_tokens": usage.cached_tokens, - } - return {key: value for key, value in fields.items() if value is not None} + return dict(response.provider_metadata or {}) or None -def _build_llm_metadata(response: LLMResponse) -> Optional[dict]: - """Assemble ``llm_metadata`` from the main call's ``provider_metadata`` plus usage. +def _make_generation_record( + response: LLMResponse, started_at: Optional[float], finished_at: Optional[float] +) -> RailCallRecord: + """Represent the main generation call as a ``RailCallRecord`` (rail_type "generation"). + + Unifies the main call with the rail calls so the log builder can treat every LLM call + the same way. + """ + duration = (finished_at - started_at) if (started_at is not None and finished_at is not None) else None + return RailCallRecord( + flow="generation", + rail_type="generation", + is_safe=True, + action_name="generate_bot_message", + task="general", + usage=response.usage, + llm_model_name=response.model, + started_at=started_at, + finished_at=finished_at, + duration=duration, + ) + + +def _record_has_llm_call(record: RailCallRecord) -> bool: + """True when a record reflects a model call (has usage or a model name).""" + return record.usage is not None or record.llm_model_name is not None + + +def _call_info(record: RailCallRecord) -> LLMCallInfo: + """Map a ``RailCallRecord`` to a ``LLMCallInfo`` for the log.""" + usage = record.usage + return LLMCallInfo( + task=record.task, + duration=record.duration, + total_tokens=usage.total_tokens if usage else None, + prompt_tokens=usage.input_tokens if usage else None, + completion_tokens=usage.output_tokens if usage else None, + started_at=record.started_at, + finished_at=record.finished_at, + llm_model_name=record.llm_model_name or "unknown", + llm_provider_name=record.llm_provider_name or "unknown", + ) + + +def _activated_rail(record: RailCallRecord) -> ActivatedRail: + """Map a ``RailCallRecord`` to an ``ActivatedRail`` with one synthetic ``ExecutedAction``. + + IORails runs one model-backed check per rail (not a Colang action chain), so each rail + gets a single ``ExecutedAction`` carrying the rail's structured verdict as + ``return_value`` and its (at most one) LLM call. + """ + action = ExecutedAction( + action_name=record.action_name or record.flow, + action_params={}, + return_value=record.return_value, + llm_calls=[_call_info(record)] if _record_has_llm_call(record) else [], + started_at=record.started_at, + finished_at=record.finished_at, + duration=record.duration, + ) + return ActivatedRail( + type=record.rail_type, + name=record.flow, + executed_actions=[action], + stop=not record.is_safe, + started_at=record.started_at, + finished_at=record.finished_at, + duration=record.duration, + ) + + +def _build_generation_stats( + records: list[RailCallRecord], call_records: list[RailCallRecord], total_duration: Optional[float] +) -> GenerationStats: + """Aggregate per-phase durations and token totals across every recorded call.""" + + def _phase_duration(*rail_types: str) -> Optional[float]: + total = sum((r.duration or 0.0) for r in records if r.rail_type in rail_types) + return total or None + + return GenerationStats( + input_rails_duration=_phase_duration("input", "tool_input"), + output_rails_duration=_phase_duration("output", "tool_output"), + generation_rails_duration=_phase_duration("generation"), + total_duration=total_duration, + llm_calls_duration=sum((r.duration or 0.0) for r in call_records) or None, + llm_calls_count=len(call_records), + llm_calls_total_prompt_tokens=sum(r.usage.input_tokens for r in call_records if r.usage), + llm_calls_total_completion_tokens=sum(r.usage.output_tokens for r in call_records if r.usage), + llm_calls_total_tokens=sum(r.usage.total_tokens for r in call_records if r.usage), + ) + - Mirrors LLMRails' ``llm_metadata`` (the provider metadata blob) and adds a - structured ``usage`` sub-key with token counts that LLMRails does not expose - outside its ``log``. Returns ``None`` when neither is present. +def _build_generation_log( + records: list[RailCallRecord], options: Optional[GenerationOptions], total_duration: Optional[float] +) -> Optional[GenerationLog]: + """Synthesize a ``GenerationLog`` from collected rail + generation records. + + Returns ``None`` unless ``options.log`` requests ``activated_rails`` or ``llm_calls``. + ``stats`` is always included when either is requested (matching LLMRails); the + per-rail ``activated_rails`` and flat ``llm_calls`` are gated on their own flags. + ``internal_events`` / ``colang_history`` are rejected earlier (Colang-only). """ - metadata = dict(response.provider_metadata or {}) - if response.usage: - usage = _usage_to_dict(response.usage) - if usage: - metadata["usage"] = usage - return metadata or None + log_options = options.log if options else None + if not log_options or not (log_options.activated_rails or log_options.llm_calls): + return None + + call_records = [record for record in records if _record_has_llm_call(record)] + generation_log = GenerationLog() + generation_log.stats = _build_generation_stats(records, call_records, total_duration) + if log_options.activated_rails: + generation_log.activated_rails = [_activated_rail(record) for record in records] + if log_options.llm_calls: + generation_log.llm_calls = [_call_info(record) for record in call_records] + return generation_log def _build_generation_response( - response_text: str, reasoning_content: Optional[str], response: LLMResponse + response_text: str, + reasoning_content: Optional[str], + response: LLMResponse, + log: Optional[GenerationLog] = None, ) -> GenerationResponse: """Build the structured ``GenerationResponse`` returned when ``options`` are supplied. @@ -218,7 +321,8 @@ def _build_generation_response( canonical ``ToolCall.to_dict()`` shape (dict arguments), matching LLMRails' ``GenerationResponse.tool_calls`` rather than the OpenAI-wire JSON-string shape used for the bare message. ``llm_output`` stays ``None`` to match LLMRails, whose - ``raw_response`` source is never populated. + ``raw_response`` source is never populated. ``log`` is set when the caller requested + log details via ``options.log``. """ result = GenerationResponse(response=[{"role": "assistant", "content": response_text}]) if reasoning_content: @@ -228,15 +332,24 @@ def _build_generation_response( llm_metadata = _build_llm_metadata(response) if llm_metadata: result.llm_metadata = llm_metadata + if log is not None: + result.log = log return result -def _finalize_refusal(structured: bool) -> Union[LLMMessage, GenerationResponse]: - """Shape the refusal message for the active return contract (structured vs bare).""" +def _finalize_refusal(structured: bool, log: Optional[GenerationLog] = None) -> Union[LLMMessage, GenerationResponse]: + """Shape the refusal message for the active return contract (structured vs bare). + + On the structured path the collected ``log`` (rails that ran up to the block) is + attached when the caller requested it. + """ message: LLMMessage = {"role": "assistant", "content": REFUSAL_MESSAGE} - if structured: - return GenerationResponse(response=[message]) - return message + if not structured: + return message + result = GenerationResponse(response=[message]) + if log is not None: + result.log = log + return result def _response_content_for_capture(result: Union[LLMMessage, GenerationResponse]) -> Optional[str]: @@ -255,9 +368,10 @@ def _raise_on_unsupported_options(options: Optional[GenerationOptions], state: o """Raise for ``GenerationOptions``/``state`` features IORails cannot honor. ``state`` and ``output_vars``/``output_data`` require Colang runtime state that - IORails does not have and raise ``ValueError``; ``log`` is deferred to a follow-up - PR and raises ``NotImplementedError``. These raise rather than silently returning - empty data, mirroring how LLMRails rejects options it cannot fulfill. + IORails does not have and raise ``ValueError``. ``log.internal_events`` / + ``log.colang_history`` are Colang-runtime-only and raise ``NotImplementedError``; + ``log.activated_rails`` / ``log.llm_calls`` are supported. These raise rather than + silently returning empty data, mirroring how LLMRails rejects options it cannot fulfill. """ if state is not None: raise ValueError("state is not supported by IORails; it is a stateless input/output rails engine") @@ -266,13 +380,11 @@ def _raise_on_unsupported_options(options: Optional[GenerationOptions], state: o if options.output_vars: raise ValueError("output_vars/output_data is not supported by IORails; it has no Colang context to return") log_options = options.log - if log_options and ( - log_options.activated_rails - or log_options.llm_calls - or log_options.internal_events - or log_options.colang_history - ): - raise NotImplementedError("GenerationResponse `log` is not supported by IORails") + if log_options and (log_options.internal_events or log_options.colang_history): + raise NotImplementedError( + "GenerationLog `internal_events` and `colang_history` are not supported by IORails " + "(no Colang runtime); use `activated_rails` and/or `llm_calls`" + ) def _coerce_generation_options(options: Optional[Union[dict, GenerationOptions]]) -> Optional[GenerationOptions]: @@ -703,25 +815,40 @@ async def _do_generate( tool_input_enabled = options.rails.tool_input if options else True tool_output_enabled = options.rails.tool_output if options else True + # Per-rail + generation records accumulated for the GenerationLog (built only when + # options.log requests it). Each rail check and the main call append their record. + t_start = time.monotonic() + records: list[RailCallRecord] = [] + + def _blocked_return() -> Union[LLMMessage, GenerationResponse]: + """Refusal shaped for the return contract, carrying the log of rails run so far.""" + log_obj = ( + _build_generation_log(records, options, time.monotonic() - t_start) if has_generation_options else None + ) + return _finalize_refusal(has_generation_options, log_obj) + # Agent/client executes tool-calls and sends results to Main LLM with prior conversation history. # Symmetric with INPUT rails log.info("[%s] Running tool result rails", req_id) tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) + records.extend(tool_result.records) if not tool_result.is_safe: log.info("[%s] Tool result blocked: %s", req_id, tool_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.INPUT) - return _finalize_refusal(has_generation_options) + return _blocked_return() if self._speculative_generation: response = await self._do_generate_speculative( - messages, req_id, llm_kwargs, request_span, input_enabled=input_enabled + messages, req_id, llm_kwargs, request_span, input_enabled=input_enabled, records_out=records ) else: - response = await self._do_generate_sequential(messages, req_id, llm_kwargs, input_enabled=input_enabled) + response = await self._do_generate_sequential( + messages, req_id, llm_kwargs, input_enabled=input_enabled, records_out=records + ) if response is None: - return _finalize_refusal(has_generation_options) + return _blocked_return() # Log raw content before reasoning extraction and think-token removal log.debug("[%s] Raw LLM response: %s", req_id, truncate(response.content)) @@ -738,11 +865,12 @@ async def _do_generate( tool_call = await self.rails_manager.are_tool_calls_safe( response.tool_calls, llm_kwargs, enabled=tool_output_enabled ) + records.extend(tool_call.records) if not tool_call.is_safe: log.info("[%s] Tool call blocked: %s", req_id, tool_call.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return _finalize_refusal(has_generation_options) + return _blocked_return() # Output rails check the final answer, not reasoning traces. # Reasoning is re-attached as tags only below so reasoning intentionally bypasses output @@ -753,14 +881,16 @@ async def _do_generate( if not is_tool_call_only: log.info("[%s] Running output rails", req_id) output_result = await self.rails_manager.is_output_safe(messages, response_text, enabled=output_enabled) + records.extend(output_result.records) if not output_result.is_safe: log.info("[%s] Output blocked: %s", req_id, output_result.reason) if self._metrics_enabled: record_request_blocked(RailDirection.OUTPUT) - return _finalize_refusal(has_generation_options) + return _blocked_return() if has_generation_options: - return _build_generation_response(response_text, reasoning_content, response) + log_obj = _build_generation_log(records, options, time.monotonic() - t_start) + return _build_generation_response(response_text, reasoning_content, response, log_obj) # Bare return path: reasoning is delivered inline as a prefix # (LLMRails legacy shape). The structured path above instead puts it in the @@ -771,11 +901,19 @@ async def _do_generate( return _build_assistant_message(response_text, response.tool_calls) async def _do_generate_sequential( - self, messages: LLMMessages, req_id: str, llm_kwargs: dict, *, input_enabled: Union[bool, list[str]] = True + self, + messages: LLMMessages, + req_id: str, + llm_kwargs: dict, + *, + input_enabled: Union[bool, list[str]] = True, + records_out: Optional[list[RailCallRecord]] = None, ) -> Optional[LLMResponse]: """Sequential path: input rails block before LLM generation starts.""" log.info("[%s] Running input rails", req_id) input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) + if records_out is not None: + records_out.extend(input_result.records) if not input_result.is_safe: log.info("[%s] Input blocked: %s", req_id, input_result.reason) if self._metrics_enabled: @@ -783,7 +921,12 @@ async def _do_generate_sequential( return None log.info("[%s] Calling main LLM", req_id) - return await self.engine_registry.model_call("main", messages, **llm_kwargs) + started_at = time.monotonic() + response = await self.engine_registry.model_call("main", messages, **llm_kwargs) + finished_at = time.monotonic() + if records_out is not None: + records_out.append(_make_generation_record(response, started_at, finished_at)) + return response async def _do_generate_speculative( self, @@ -793,6 +936,7 @@ async def _do_generate_speculative( request_span: Optional["Span"] = None, *, input_enabled: Union[bool, list[str]] = True, + records_out: Optional[list[RailCallRecord]] = None, ) -> Optional[LLMResponse]: """Speculative path: input rails and LLM generation race concurrently.""" log.info("[%s] Speculative generation: launching input rails + LLM concurrently", req_id) @@ -802,7 +946,7 @@ async def _do_generate_speculative( try: response = await self._parallel_input_rail_and_response_generation( - rails_task, gen_task, req_id, request_span + rails_task, gen_task, req_id, request_span, records_out=records_out ) except BaseException as outer_exc: for t in (rails_task, gen_task): @@ -835,6 +979,8 @@ async def _parallel_input_rail_and_response_generation( gen_task: asyncio.Task, req_id: str, request_span: Optional["Span"] = None, + *, + records_out: Optional[list[RailCallRecord]] = None, ) -> Optional[LLMResponse]: """Race input rails against LLM generation, return LLMResponse or None (rejected).""" done, _ = await asyncio.wait({rails_task, gen_task}, return_when=asyncio.FIRST_COMPLETED) @@ -847,6 +993,8 @@ async def _parallel_input_rail_and_response_generation( if rails_task in done: input_result = rails_task.result() + if records_out is not None: + records_out.extend(input_result.records) if not input_result.is_safe: log.info("[%s] Input blocked (speculative): %s", req_id, input_result.reason) @@ -872,6 +1020,8 @@ async def _parallel_input_rail_and_response_generation( response = gen_task.result() input_result = await rails_task + if records_out is not None: + records_out.extend(input_result.records) if not input_result.is_safe: log.info("[%s] Input blocked (speculative, gen-first): %s", req_id, input_result.reason) @@ -885,6 +1035,10 @@ async def _parallel_input_rail_and_response_generation( set_speculative_span_attrs(request_span, first_completed, "none") log.debug("[%s] Main LLM response: %s", req_id, truncate(response.content)) + # Speculative races the main call, so per-call timing isn't tracked here; the + # generation record carries usage/model without a duration. + if records_out is not None: + records_out.append(_make_generation_record(response, None, None)) return response def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = None) -> RailsResult: diff --git a/nemoguardrails/guardrails/rail_action.py b/nemoguardrails/guardrails/rail_action.py index 5d4a74ade1..321f7192c2 100644 --- a/nemoguardrails/guardrails/rail_action.py +++ b/nemoguardrails/guardrails/rail_action.py @@ -23,7 +23,10 @@ from __future__ import annotations import logging +import time from abc import ABC, abstractmethod +from contextvars import ContextVar +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Optional, Union from nemoguardrails.guardrails.engine_registry import EngineRegistry @@ -36,7 +39,7 @@ from nemoguardrails.guardrails.telemetry import action_span, record_span_error from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.rails.llm.config import _get_flow_model, _get_flow_name -from nemoguardrails.types import LLMResponse +from nemoguardrails.types import LLMResponse, UsageInfo if TYPE_CHECKING: from opentelemetry.trace import Tracer @@ -44,6 +47,33 @@ log = logging.getLogger(__name__) +@dataclass(frozen=True) +class RailLLMCall: + """Usage/model/timing captured for the LLM call a rail made, for GenerationLog. + + Set by :meth:`RailAction._get_llm_response` and read by RailsManager right after the + rail runs (same async task), so the caller can build the per-rail ``RailCallRecord``. + """ + + usage: Optional[UsageInfo] + llm_model_name: Optional[str] + started_at: float + finished_at: float + duration: float + + +# Request-scoped: the last LLM call a rail made. ``RailAction.run`` clears it at the +# start of each rail so a rail that makes no model call leaves it None. +_rail_llm_call_var: ContextVar[Optional[RailLLMCall]] = ContextVar("rail_llm_call", default=None) + + +def take_rail_llm_call() -> Optional[RailLLMCall]: + """Return and clear the LLM call captured for the rail that just ran.""" + call = _rail_llm_call_var.get() + _rail_llm_call_var.set(None) + return call + + class RailAction(ABC): """Base class for all IORails rail actions. @@ -81,6 +111,9 @@ async def run( bot_response: Optional[str] = None, ) -> RailResult: """Execute the full rail pipeline and return a safety result.""" + # Clear any capture from a prior rail on this task; a rail that makes no model + # call then leaves it None and produces a record with no LLM call. + _rail_llm_call_var.set(None) with action_span(self._tracer, self.action_name) as span: req_id = get_request_id() base_flow = _get_flow_name(flow) @@ -153,10 +186,26 @@ async def _get_llm_response( messages: list[dict], **kwargs: Any, ) -> LLMResponse: - """Call an LLM via EngineRegistry and return the structured response.""" + """Call an LLM via EngineRegistry and return the structured response. + + Captures usage/model/timing into a request-scoped contextvar so RailsManager can + record this call in the GenerationLog after the rail runs. + """ if not model_type: raise RuntimeError("model_type is required for LLM calls") - return await self.engine_registry.model_call(model_type, messages, **kwargs) + started_at = time.monotonic() + response = await self.engine_registry.model_call(model_type, messages, **kwargs) + finished_at = time.monotonic() + _rail_llm_call_var.set( + RailLLMCall( + usage=response.usage, + llm_model_name=response.model, + started_at=started_at, + finished_at=finished_at, + duration=finished_at - started_at, + ) + ) + return response async def _get_api_response( self, diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index d296289ff1..80122823cb 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -37,11 +37,12 @@ from nemoguardrails.guardrails.actions.topic_safety_action import TopicSafetyInputAction from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.guardrails_types import ( + RailCallRecord, RailDirection, RailResult, get_request_id, ) -from nemoguardrails.guardrails.rail_action import RailAction +from nemoguardrails.guardrails.rail_action import RailAction, RailLLMCall, take_rail_llm_call from nemoguardrails.guardrails.telemetry import mark_rail_stop, rail_span, set_rail_content from nemoguardrails.guardrails.tool_rail_action import ToolRailAction from nemoguardrails.guardrails.tool_schema import ToolExchange, Toolset @@ -79,6 +80,28 @@ _ToolActionT = TypeVar("_ToolActionT", bound=ToolRailAction) +def _rail_call_record(flow: str, rail_type: str, result: RailResult, call: Optional[RailLLMCall]) -> RailCallRecord: + """Build the per-rail GenerationLog record from a rail's result + its captured LLM call. + + ``call`` is None for model-free rails (e.g. tool validators); ``return_value`` uses the + action's structured verdict when present, else a minimal ``{"allowed": is_safe}``. + """ + verdict = result.return_value if result.return_value is not None else {"allowed": result.is_safe} + return RailCallRecord( + flow=flow, + rail_type=rail_type, + is_safe=result.is_safe, + action_name=_get_flow_name(flow), + return_value=verdict, + task=flow, + usage=call.usage if call else None, + llm_model_name=call.llm_model_name if call else None, + started_at=call.started_at if call else None, + finished_at=call.finished_at if call else None, + duration=call.duration if call else None, + ) + + class RailsManager: """Orchestrates input and output safety checks for IORails. @@ -304,8 +327,10 @@ async def _run_rail( with rail_span(self._tracer, flow, direction) as span: action = self._actions[flow] result = await action.run(flow, messages, bot_response) + call = take_rail_llm_call() if not result.is_safe and result.triggered_rail is None: result = replace(result, triggered_rail=_get_flow_name(flow) or flow) + result = replace(result, records=(_rail_call_record(flow, direction.value.lower(), result, call),)) mark_rail_stop(span, result.is_safe) # Capture rail input + block reason after the action runs. # RailAction.run() catches its own exceptions and returns @@ -324,6 +349,8 @@ async def _run_tool_call_rail(self, flow: str, tool_calls: list[ToolCall], tools """Dispatch a single tool-call rail to its action, wrapped in an OUTPUT rail span.""" with rail_span(self._tracer, flow, RailDirection.OUTPUT) as span: result = await self._tool_call_actions[flow].run(toolset, tool_calls) + # Tool rails are model-free, so no LLM call is captured (usage stays None). + result = replace(result, records=(_rail_call_record(flow, "tool_output", result, take_rail_llm_call()),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: set_rail_content( @@ -346,6 +373,8 @@ async def _run_tool_result_rail(self, flow: str, exchanges: list[ToolExchange]) result = await action.run(exchange.results, exchange.calls) if not result.is_safe: break + # Tool rails are model-free, so no LLM call is captured (usage stays None). + result = replace(result, records=(_rail_call_record(flow, "tool_input", result, take_rail_llm_call()),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: all_results = [r for exchange in exchanges for r in exchange.results] @@ -368,14 +397,16 @@ async def _run_rails_sequential( """Run rail coroutines sequentially, short-circuiting on first unsafe result.""" req_id = get_request_id() remaining = iter(rails.items()) + collected: list[RailCallRecord] = [] try: for flow, coro in remaining: result = await coro + collected.extend(result.records) log.debug("[%s] %s flow %s result %s", req_id, direction.value, flow, result) if not result.is_safe: log.info("[%s] %s flow %s blocked", req_id, direction.value, flow) - return result - return RailResult(is_safe=True) + return replace(result, records=tuple(collected)) + return RailResult(is_safe=True, records=tuple(collected)) finally: for _, coro in remaining: coro.close() @@ -391,12 +422,14 @@ async def _run_rails_parallel( tasks = list(task_to_flow.keys()) task_order = {task: i for i, task in enumerate(tasks)} pending_tasks: set[asyncio.Task] = set(tasks) + collected: list[RailCallRecord] = [] try: while pending_tasks: done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED) for task in sorted(done, key=lambda t: task_order[t]): result = task.result() + collected.extend(result.records) flow = task_to_flow[task] log.debug("[%s] %s flow %s result %s", req_id, direction.value, flow, result) if not result.is_safe: @@ -411,8 +444,8 @@ async def _run_rails_parallel( t.cancel() if pending_tasks: await asyncio.wait(pending_tasks) - return result - return RailResult(is_safe=True) + return replace(result, records=tuple(collected)) + return RailResult(is_safe=True, records=tuple(collected)) except BaseException: for t in tasks: if not t.done(): diff --git a/tests/guardrails/test_iorails_generation_log.py b/tests/guardrails/test_iorails_generation_log.py new file mode 100644 index 0000000000..db52bea636 --- /dev/null +++ b/tests/guardrails/test_iorails_generation_log.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GenerationLog + token usage for IORails (PR B), matching LLMRails' log contract. + +When ``options.log`` requests ``activated_rails`` and/or ``llm_calls``, IORails +synthesizes a ``GenerationLog`` from the per-rail ``RailCallRecord``s carried on each +``RailResult`` plus the main generation call: ``activated_rails`` (one synthetic +``ExecutedAction`` per rail carrying the real verdict as ``return_value``), a flat +``llm_calls`` list (main + every rail), and aggregate ``stats``. ``internal_events`` +and ``colang_history`` are Colang-runtime-only and raise ``NotImplementedError``. +Token usage lives only in ``log`` — it is no longer surfaced under ``llm_metadata``. +""" + +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio + +from nemoguardrails.guardrails.guardrails_types import RailCallRecord, RailResult +from nemoguardrails.guardrails.iorails import IORails +from nemoguardrails.rails.llm.options import GenerationResponse +from nemoguardrails.types import LLMResponse, UsageInfo +from tests.guardrails.async_helpers import started_iorails +from tests.guardrails.test_data import NEMOGUARDS_CONFIG + +_USER = [{"role": "user", "content": "hi"}] + +_CS_INPUT_FLOW = "content safety check input $model=content_safety" +_CS_OUTPUT_FLOW = "content safety check output $model=content_safety" + + +@pytest_asyncio.fixture +async def iorails(): + """Started IORails instance with worker-queue teardown after each test.""" + async with started_iorails(NEMOGUARDS_CONFIG) as iorails: + yield iorails + + +def _input_record(*, is_safe: bool = True, return_value=None) -> RailCallRecord: + """A content-safety input-rail record with a NeMoGuard-style verdict + usage.""" + return RailCallRecord( + flow=_CS_INPUT_FLOW, + rail_type="input", + is_safe=is_safe, + action_name="content_safety_check_input", + return_value=return_value if return_value is not None else {"allowed": is_safe, "policy_violations": []}, + task="content_safety_check_input $model=content_safety", + usage=UsageInfo(input_tokens=762, output_tokens=8, total_tokens=770), + llm_model_name="nvidia/llama-3.1-nemoguard-8b-content-safety", + llm_provider_name="nim", + started_at=1.0, + finished_at=1.5, + duration=0.5, + ) + + +def _output_record() -> RailCallRecord: + """A content-safety output-rail record with usage.""" + return RailCallRecord( + flow=_CS_OUTPUT_FLOW, + rail_type="output", + is_safe=True, + action_name="content_safety_check_output", + return_value={"allowed": True, "policy_violations": []}, + task="content_safety_check_output $model=content_safety", + usage=UsageInfo(input_tokens=855, output_tokens=15, total_tokens=870), + llm_model_name="nvidia/llama-3.1-nemoguard-8b-content-safety", + llm_provider_name="nim", + started_at=3.0, + finished_at=3.4, + duration=0.4, + ) + + +def _stub_pipeline(iorails: IORails, *, input_records=(), output_records=(), input_safe=True) -> None: + """Stub input/output rails to return given records, and a main call with usage.""" + iorails.rails_manager.is_input_safe = AsyncMock( + return_value=RailResult( + is_safe=input_safe, + reason=None if input_safe else "unsafe", + triggered_rail=None if input_safe else _CS_INPUT_FLOW, + records=tuple(input_records), + ) + ) + iorails.rails_manager.is_output_safe = AsyncMock( + return_value=RailResult(is_safe=True, records=tuple(output_records)) + ) + iorails.engine_registry.model_call = AsyncMock( + return_value=LLMResponse(content="Hi", usage=UsageInfo(input_tokens=99, output_tokens=50, total_tokens=149)) + ) + + +class TestLogGating: + """`log` is only built when explicitly requested.""" + + @pytest.mark.asyncio + async def test_log_none_when_not_requested(self, iorails): + """With options but no log flags, `res.log` stays None.""" + _stub_pipeline(iorails, input_records=[_input_record()]) + + result = await iorails.generate_async(_USER, options={}) + + assert isinstance(result, GenerationResponse) + assert result.log is None + + +class TestLlmCallsAndStats: + """Flat `llm_calls` list and aggregate `stats` cover the main call plus every rail.""" + + @pytest.mark.asyncio + async def test_llm_calls_and_stats_aggregate_main_and_rails(self, iorails): + """`log.llm_calls` lists main + each rail call; `log.stats` sums their tokens.""" + _stub_pipeline(iorails, input_records=[_input_record()], output_records=[_output_record()]) + + result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) + + assert result.log is not None + stats = result.log.stats + assert stats.llm_calls_count == 3 + assert stats.llm_calls_total_prompt_tokens == 762 + 855 + 99 + assert stats.llm_calls_total_completion_tokens == 8 + 15 + 50 + assert stats.llm_calls_total_tokens == 770 + 870 + 149 + + totals = sorted(call.total_tokens for call in result.log.llm_calls) + assert totals == [149, 770, 870] + + +class TestActivatedRails: + """`activated_rails` carries one synthetic action per rail with the real verdict.""" + + @pytest.mark.asyncio + async def test_activated_rail_carries_verdict_as_return_value(self, iorails): + """The rail's structured verdict is preserved as executed_actions[0].return_value.""" + verdict = {"allowed": False, "policy_violations": ["Violence", "Criminal Planning/Confessions"]} + _stub_pipeline(iorails, input_records=[_input_record(is_safe=False, return_value=verdict)], input_safe=False) + + result = await iorails.generate_async(_USER, options={"log": {"activated_rails": True}}) + + assert result.log is not None + rail = next(r for r in result.log.activated_rails if r.name == _CS_INPUT_FLOW) + assert rail.type == "input" + assert rail.executed_actions[0].action_name == "content_safety_check_input" + assert rail.executed_actions[0].return_value == verdict + + @pytest.mark.asyncio + async def test_stop_set_on_blocking_rail(self, iorails): + """A blocking input rail is marked stop=True, and the blocked request still logs it.""" + _stub_pipeline(iorails, input_records=[_input_record(is_safe=False)], input_safe=False) + + result = await iorails.generate_async(_USER, options={"log": {"activated_rails": True}}) + + assert isinstance(result, GenerationResponse) + assert result.log is not None + rail = next(r for r in result.log.activated_rails if r.name == _CS_INPUT_FLOW) + assert rail.stop is True + + +class TestUnsupportedLogOptions: + """Colang-runtime-only log fields raise NotImplementedError.""" + + @pytest.mark.asyncio + async def test_internal_events_raises(self, iorails): + """Requesting internal_events raises — IORails runs no Colang event stream.""" + _stub_pipeline(iorails, input_records=[_input_record()]) + + with pytest.raises(NotImplementedError): + await iorails.generate_async(_USER, options={"log": {"internal_events": True}}) + + @pytest.mark.asyncio + async def test_colang_history_raises(self, iorails): + """Requesting colang_history raises — IORails produces no Colang transcript.""" + _stub_pipeline(iorails, input_records=[_input_record()]) + + with pytest.raises(NotImplementedError): + await iorails.generate_async(_USER, options={"log": {"colang_history": True}}) + + +class TestUsageRemovedFromLlmMetadata: + """Token usage moved to `log`; `llm_metadata` is a pure provider_metadata passthrough.""" + + @pytest.mark.asyncio + async def test_llm_metadata_has_no_usage_key(self, iorails): + """provider_metadata is surfaced verbatim; no `usage` sub-key is added.""" + 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="Hi", + provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}, + usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15), + ) + ) + + result = await iorails.generate_async(_USER, options={}) + + assert result.llm_metadata == {"response_headers": {"nvcf-status": "fulfilled"}} + + @pytest.mark.asyncio + async def test_llm_metadata_none_when_only_usage(self, iorails): + """With usage but no provider_metadata, llm_metadata is None (usage no longer graft-in).""" + 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="Hi", usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15)) + ) + + result = await iorails.generate_async(_USER, options={}) + + assert result.llm_metadata is None diff --git a/tests/guardrails/test_iorails_generation_log_capture.py b/tests/guardrails/test_iorails_generation_log_capture.py new file mode 100644 index 0000000000..687468de1c --- /dev/null +++ b/tests/guardrails/test_iorails_generation_log_capture.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real rail-call capture → GenerationLog for IORails (PR B, Phase 2b). + +Exercises the ACTUAL RailAction + RailsManager pipeline (not stubbed at the +rails-manager level): a rail's model call is captured (usage/model/timing) and its +verdict preserved onto the RailResult, then surfaced as per-rail RailCallRecords that +IORails turns into GenerationLog entries. Only the engine's ``model_call`` is mocked, +so the real content-safety actions parse the responses and produce real records. +""" + +import json +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio + +from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE +from nemoguardrails.rails.llm.options import GenerationResponse +from nemoguardrails.types import LLMResponse, UsageInfo +from tests.guardrails.async_helpers import started_iorails +from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG + +_USER = [{"role": "user", "content": "hi"}] +_CS_INPUT = "content safety check input $model=content_safety" +_CS_MODEL = "nvidia/llama-3.1-nemoguard-8b-content-safety" +_MAIN_MODEL = "meta/llama-3.3-70b-instruct" + +# {"User Safety": "safe", "Response Safety": "safe"} parses safe for both the input +# parser (reads User Safety) and the output parser (reads Response Safety). +_SAFE_BOTH = json.dumps({"User Safety": "safe", "Response Safety": "safe"}) +_UNSAFE_INPUT = json.dumps({"User Safety": "unsafe", "Safety Categories": "S1: Violence"}) + + +@pytest_asyncio.fixture +async def iorails(): + """Started IORails on a content-safety-only config (input + output rails).""" + async with started_iorails(CONTENT_SAFETY_CONFIG) as engine: + yield engine + + +class TestRailRecordCapture: + """A real rail run captures its LLM call + verdict onto RailResult.records.""" + + @pytest.mark.asyncio + async def test_is_input_safe_captures_usage_and_verdict(self, iorails): + """is_input_safe returns a record carrying the rail's tokens, model, and verdict.""" + iorails.engine_registry.model_call = AsyncMock( + return_value=LLMResponse( + content=_SAFE_BOTH, + usage=UsageInfo(input_tokens=762, output_tokens=8, total_tokens=770), + model=_CS_MODEL, + ) + ) + + result = await iorails.rails_manager.is_input_safe(_USER) + + assert result.is_safe is True + assert len(result.records) == 1 + record = result.records[0] + assert record.flow == _CS_INPUT + assert record.rail_type == "input" + assert record.usage.total_tokens == 770 + assert record.llm_model_name == _CS_MODEL + assert record.return_value == {"allowed": True, "policy_violations": []} + + +class TestGenerationLogEndToEnd: + """Full generate_async path builds a GenerationLog from real rail + main-call records.""" + + @pytest.mark.asyncio + async def test_stats_and_activated_rails(self, iorails): + """log covers input CS + main + output CS calls, with the content-safety verdict.""" + + async def _model_call(model_type, messages, **kwargs): + if model_type == "content_safety": + return LLMResponse( + content=_SAFE_BOTH, + usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), + model=_CS_MODEL, + ) + return LLMResponse( + content="Hi", + usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), + model=_MAIN_MODEL, + ) + + iorails.engine_registry.model_call = AsyncMock(side_effect=_model_call) + + result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True, "activated_rails": True}}) + + assert isinstance(result, GenerationResponse) + assert result.log is not None + stats = result.log.stats + assert stats.llm_calls_count == 3 + assert stats.llm_calls_total_prompt_tokens == 100 + 20 + 100 + assert stats.llm_calls_total_completion_tokens == 10 + 5 + 10 + assert stats.llm_calls_total_tokens == 110 + 25 + 110 + + cs_rail = next(rail for rail in result.log.activated_rails if rail.name == _CS_INPUT) + assert cs_rail.type == "input" + assert cs_rail.executed_actions[0].return_value == {"allowed": True, "policy_violations": []} + assert any(rail.type == "generation" for rail in result.log.activated_rails) + + @pytest.mark.asyncio + async def test_blocked_input_logs_verdict_and_stop(self, iorails): + """A blocked input rail logs its unsafe verdict + stop, and the request refuses.""" + iorails.engine_registry.model_call = AsyncMock( + return_value=LLMResponse( + content=_UNSAFE_INPUT, + usage=UsageInfo(input_tokens=762, output_tokens=22, total_tokens=784), + model=_CS_MODEL, + ) + ) + + result = await iorails.generate_async(_USER, options={"log": {"activated_rails": True}}) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] + assert result.log is not None + cs_rail = next(rail for rail in result.log.activated_rails if rail.name == _CS_INPUT) + assert cs_rail.stop is True + verdict = cs_rail.executed_actions[0].return_value + assert verdict["allowed"] is False + assert len(verdict["policy_violations"]) >= 1 diff --git a/tests/guardrails/test_iorails_generation_response.py b/tests/guardrails/test_iorails_generation_response.py index 4c7c7dbe76..ddd81126e4 100644 --- a/tests/guardrails/test_iorails_generation_response.py +++ b/tests/guardrails/test_iorails_generation_response.py @@ -213,11 +213,11 @@ async def test_no_tool_calls_field_is_none(self, iorails): class TestLLMMetadata: - """``llm_metadata`` carries main-call ``provider_metadata`` plus a ``usage`` sub-key.""" + """``llm_metadata`` is the main-call ``provider_metadata`` verbatim; usage lives in ``log``.""" @pytest.mark.asyncio - async def test_provider_metadata_and_usage_merged(self, iorails): - """provider_metadata is surfaced and structured usage is added under ``usage``.""" + async def test_provider_metadata_passthrough(self, iorails): + """provider_metadata is surfaced verbatim; token usage is NOT added under ``usage``.""" _stub_safe_rails(iorails) _stub_model( iorails, @@ -231,14 +231,11 @@ async def test_provider_metadata_and_usage_merged(self, iorails): result = await iorails.generate_async(_USER, options={}) assert isinstance(result, GenerationResponse) - assert result.llm_metadata == { - "response_headers": {"nvcf-status": "fulfilled"}, - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - } + assert result.llm_metadata == {"response_headers": {"nvcf-status": "fulfilled"}} @pytest.mark.asyncio - async def test_usage_only_when_no_provider_metadata(self, iorails): - """With usage but no provider_metadata, ``llm_metadata`` holds just the ``usage`` sub-key.""" + async def test_usage_alone_does_not_populate_metadata(self, iorails): + """With usage but no provider_metadata, ``llm_metadata`` is None — usage moved to log.""" _stub_safe_rails(iorails) _stub_model( iorails, @@ -248,32 +245,7 @@ async def test_usage_only_when_no_provider_metadata(self, iorails): result = await iorails.generate_async(_USER, options={}) assert isinstance(result, GenerationResponse) - assert result.llm_metadata == {"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}} - - @pytest.mark.asyncio - async def test_usage_includes_reasoning_and_cached_tokens_when_present(self, iorails): - """Non-None ``reasoning_tokens``/``cached_tokens`` appear in the ``usage`` sub-key.""" - _stub_safe_rails(iorails) - _stub_model( - iorails, - LLMResponse( - content="Hello", - usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15, reasoning_tokens=3, cached_tokens=2), - ), - ) - - result = await iorails.generate_async(_USER, options={}) - - assert isinstance(result, GenerationResponse) - assert result.llm_metadata == { - "usage": { - "input_tokens": 10, - "output_tokens": 5, - "total_tokens": 15, - "reasoning_tokens": 3, - "cached_tokens": 2, - } - } + assert result.llm_metadata is None @pytest.mark.asyncio async def test_no_metadata_or_usage_is_none(self, iorails): From 07c4b29e9f925cea35515cdbd237d5e4f716bf3f Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:06:26 -0500 Subject: [PATCH 04/11] Change test_log_request_flag_raises to test colang-only log fields raise --- tests/guardrails/test_iorails_generation_response.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/guardrails/test_iorails_generation_response.py b/tests/guardrails/test_iorails_generation_response.py index ddd81126e4..e70ac12b3a 100644 --- a/tests/guardrails/test_iorails_generation_response.py +++ b/tests/guardrails/test_iorails_generation_response.py @@ -299,13 +299,18 @@ async def test_output_vars_list_raises(self, iorails): await iorails.generate_async(_USER, options={"output_vars": ["relevant_chunks"]}) @pytest.mark.asyncio - async def test_log_request_flag_raises(self, iorails): - """Requesting any ``log`` detail raises NotImplementedError while ``log`` is deferred.""" + async def test_colang_only_log_flag_raises(self, iorails): + """Colang-runtime-only log details raise NotImplementedError. + + ``activated_rails``/``llm_calls`` are supported and produce a GenerationLog; + only ``internal_events`` and ``colang_history`` (which need the Colang runtime) + are rejected. + """ _stub_safe_rails(iorails) _stub_model(iorails, LLMResponse(content="Hello")) with pytest.raises(NotImplementedError): - await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) + await iorails.generate_async(_USER, options={"log": {"internal_events": True}}) @pytest.mark.asyncio async def test_state_argument_raises(self, iorails): From 604deabc45ed30523e0bf74efe4a3a5040463372 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:08:53 -0500 Subject: [PATCH 05/11] Add tests for missing coverage --- .../guardrails/test_iorails_generation_log.py | 55 ++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/tests/guardrails/test_iorails_generation_log.py b/tests/guardrails/test_iorails_generation_log.py index db52bea636..9bece958af 100644 --- a/tests/guardrails/test_iorails_generation_log.py +++ b/tests/guardrails/test_iorails_generation_log.py @@ -30,7 +30,7 @@ import pytest_asyncio from nemoguardrails.guardrails.guardrails_types import RailCallRecord, RailResult -from nemoguardrails.guardrails.iorails import IORails +from nemoguardrails.guardrails.iorails import IORails, _activated_rail from nemoguardrails.rails.llm.options import GenerationResponse from nemoguardrails.types import LLMResponse, UsageInfo from tests.guardrails.async_helpers import started_iorails @@ -220,3 +220,56 @@ async def test_llm_metadata_none_when_only_usage(self, iorails): result = await iorails.generate_async(_USER, options={}) assert result.llm_metadata is None + + +class TestActivatedRailHelper: + """`_activated_rail` maps a RailCallRecord to an ActivatedRail + synthetic ExecutedAction.""" + + def test_record_with_llm_call(self): + """A model-backed rail yields one ExecutedAction with its verdict and a single LLMCallInfo.""" + record = RailCallRecord( + flow=_CS_INPUT_FLOW, + rail_type="input", + is_safe=False, + action_name="content_safety_check_input", + return_value={"allowed": False, "policy_violations": ["Violence"]}, + task="content_safety_check_input $model=content_safety", + usage=UsageInfo(input_tokens=762, output_tokens=22, total_tokens=784), + llm_model_name="nvidia/llama-3.1-nemoguard-8b-content-safety", + llm_provider_name="nim", + started_at=1.0, + finished_at=1.5, + duration=0.5, + ) + + rail = _activated_rail(record) + + assert rail.type == "input" + assert rail.name == _CS_INPUT_FLOW + assert rail.stop is True + assert rail.duration == 0.5 + action = rail.executed_actions[0] + assert action.action_name == "content_safety_check_input" + assert action.return_value == {"allowed": False, "policy_violations": ["Violence"]} + assert len(action.llm_calls) == 1 + assert action.llm_calls[0].total_tokens == 784 + assert action.llm_calls[0].llm_model_name == "nvidia/llama-3.1-nemoguard-8b-content-safety" + + def test_record_without_llm_call(self): + """A model-free rail (usage=None) yields an empty llm_calls list; action_name falls back to flow.""" + record = RailCallRecord( + flow="tool call validation", + rail_type="tool_output", + is_safe=True, + action_name=None, + return_value={"allowed": True}, + ) + + rail = _activated_rail(record) + + assert rail.type == "tool_output" + assert rail.stop is False + action = rail.executed_actions[0] + assert action.action_name == "tool call validation" + assert action.llm_calls == [] + assert action.return_value == {"allowed": True} From 95c0a741db678ab4822f687c20deea8fa3753b5b Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Thu, 16 Jul 2026 22:14:35 -0500 Subject: [PATCH 06/11] Add missing GenerationResponse fields --- .../actions/jailbreak_detection_action.py | 7 +-- .../guardrails/actions/topic_safety_action.py | 4 +- nemoguardrails/guardrails/engine_registry.py | 4 ++ nemoguardrails/guardrails/guardrails_types.py | 2 + nemoguardrails/guardrails/iorails.py | 32 +++++++++---- nemoguardrails/guardrails/rail_action.py | 48 +++++++++++++++---- nemoguardrails/guardrails/rails_manager.py | 13 +++-- .../guardrails/test_iorails_generation_log.py | 18 +++++++ .../test_iorails_generation_log_capture.py | 10 ++++ ...est_jailbreak_detection_iorails_actions.py | 12 ++--- .../test_topic_safety_iorails_actions.py | 8 +++- 11 files changed, 122 insertions(+), 36 deletions(-) diff --git a/nemoguardrails/guardrails/actions/jailbreak_detection_action.py b/nemoguardrails/guardrails/actions/jailbreak_detection_action.py index edc806573a..2c974d79a0 100644 --- a/nemoguardrails/guardrails/actions/jailbreak_detection_action.py +++ b/nemoguardrails/guardrails/actions/jailbreak_detection_action.py @@ -41,6 +41,7 @@ def _parse_response(self, response: Any) -> RailResult: raise RuntimeError(f"Jailbreak response missing 'jailbreak' field: {response}") score = response.get("score", "unknown") - if response["jailbreak"]: - return RailResult(is_safe=False, reason=f"Score: {score}") - return RailResult(is_safe=True, reason=f"Score: {score}") + is_jailbreak = response["jailbreak"] + if is_jailbreak: + return RailResult(is_safe=False, reason=f"Score: {score}", return_value=is_jailbreak) + return RailResult(is_safe=True, reason=f"Score: {score}", return_value=is_jailbreak) diff --git a/nemoguardrails/guardrails/actions/topic_safety_action.py b/nemoguardrails/guardrails/actions/topic_safety_action.py index 076824fa02..cee15f0fe6 100644 --- a/nemoguardrails/guardrails/actions/topic_safety_action.py +++ b/nemoguardrails/guardrails/actions/topic_safety_action.py @@ -63,5 +63,5 @@ async def _get_response(self, model_type: Optional[str], prompt: Any) -> str: def _parse_response(self, response: Any) -> RailResult: if response.lower().strip() == "off-topic": - return RailResult(is_safe=False, reason="Topic safety: off-topic") - return RailResult(is_safe=True) + return RailResult(is_safe=False, reason="Topic safety: off-topic", return_value={"on_topic": False}) + return RailResult(is_safe=True, return_value={"on_topic": True}) diff --git a/nemoguardrails/guardrails/engine_registry.py b/nemoguardrails/guardrails/engine_registry.py index a9bf7ef254..1b7ff6f242 100644 --- a/nemoguardrails/guardrails/engine_registry.py +++ b/nemoguardrails/guardrails/engine_registry.py @@ -172,6 +172,10 @@ def _get_engine(self, name: str, expected_type: type[_EngineT]) -> _EngineT: raise TypeError(f"Engine '{name}' is {type(engine).__name__}, expected {expected_type.__name__}") return engine + def provider_name(self, model_type: str) -> str: + """Return the provider/engine name (e.g. 'nim', 'openai') for a model engine.""" + return self._get_engine(model_type, ModelEngine).model_config.engine or "unknown" + async def model_call(self, model_type: str, messages: list[dict], **kwargs: Any) -> LLMResponse: """Route a chat completion request to the named model engine. diff --git a/nemoguardrails/guardrails/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index f5e613a925..f48d174968 100644 --- a/nemoguardrails/guardrails/guardrails_types.py +++ b/nemoguardrails/guardrails/guardrails_types.py @@ -48,9 +48,11 @@ class RailCallRecord: flow: str rail_type: str is_safe: bool + made_call: bool = False action_name: Optional[str] = None return_value: Any = None task: Optional[str] = None + request_id: Optional[str] = None usage: Optional[UsageInfo] = None llm_model_name: Optional[str] = None llm_provider_name: Optional[str] = None diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 21ef59b84a..0ef146f90d 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -192,22 +192,29 @@ def _build_llm_metadata(response: LLMResponse) -> Optional[dict]: def _make_generation_record( - response: LLMResponse, started_at: Optional[float], finished_at: Optional[float] + response: LLMResponse, + started_at: Optional[float], + finished_at: Optional[float], + duration: Optional[float], + provider_name: Optional[str], ) -> RailCallRecord: """Represent the main generation call as a ``RailCallRecord`` (rail_type "generation"). Unifies the main call with the rail calls so the log builder can treat every LLM call - the same way. + the same way. ``started_at``/``finished_at`` are wall-clock timestamps; ``duration`` + is a monotonic delta. """ - duration = (finished_at - started_at) if (started_at is not None and finished_at is not None) else None return RailCallRecord( flow="generation", rail_type="generation", is_safe=True, + made_call=True, action_name="generate_bot_message", task="general", + request_id=response.request_id, usage=response.usage, llm_model_name=response.model, + llm_provider_name=provider_name, started_at=started_at, finished_at=finished_at, duration=duration, @@ -215,8 +222,8 @@ def _make_generation_record( def _record_has_llm_call(record: RailCallRecord) -> bool: - """True when a record reflects a model call (has usage or a model name).""" - return record.usage is not None or record.llm_model_name is not None + """True when a record reflects a model/API call (made a call, or carries usage/model).""" + return record.made_call or record.usage is not None or record.llm_model_name is not None def _call_info(record: RailCallRecord) -> LLMCallInfo: @@ -230,6 +237,7 @@ def _call_info(record: RailCallRecord) -> LLMCallInfo: completion_tokens=usage.output_tokens if usage else None, started_at=record.started_at, finished_at=record.finished_at, + id=record.request_id, llm_model_name=record.llm_model_name or "unknown", llm_provider_name=record.llm_provider_name or "unknown", ) @@ -921,11 +929,14 @@ async def _do_generate_sequential( return None log.info("[%s] Calling main LLM", req_id) - started_at = time.monotonic() + started_at = time.time() + t0 = time.monotonic() response = await self.engine_registry.model_call("main", messages, **llm_kwargs) - finished_at = time.monotonic() + duration = time.monotonic() - t0 + finished_at = time.time() if records_out is not None: - records_out.append(_make_generation_record(response, started_at, finished_at)) + provider = self.engine_registry.provider_name("main") + records_out.append(_make_generation_record(response, started_at, finished_at, duration, provider)) return response async def _do_generate_speculative( @@ -1036,9 +1047,10 @@ async def _parallel_input_rail_and_response_generation( log.debug("[%s] Main LLM response: %s", req_id, truncate(response.content)) # Speculative races the main call, so per-call timing isn't tracked here; the - # generation record carries usage/model without a duration. + # generation record carries usage/model without timestamps or a duration. if records_out is not None: - records_out.append(_make_generation_record(response, None, None)) + provider = self.engine_registry.provider_name("main") + records_out.append(_make_generation_record(response, None, None, None, provider)) return response def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = None) -> RailsResult: diff --git a/nemoguardrails/guardrails/rail_action.py b/nemoguardrails/guardrails/rail_action.py index 321f7192c2..5680085413 100644 --- a/nemoguardrails/guardrails/rail_action.py +++ b/nemoguardrails/guardrails/rail_action.py @@ -49,14 +49,19 @@ @dataclass(frozen=True) class RailLLMCall: - """Usage/model/timing captured for the LLM call a rail made, for GenerationLog. + """Usage/model/timing captured for the call a rail made, for GenerationLog. - Set by :meth:`RailAction._get_llm_response` and read by RailsManager right after the - rail runs (same async task), so the caller can build the per-rail ``RailCallRecord``. + Set by :meth:`RailAction._get_llm_response` (LLM rails) or + :meth:`RailAction._get_api_response` (API rails, e.g. jailbreak — usage/model None), + and read by RailsManager right after the rail runs (same async task), so the caller + can build the per-rail ``RailCallRecord``. ``started_at``/``finished_at`` are + wall-clock (``time.time()``) timestamps; ``duration`` is a monotonic delta. """ usage: Optional[UsageInfo] llm_model_name: Optional[str] + request_id: Optional[str] + provider_name: Optional[str] started_at: float finished_at: float duration: float @@ -67,7 +72,7 @@ class RailLLMCall: _rail_llm_call_var: ContextVar[Optional[RailLLMCall]] = ContextVar("rail_llm_call", default=None) -def take_rail_llm_call() -> Optional[RailLLMCall]: +def get_and_clear_rail_llm_call_contextvar() -> Optional[RailLLMCall]: """Return and clear the LLM call captured for the rail that just ran.""" call = _rail_llm_call_var.get() _rail_llm_call_var.set(None) @@ -193,16 +198,20 @@ async def _get_llm_response( """ if not model_type: raise RuntimeError("model_type is required for LLM calls") - started_at = time.monotonic() + started_at = time.time() + t0 = time.monotonic() response = await self.engine_registry.model_call(model_type, messages, **kwargs) - finished_at = time.monotonic() + duration = time.monotonic() - t0 + finished_at = time.time() _rail_llm_call_var.set( RailLLMCall( usage=response.usage, llm_model_name=response.model, + request_id=response.request_id, + provider_name=self.engine_registry.provider_name(model_type), started_at=started_at, finished_at=finished_at, - duration=finished_at - started_at, + duration=duration, ) ) return response @@ -213,8 +222,29 @@ async def _get_api_response( body: dict[str, Any], **kwargs: Any, ) -> dict[str, Any]: - """Call an API endpoint via EngineRegistry and return the response dict.""" - return await self.engine_registry.api_call(api_name, body, **kwargs) + """Call an API endpoint via EngineRegistry and return the response dict. + + Records the call (with no token usage or model name) so API-backed rails such as + jailbreak detection still appear in the GenerationLog's ``llm_calls`` and counts, + matching LLMRails. + """ + started_at = time.time() + t0 = time.monotonic() + response = await self.engine_registry.api_call(api_name, body, **kwargs) + duration = time.monotonic() - t0 + finished_at = time.time() + _rail_llm_call_var.set( + RailLLMCall( + usage=None, + llm_model_name=None, + request_id=None, + provider_name=None, + started_at=started_at, + finished_at=finished_at, + duration=duration, + ) + ) + return response async def _get_local_response(self, **kwargs: Any) -> Any: """Run a local/in-process check. Override in subclasses that need it.""" diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index 80122823cb..20bb27e329 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -42,7 +42,7 @@ RailResult, get_request_id, ) -from nemoguardrails.guardrails.rail_action import RailAction, RailLLMCall, take_rail_llm_call +from nemoguardrails.guardrails.rail_action import RailAction, RailLLMCall, get_and_clear_rail_llm_call_contextvar from nemoguardrails.guardrails.telemetry import mark_rail_stop, rail_span, set_rail_content from nemoguardrails.guardrails.tool_rail_action import ToolRailAction from nemoguardrails.guardrails.tool_schema import ToolExchange, Toolset @@ -91,11 +91,14 @@ def _rail_call_record(flow: str, rail_type: str, result: RailResult, call: Optio flow=flow, rail_type=rail_type, is_safe=result.is_safe, + made_call=call is not None, action_name=_get_flow_name(flow), return_value=verdict, task=flow, + request_id=call.request_id if call else None, usage=call.usage if call else None, llm_model_name=call.llm_model_name if call else None, + llm_provider_name=call.provider_name if call else None, started_at=call.started_at if call else None, finished_at=call.finished_at if call else None, duration=call.duration if call else None, @@ -327,7 +330,7 @@ async def _run_rail( with rail_span(self._tracer, flow, direction) as span: action = self._actions[flow] result = await action.run(flow, messages, bot_response) - call = take_rail_llm_call() + call = get_and_clear_rail_llm_call_contextvar() if not result.is_safe and result.triggered_rail is None: result = replace(result, triggered_rail=_get_flow_name(flow) or flow) result = replace(result, records=(_rail_call_record(flow, direction.value.lower(), result, call),)) @@ -350,7 +353,8 @@ async def _run_tool_call_rail(self, flow: str, tool_calls: list[ToolCall], tools with rail_span(self._tracer, flow, RailDirection.OUTPUT) as span: result = await self._tool_call_actions[flow].run(toolset, tool_calls) # Tool rails are model-free, so no LLM call is captured (usage stays None). - result = replace(result, records=(_rail_call_record(flow, "tool_output", result, take_rail_llm_call()),)) + call = get_and_clear_rail_llm_call_contextvar() + result = replace(result, records=(_rail_call_record(flow, "tool_output", result, call),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: set_rail_content( @@ -374,7 +378,8 @@ async def _run_tool_result_rail(self, flow: str, exchanges: list[ToolExchange]) if not result.is_safe: break # Tool rails are model-free, so no LLM call is captured (usage stays None). - result = replace(result, records=(_rail_call_record(flow, "tool_input", result, take_rail_llm_call()),)) + call = get_and_clear_rail_llm_call_contextvar() + result = replace(result, records=(_rail_call_record(flow, "tool_input", result, call),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: all_results = [r for exchange in exchanges for r in exchange.results] diff --git a/tests/guardrails/test_iorails_generation_log.py b/tests/guardrails/test_iorails_generation_log.py index 9bece958af..3ff878542d 100644 --- a/tests/guardrails/test_iorails_generation_log.py +++ b/tests/guardrails/test_iorails_generation_log.py @@ -273,3 +273,21 @@ def test_record_without_llm_call(self): assert action.action_name == "tool call validation" assert action.llm_calls == [] assert action.return_value == {"allowed": True} + + def test_api_rail_made_call_without_usage_counts(self): + """An API rail (made_call=True, usage=None, e.g. jailbreak) still yields one llm_call.""" + record = RailCallRecord( + flow="jailbreak detection model", + rail_type="input", + is_safe=True, + made_call=True, + action_name="jailbreak detection model", + return_value=False, + task="jailbreak detection model", + ) + + rail = _activated_rail(record) + + assert len(rail.executed_actions[0].llm_calls) == 1 + assert rail.executed_actions[0].llm_calls[0].total_tokens is None + assert rail.executed_actions[0].return_value is False diff --git a/tests/guardrails/test_iorails_generation_log_capture.py b/tests/guardrails/test_iorails_generation_log_capture.py index 687468de1c..c3699f930d 100644 --- a/tests/guardrails/test_iorails_generation_log_capture.py +++ b/tests/guardrails/test_iorails_generation_log_capture.py @@ -63,6 +63,7 @@ async def test_is_input_safe_captures_usage_and_verdict(self, iorails): content=_SAFE_BOTH, usage=UsageInfo(input_tokens=762, output_tokens=8, total_tokens=770), model=_CS_MODEL, + request_id="req-cs-input", ) ) @@ -73,8 +74,11 @@ async def test_is_input_safe_captures_usage_and_verdict(self, iorails): record = result.records[0] assert record.flow == _CS_INPUT assert record.rail_type == "input" + assert record.made_call is True assert record.usage.total_tokens == 770 assert record.llm_model_name == _CS_MODEL + assert record.llm_provider_name == "nim" + assert record.request_id == "req-cs-input" assert record.return_value == {"allowed": True, "policy_violations": []} @@ -91,11 +95,13 @@ async def _model_call(model_type, messages, **kwargs): content=_SAFE_BOTH, usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), model=_CS_MODEL, + request_id="req-cs", ) return LLMResponse( content="Hi", usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), model=_MAIN_MODEL, + request_id="req-main", ) iorails.engine_registry.model_call = AsyncMock(side_effect=_model_call) @@ -113,6 +119,10 @@ async def _model_call(model_type, messages, **kwargs): cs_rail = next(rail for rail in result.log.activated_rails if rail.name == _CS_INPUT) assert cs_rail.type == "input" assert cs_rail.executed_actions[0].return_value == {"allowed": True, "policy_violations": []} + # id + provider are threaded through the capture (parity fixes). + cs_call = cs_rail.executed_actions[0].llm_calls[0] + assert cs_call.id == "req-cs" + assert cs_call.llm_provider_name == "nim" assert any(rail.type == "generation" for rail in result.log.activated_rails) @pytest.mark.asyncio diff --git a/tests/guardrails/test_jailbreak_detection_iorails_actions.py b/tests/guardrails/test_jailbreak_detection_iorails_actions.py index 57c00d38cc..59837b8560 100644 --- a/tests/guardrails/test_jailbreak_detection_iorails_actions.py +++ b/tests/guardrails/test_jailbreak_detection_iorails_actions.py @@ -52,14 +52,14 @@ def test_creates_api_payload(self, action): class TestJailbreakParseResponse: def test_safe(self, action): - assert action._parse_response({"jailbreak": False, "score": 0.1}) == RailResult( - is_safe=True, reason="Score: 0.1" - ) + result = action._parse_response({"jailbreak": False, "score": 0.1}) + assert result == RailResult(is_safe=True, reason="Score: 0.1") + assert result.return_value is False def test_jailbreak_detected(self, action): - assert action._parse_response({"jailbreak": True, "score": 0.95}) == RailResult( - is_safe=False, reason="Score: 0.95" - ) + result = action._parse_response({"jailbreak": True, "score": 0.95}) + assert result == RailResult(is_safe=False, reason="Score: 0.95") + assert result.return_value is True def test_missing_jailbreak_field_raises(self, action): with pytest.raises(RuntimeError, match="missing 'jailbreak' field"): diff --git a/tests/guardrails/test_topic_safety_iorails_actions.py b/tests/guardrails/test_topic_safety_iorails_actions.py index 6504552dd8..7abbfd9a9b 100644 --- a/tests/guardrails/test_topic_safety_iorails_actions.py +++ b/tests/guardrails/test_topic_safety_iorails_actions.py @@ -92,10 +92,14 @@ def test_multi_turn_messages_included(self, action): class TestTopicSafetyParseResponse: def test_on_topic(self, action): - assert action._parse_response("on-topic") == RailResult(is_safe=True) + result = action._parse_response("on-topic") + assert result == RailResult(is_safe=True) + assert result.return_value == {"on_topic": True} def test_off_topic(self, action): - assert action._parse_response("off-topic") == RailResult(is_safe=False, reason="Topic safety: off-topic") + result = action._parse_response("off-topic") + assert result == RailResult(is_safe=False, reason="Topic safety: off-topic") + assert result.return_value == {"on_topic": False} @pytest.mark.parametrize("text", ["Off-Topic", " off-topic \n", "OFF-TOPIC"]) def test_off_topic_variants(self, action, text): From 3001e83dfd5db5140d3c1805eb422da9cbedd076 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 10:22:54 -0500 Subject: [PATCH 07/11] Match LLMRails underscore-separated task names rather than space-separated --- nemoguardrails/guardrails/rails_manager.py | 12 +++++-- tests/guardrails/test_rails_manager.py | 37 +++++++++++++++++++++- 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index 20bb27e329..4a4e19ed76 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -47,7 +47,7 @@ from nemoguardrails.guardrails.tool_rail_action import ToolRailAction from nemoguardrails.guardrails.tool_schema import ToolExchange, Toolset from nemoguardrails.llm.taskmanager import LLMTaskManager -from nemoguardrails.rails.llm.config import _get_flow_name +from nemoguardrails.rails.llm.config import _get_flow_model, _get_flow_name from nemoguardrails.types import ToolCall if TYPE_CHECKING: @@ -87,14 +87,20 @@ def _rail_call_record(flow: str, rail_type: str, result: RailResult, call: Optio action's structured verdict when present, else a minimal ``{"allowed": is_safe}``. """ verdict = result.return_value if result.return_value is not None else {"allowed": result.is_safe} + # GenerationLog parity with LLMRails: action_name/task use the prompt-template key + # (underscores) rather than the space-separated Colang flow name; ``flow`` keeps spaces. + base_name = _get_flow_name(flow) or flow + model = _get_flow_model(flow) + action_name = base_name.replace(" ", "_") + task = f"{action_name} $model={model}" if model else action_name return RailCallRecord( flow=flow, rail_type=rail_type, is_safe=result.is_safe, made_call=call is not None, - action_name=_get_flow_name(flow), + action_name=action_name, return_value=verdict, - task=flow, + task=task, request_id=call.request_id if call else None, usage=call.usage if call else None, llm_model_name=call.llm_model_name if call else None, diff --git a/tests/guardrails/test_rails_manager.py b/tests/guardrails/test_rails_manager.py index f3957d0697..fe39e6ea98 100644 --- a/tests/guardrails/test_rails_manager.py +++ b/tests/guardrails/test_rails_manager.py @@ -32,7 +32,7 @@ from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult -from nemoguardrails.guardrails.rails_manager import RailsManager +from nemoguardrails.guardrails.rails_manager import RailsManager, _rail_call_record from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.tracing.constants import GuardrailsAttributes @@ -971,3 +971,38 @@ async def test_safe_result_has_no_triggered_rail(self, content_safety_rails_mana result = await content_safety_rails_manager.is_input_safe(MESSAGES) assert result.is_safe assert result.triggered_rail is None + + +class TestRailCallRecordNaming: + """`_rail_call_record` names task/action_name in LLMRails' underscore form. + + The GenerationLog's ``executed_actions[].action_name`` and ``llm_calls[].task`` + must match LLMRails, which uses the prompt-template key (underscores) rather than + the space-separated Colang flow name. ``flow`` itself keeps the space form. + """ + + def test_modelled_rail_uses_underscore_task_and_action_name(self): + """A ``$model=`` flow yields an underscore action_name plus an underscore ``$model`` task.""" + record = _rail_call_record( + flow="content safety check input $model=content_safety", + rail_type="input", + result=RailResult(is_safe=True), + call=None, + ) + + assert record.flow == "content safety check input $model=content_safety" + assert record.action_name == "content_safety_check_input" + assert record.task == "content_safety_check_input $model=content_safety" + + def test_modelless_rail_uses_underscore_action_name_and_task(self): + """A model-free flow (jailbreak) yields an underscore task with no ``$model`` suffix.""" + record = _rail_call_record( + flow="jailbreak detection model", + rail_type="input", + result=RailResult(is_safe=True), + call=None, + ) + + assert record.flow == "jailbreak detection model" + assert record.action_name == "jailbreak_detection_model" + assert record.task == "jailbreak_detection_model" From a04d242c518b0a1e47280d949f51dc656f411e82 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:09:55 -0500 Subject: [PATCH 08/11] Add prompt and response to GenerationResponse activated_rails.executed_actions.llm_calls.{prompt, completion} --- examples/configs/nemoguards/config.yml | 2 +- nemoguardrails/guardrails/guardrails_types.py | 12 +++++ nemoguardrails/guardrails/iorails.py | 24 +++++++-- nemoguardrails/guardrails/rail_action.py | 7 +++ nemoguardrails/guardrails/rails_manager.py | 2 + .../guardrails/test_iorails_generation_log.py | 37 +++++++++++++ .../test_iorails_generation_log_capture.py | 53 +++++++++++++++++++ tests/guardrails/test_rails_manager.py | 21 +++++++- 8 files changed, 151 insertions(+), 7 deletions(-) diff --git a/examples/configs/nemoguards/config.yml b/examples/configs/nemoguards/config.yml index 878f8061a2..36af0ac8a9 100644 --- a/examples/configs/nemoguards/config.yml +++ b/examples/configs/nemoguards/config.yml @@ -1,7 +1,7 @@ models: - type: main engine: nim - model: meta/llama-3.3-70b-instruct + model: nvidia/nemotron-3-ultra-550b-a55b - type: content_safety engine: nim diff --git a/nemoguardrails/guardrails/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index f48d174968..9dd3825815 100644 --- a/nemoguardrails/guardrails/guardrails_types.py +++ b/nemoguardrails/guardrails/guardrails_types.py @@ -56,6 +56,8 @@ class RailCallRecord: usage: Optional[UsageInfo] = None llm_model_name: Optional[str] = None llm_provider_name: Optional[str] = None + prompt: Optional[str] = None + completion: Optional[str] = None started_at: Optional[float] = None finished_at: Optional[float] = None duration: Optional[float] = None @@ -127,3 +129,13 @@ def truncate(text: object, max_len: int | None = None) -> str: if len(s) <= limit: return s return s[:limit] + "..." + + +def serialize_prompt(messages: list[dict]) -> str: + """Render a chat message list to a role-labeled string for GenerationLog's ``prompt``. + + Content parity with LLMRails' logged prompt, not byte-for-byte format parity: each + message becomes ``": "`` and messages are blank-line separated. A + message with no content (e.g. a reasoning-only or tool-call turn) renders as empty. + """ + return "\n\n".join(f"{m.get('role', '')}: {m.get('content') or ''}" for m in messages) diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 0ef146f90d..aed5bb5ef8 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -40,6 +40,7 @@ RailCallRecord, RailDirection, get_request_id, + serialize_prompt, truncate, ) from nemoguardrails.guardrails.rails_manager import RailsManager @@ -197,12 +198,14 @@ def _make_generation_record( finished_at: Optional[float], duration: Optional[float], provider_name: Optional[str], + prompt: Optional[str], ) -> RailCallRecord: """Represent the main generation call as a ``RailCallRecord`` (rail_type "generation"). Unifies the main call with the rail calls so the log builder can treat every LLM call - the same way. ``started_at``/``finished_at`` are wall-clock timestamps; ``duration`` - is a monotonic delta. + the same way. ``prompt`` is the serialized main-model messages; ``completion`` is the + response content. ``started_at``/``finished_at`` are wall-clock timestamps; + ``duration`` is a monotonic delta. """ return RailCallRecord( flow="generation", @@ -215,6 +218,8 @@ def _make_generation_record( usage=response.usage, llm_model_name=response.model, llm_provider_name=provider_name, + prompt=prompt, + completion=response.content, started_at=started_at, finished_at=finished_at, duration=duration, @@ -238,6 +243,8 @@ def _call_info(record: RailCallRecord) -> LLMCallInfo: started_at=record.started_at, finished_at=record.finished_at, id=record.request_id, + prompt=record.prompt, + completion=record.completion, llm_model_name=record.llm_model_name or "unknown", llm_provider_name=record.llm_provider_name or "unknown", ) @@ -936,7 +943,8 @@ async def _do_generate_sequential( finished_at = time.time() if records_out is not None: provider = self.engine_registry.provider_name("main") - records_out.append(_make_generation_record(response, started_at, finished_at, duration, provider)) + prompt = serialize_prompt(messages) + records_out.append(_make_generation_record(response, started_at, finished_at, duration, provider, prompt)) return response async def _do_generate_speculative( @@ -957,7 +965,12 @@ async def _do_generate_speculative( try: response = await self._parallel_input_rail_and_response_generation( - rails_task, gen_task, req_id, request_span, records_out=records_out + rails_task, + gen_task, + req_id, + request_span, + records_out=records_out, + main_prompt=serialize_prompt(messages), ) except BaseException as outer_exc: for t in (rails_task, gen_task): @@ -992,6 +1005,7 @@ async def _parallel_input_rail_and_response_generation( request_span: Optional["Span"] = None, *, records_out: Optional[list[RailCallRecord]] = None, + main_prompt: str = "", ) -> Optional[LLMResponse]: """Race input rails against LLM generation, return LLMResponse or None (rejected).""" done, _ = await asyncio.wait({rails_task, gen_task}, return_when=asyncio.FIRST_COMPLETED) @@ -1050,7 +1064,7 @@ async def _parallel_input_rail_and_response_generation( # generation record carries usage/model without timestamps or a duration. if records_out is not None: provider = self.engine_registry.provider_name("main") - records_out.append(_make_generation_record(response, None, None, None, provider)) + records_out.append(_make_generation_record(response, None, None, None, provider, main_prompt)) return response def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = None) -> RailsResult: diff --git a/nemoguardrails/guardrails/rail_action.py b/nemoguardrails/guardrails/rail_action.py index 5680085413..761955c6a6 100644 --- a/nemoguardrails/guardrails/rail_action.py +++ b/nemoguardrails/guardrails/rail_action.py @@ -34,6 +34,7 @@ LLMMessages, RailResult, get_request_id, + serialize_prompt, truncate, ) from nemoguardrails.guardrails.telemetry import action_span, record_span_error @@ -62,6 +63,8 @@ class RailLLMCall: llm_model_name: Optional[str] request_id: Optional[str] provider_name: Optional[str] + prompt: Optional[str] + completion: Optional[str] started_at: float finished_at: float duration: float @@ -209,6 +212,8 @@ async def _get_llm_response( llm_model_name=response.model, request_id=response.request_id, provider_name=self.engine_registry.provider_name(model_type), + prompt=serialize_prompt(messages), + completion=response.content, started_at=started_at, finished_at=finished_at, duration=duration, @@ -239,6 +244,8 @@ async def _get_api_response( llm_model_name=None, request_id=None, provider_name=None, + prompt=None, + completion=None, started_at=started_at, finished_at=finished_at, duration=duration, diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index 4a4e19ed76..e6b5dd6e75 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -105,6 +105,8 @@ def _rail_call_record(flow: str, rail_type: str, result: RailResult, call: Optio usage=call.usage if call else None, llm_model_name=call.llm_model_name if call else None, llm_provider_name=call.provider_name if call else None, + prompt=call.prompt if call else None, + completion=call.completion if call else None, started_at=call.started_at if call else None, finished_at=call.finished_at if call else None, duration=call.duration if call else None, diff --git a/tests/guardrails/test_iorails_generation_log.py b/tests/guardrails/test_iorails_generation_log.py index 3ff878542d..aedd07c2d6 100644 --- a/tests/guardrails/test_iorails_generation_log.py +++ b/tests/guardrails/test_iorails_generation_log.py @@ -291,3 +291,40 @@ def test_api_rail_made_call_without_usage_counts(self): assert len(rail.executed_actions[0].llm_calls) == 1 assert rail.executed_actions[0].llm_calls[0].total_tokens is None assert rail.executed_actions[0].return_value is False + + def test_record_maps_prompt_and_completion(self): + """A record's captured prompt/completion surface on the mapped LLMCallInfo.""" + record = RailCallRecord( + flow=_CS_INPUT_FLOW, + rail_type="input", + is_safe=True, + made_call=True, + action_name="content_safety_check_input", + return_value={"allowed": True}, + task="content_safety_check_input $model=content_safety", + usage=UsageInfo(input_tokens=10, output_tokens=2, total_tokens=12), + prompt="user: hi", + completion='{"User Safety": "safe"}', + ) + + call = _activated_rail(record).executed_actions[0].llm_calls[0] + + assert call.prompt == "user: hi" + assert call.completion == '{"User Safety": "safe"}' + + def test_api_rail_record_has_no_prompt_or_completion(self): + """An API rail (jailbreak) captures no content, so prompt/completion map to None.""" + record = RailCallRecord( + flow="jailbreak detection model", + rail_type="input", + is_safe=True, + made_call=True, + action_name="jailbreak_detection_model", + return_value=False, + task="jailbreak_detection_model", + ) + + call = _activated_rail(record).executed_actions[0].llm_calls[0] + + assert call.prompt is None + assert call.completion is None diff --git a/tests/guardrails/test_iorails_generation_log_capture.py b/tests/guardrails/test_iorails_generation_log_capture.py index c3699f930d..4545510818 100644 --- a/tests/guardrails/test_iorails_generation_log_capture.py +++ b/tests/guardrails/test_iorails_generation_log_capture.py @@ -125,6 +125,59 @@ async def _model_call(model_type, messages, **kwargs): assert cs_call.llm_provider_name == "nim" assert any(rail.type == "generation" for rail in result.log.activated_rails) + @pytest.mark.asyncio + async def test_llm_calls_capture_prompt_and_completion(self, iorails): + """Each LLM-backed call in the log carries its serialized prompt and raw completion.""" + + async def _model_call(model_type, messages, **kwargs): + if model_type == "content_safety": + return LLMResponse( + content=_SAFE_BOTH, + usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), + model=_CS_MODEL, + ) + return LLMResponse( + content="Hi", + usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), + model=_MAIN_MODEL, + ) + + iorails.engine_registry.model_call = AsyncMock(side_effect=_model_call) + + result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) + + assert isinstance(result, GenerationResponse) + assert result.log is not None + llm_calls = result.log.llm_calls or [] + cs_call = next(c for c in llm_calls if c.task and "content_safety_check_input" in c.task) + main_call = next(c for c in llm_calls if c.task == "general") + + assert cs_call.completion == _SAFE_BOTH + assert cs_call.prompt is not None + assert "hi" in cs_call.prompt + assert main_call.completion == "Hi" + assert main_call.prompt is not None + assert "hi" in main_call.prompt + + @pytest.mark.asyncio + async def test_prompt_serializes_role_and_content(self, iorails): + """The serialized prompt is role-labeled so a reader can tell system from user turns.""" + iorails.engine_registry.model_call = AsyncMock( + return_value=LLMResponse( + content=_SAFE_BOTH, + usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), + model=_CS_MODEL, + ) + ) + + result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) + + assert result.log is not None + llm_calls = result.log.llm_calls or [] + cs_call = next(c for c in llm_calls if c.task and "content_safety_check_input" in c.task) + assert cs_call.prompt is not None + assert "user:" in cs_call.prompt or "system:" in cs_call.prompt + @pytest.mark.asyncio async def test_blocked_input_logs_verdict_and_stop(self, iorails): """A blocked input rail logs its unsafe verdict + stop, and the request refuses.""" diff --git a/tests/guardrails/test_rails_manager.py b/tests/guardrails/test_rails_manager.py index fe39e6ea98..7bcd92f17b 100644 --- a/tests/guardrails/test_rails_manager.py +++ b/tests/guardrails/test_rails_manager.py @@ -31,7 +31,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from nemoguardrails.guardrails.engine_registry import EngineRegistry -from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult +from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult, serialize_prompt from nemoguardrails.guardrails.rails_manager import RailsManager, _rail_call_record from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.rails.llm.config import RailsConfig @@ -1006,3 +1006,22 @@ def test_modelless_rail_uses_underscore_action_name_and_task(self): assert record.flow == "jailbreak detection model" assert record.action_name == "jailbreak_detection_model" assert record.task == "jailbreak_detection_model" + + +class TestSerializePrompt: + """`serialize_prompt` renders a message list to a role-labeled string for the log.""" + + def test_role_labeled_join(self): + """Each message renders as ': ', blank-line separated.""" + out = serialize_prompt( + [ + {"role": "system", "content": "be nice"}, + {"role": "user", "content": "hi"}, + ] + ) + assert out == "system: be nice\n\nuser: hi" + + def test_missing_content_renders_empty(self): + """A message with no content (reasoning/tool turn) renders as empty, not an error.""" + out = serialize_prompt([{"role": "assistant", "content": None}]) + assert out == "assistant: " From 559ace5134a5b352c8465b8084131db7b16ae53d Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:29:51 -0500 Subject: [PATCH 09/11] Add missing test coverage --- .../test_iorails_generation_response.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/guardrails/test_iorails_generation_response.py b/tests/guardrails/test_iorails_generation_response.py index e70ac12b3a..a9dde9c51a 100644 --- a/tests/guardrails/test_iorails_generation_response.py +++ b/tests/guardrails/test_iorails_generation_response.py @@ -31,7 +31,7 @@ import pytest_asyncio from nemoguardrails.guardrails.guardrails_types import RailResult -from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails +from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails, _response_content_for_capture from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.rails.llm.options import GenerationOptions, GenerationResponse from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction, UsageInfo @@ -380,3 +380,30 @@ def test_sync_generate_with_options_returns_generation_response(self, iorails_sy result = iorails_sync.generate(_USER, options={}) assert isinstance(result, GenerationResponse) + + +class TestResponseContentForCapture: + """`_response_content_for_capture` extracts assistant text from either return shape.""" + + def test_generation_response_list(self): + """A structured (list) response yields the last assistant message's content.""" + result = GenerationResponse(response=[{"role": "assistant", "content": "hi there"}]) + assert _response_content_for_capture(result) == "hi there" + + def test_generation_response_str(self): + """A structured (str) response is returned as-is.""" + result = GenerationResponse(response="hi there") + assert _response_content_for_capture(result) == "hi there" + + def test_generation_response_empty_list(self): + """An empty response list has no content to capture.""" + assert _response_content_for_capture(GenerationResponse(response=[])) is None + + def test_generation_response_non_str_content(self): + """A non-string message content is not captured.""" + result = GenerationResponse(response=[{"role": "assistant", "content": None}]) + assert _response_content_for_capture(result) is None + + def test_bare_message_dict(self): + """The bare-dict return path reads content off the message directly.""" + assert _response_content_for_capture({"role": "assistant", "content": "hi there"}) == "hi there" From 7912d6c12744646b974352ae955f410a9a86fec7 Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:38:26 -0500 Subject: [PATCH 10/11] Compact and clean up tests --- .../test_iorails_generation_log_capture.py | 57 ++-- .../test_iorails_generation_response.py | 276 +++++++----------- 2 files changed, 127 insertions(+), 206 deletions(-) diff --git a/tests/guardrails/test_iorails_generation_log_capture.py b/tests/guardrails/test_iorails_generation_log_capture.py index 4545510818..9dfd979569 100644 --- a/tests/guardrails/test_iorails_generation_log_capture.py +++ b/tests/guardrails/test_iorails_generation_log_capture.py @@ -45,6 +45,27 @@ _UNSAFE_INPUT = json.dumps({"User Safety": "unsafe", "Safety Categories": "S1: Violence"}) +def _cs_and_main_model_call(*, cs_request_id=None, main_request_id=None): + """Build a ``model_call`` side_effect: content-safety returns _SAFE_BOTH, the main model returns "Hi".""" + + async def _model_call(model_type, messages, **kwargs): + if model_type == "content_safety": + return LLMResponse( + content=_SAFE_BOTH, + usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), + model=_CS_MODEL, + request_id=cs_request_id, + ) + return LLMResponse( + content="Hi", + usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), + model=_MAIN_MODEL, + request_id=main_request_id, + ) + + return _model_call + + @pytest_asyncio.fixture async def iorails(): """Started IORails on a content-safety-only config (input + output rails).""" @@ -88,23 +109,9 @@ class TestGenerationLogEndToEnd: @pytest.mark.asyncio async def test_stats_and_activated_rails(self, iorails): """log covers input CS + main + output CS calls, with the content-safety verdict.""" - - async def _model_call(model_type, messages, **kwargs): - if model_type == "content_safety": - return LLMResponse( - content=_SAFE_BOTH, - usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), - model=_CS_MODEL, - request_id="req-cs", - ) - return LLMResponse( - content="Hi", - usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), - model=_MAIN_MODEL, - request_id="req-main", - ) - - iorails.engine_registry.model_call = AsyncMock(side_effect=_model_call) + iorails.engine_registry.model_call = AsyncMock( + side_effect=_cs_and_main_model_call(cs_request_id="req-cs", main_request_id="req-main") + ) result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True, "activated_rails": True}}) @@ -128,21 +135,7 @@ async def _model_call(model_type, messages, **kwargs): @pytest.mark.asyncio async def test_llm_calls_capture_prompt_and_completion(self, iorails): """Each LLM-backed call in the log carries its serialized prompt and raw completion.""" - - async def _model_call(model_type, messages, **kwargs): - if model_type == "content_safety": - return LLMResponse( - content=_SAFE_BOTH, - usage=UsageInfo(input_tokens=100, output_tokens=10, total_tokens=110), - model=_CS_MODEL, - ) - return LLMResponse( - content="Hi", - usage=UsageInfo(input_tokens=20, output_tokens=5, total_tokens=25), - model=_MAIN_MODEL, - ) - - iorails.engine_registry.model_call = AsyncMock(side_effect=_model_call) + iorails.engine_registry.model_call = AsyncMock(side_effect=_cs_and_main_model_call()) result = await iorails.generate_async(_USER, options={"log": {"llm_calls": True}}) diff --git a/tests/guardrails/test_iorails_generation_response.py b/tests/guardrails/test_iorails_generation_response.py index a9dde9c51a..2586ee9a55 100644 --- a/tests/guardrails/test_iorails_generation_response.py +++ b/tests/guardrails/test_iorails_generation_response.py @@ -19,10 +19,11 @@ ``GenerationResponse`` instead of a bare ``LLMMessage`` dict, mirroring LLMRails' ``if gen_options:`` branch. When ``options`` is absent the bare dict is returned unchanged. The structured path populates ``response``, ``reasoning_content``, -``tool_calls`` (as ``ToolCall.to_dict()`` with dict arguments), and -``llm_metadata`` (main-call ``provider_metadata`` plus a ``usage`` sub-key); -``llm_output`` is always ``None`` (parity with LLMRails' unwired ``raw_response``); -``log`` is deferred; and ``output_vars``/``state``/``log`` request options raise. +``tool_calls`` (as ``ToolCall.to_dict()`` with dict arguments), ``log`` (from +per-rail records), and ``llm_metadata`` (main-call ``provider_metadata`` only — +token usage lives in ``log``); ``llm_output`` is always ``None`` (parity with +LLMRails' unwired ``raw_response``). ``output_vars``/``state`` raise ``ValueError`` +and ``log.internal_events``/``log.colang_history`` raise ``NotImplementedError``. """ from unittest.mock import AsyncMock, patch @@ -66,27 +67,32 @@ def _stub_model(iorails: IORails, response: LLMResponse) -> None: _USER = [{"role": "user", "content": "hi"}] +_WEATHER_TOOL_CALL = ToolCall( + id="call_1", + type="function", + function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), +) -class TestStructuredResponseTrigger: - """``options`` presence decides GenerationResponse vs. bare dict.""" - @pytest.mark.asyncio - async def test_options_dict_returns_generation_response(self, iorails): - """A dict ``options`` argument switches the return type to GenerationResponse.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) +async def _generate_structured(iorails: IORails, response: LLMResponse, *, options=None, **kwargs): + """Stub safe rails + a fixed model response, then run the structured (``options``) path.""" + _stub_safe_rails(iorails) + _stub_model(iorails, response) + return await iorails.generate_async(_USER, options={} if options is None else options, **kwargs) - result = await iorails.generate_async(_USER, options={"llm_params": {"temperature": 0.5}}) - assert isinstance(result, GenerationResponse) +class TestStructuredResponseTrigger: + """``options`` presence decides GenerationResponse vs. bare dict.""" @pytest.mark.asyncio - async def test_generation_options_returns_generation_response(self, iorails): - """A GenerationOptions instance also switches the return type.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) - - result = await iorails.generate_async(_USER, options=GenerationOptions()) + @pytest.mark.parametrize( + "options", + [{"llm_params": {"temperature": 0.5}}, GenerationOptions()], + ids=["dict", "GenerationOptions"], + ) + async def test_options_returns_generation_response(self, iorails, options): + """Any ``options`` value (dict or GenerationOptions) switches the return type to GenerationResponse.""" + result = await _generate_structured(iorails, LLMResponse(content="Hello"), options=options) assert isinstance(result, GenerationResponse) @@ -108,11 +114,9 @@ class TestResponseField: @pytest.mark.asyncio async def test_plain_content_wrapped_in_list(self, iorails): """``response`` is ``[{"role":"assistant","content": text}]``.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello there")) - - result = await iorails.generate_async(_USER, options={}) + result = await _generate_structured(iorails, LLMResponse(content="Hello there")) + assert isinstance(result, GenerationResponse) assert result.response == [{"role": "assistant", "content": "Hello there"}] @@ -120,25 +124,17 @@ class TestReasoningContent: """Reasoning goes to ``reasoning_content`` with clean content (no inline ).""" @pytest.mark.asyncio - async def test_native_reasoning_in_field_content_clean(self, iorails): - """Provider ``reasoning`` populates ``reasoning_content``; ``response`` content has no prefix.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello", reasoning="thinking step")) - - result = await iorails.generate_async(_USER, options={}) - - assert isinstance(result, GenerationResponse) - assert result.reasoning_content == "thinking step" - assert result.response == [{"role": "assistant", "content": "Hello"}] - iorails.rails_manager.is_output_safe.assert_called_once_with(_USER, "Hello", enabled=True) - - @pytest.mark.asyncio - async def test_inline_think_tags_extracted_to_field(self, iorails): - """Inline tags are stripped into ``reasoning_content`` and never reach output rails.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="thinking stepHello")) - - result = await iorails.generate_async(_USER, options={}) + @pytest.mark.parametrize( + "response", + [ + LLMResponse(content="Hello", reasoning="thinking step"), + LLMResponse(content="thinking stepHello"), + ], + ids=["native-reasoning-field", "inline-think-tags"], + ) + async def test_reasoning_extracted_content_clean(self, iorails, response): + """Reasoning (native field or inline ) goes to ``reasoning_content``; response content stays clean and output rails see only the clean text.""" + result = await _generate_structured(iorails, response) assert isinstance(result, GenerationResponse) assert result.reasoning_content == "thinking step" @@ -148,10 +144,7 @@ async def test_inline_think_tags_extracted_to_field(self, iorails): @pytest.mark.asyncio async def test_no_reasoning_field_is_none(self, iorails): """Absent reasoning leaves ``reasoning_content`` as None.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="plain answer")) - - result = await iorails.generate_async(_USER, options={}) + result = await _generate_structured(iorails, LLMResponse(content="plain answer")) assert isinstance(result, GenerationResponse) assert result.reasoning_content is None @@ -164,15 +157,7 @@ class TestToolCalls: @pytest.mark.asyncio async def test_tool_calls_serialized_with_dict_arguments(self, iorails): """``tool_calls`` is a list of ``to_dict()`` entries whose ``arguments`` stay a dict, not a JSON string.""" - _stub_safe_rails(iorails) - tool_call = ToolCall( - id="call_1", - type="function", - function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), - ) - _stub_model(iorails, LLMResponse(content="", tool_calls=[tool_call])) - - result = await iorails.generate_async(_USER, options={}) + result = await _generate_structured(iorails, LLMResponse(content="", tool_calls=[_WEATHER_TOOL_CALL])) assert isinstance(result, GenerationResponse) assert result.tool_calls == [ @@ -186,15 +171,7 @@ async def test_tool_calls_serialized_with_dict_arguments(self, iorails): @pytest.mark.asyncio async def test_response_message_has_no_tool_calls_key(self, iorails): """In the structured path tool calls live only in the top-level field, not on the message.""" - _stub_safe_rails(iorails) - tool_call = ToolCall( - id="call_1", - type="function", - function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), - ) - _stub_model(iorails, LLMResponse(content="", tool_calls=[tool_call])) - - result = await iorails.generate_async(_USER, options={}) + result = await _generate_structured(iorails, LLMResponse(content="", tool_calls=[_WEATHER_TOOL_CALL])) assert isinstance(result, GenerationResponse) assert result.response == [{"role": "assistant", "content": ""}] @@ -203,10 +180,7 @@ async def test_response_message_has_no_tool_calls_key(self, iorails): @pytest.mark.asyncio async def test_no_tool_calls_field_is_none(self, iorails): """A text-only response leaves ``tool_calls`` as None.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) - - result = await iorails.generate_async(_USER, options={}) + result = await _generate_structured(iorails, LLMResponse(content="Hello")) assert isinstance(result, GenerationResponse) assert result.tool_calls is None @@ -216,47 +190,28 @@ class TestLLMMetadata: """``llm_metadata`` is the main-call ``provider_metadata`` verbatim; usage lives in ``log``.""" @pytest.mark.asyncio - async def test_provider_metadata_passthrough(self, iorails): - """provider_metadata is surfaced verbatim; token usage is NOT added under ``usage``.""" - _stub_safe_rails(iorails) - _stub_model( - iorails, - LLMResponse( - content="Hello", - provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}, - usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15), + @pytest.mark.parametrize( + "response, expected", + [ + ( + LLMResponse( + content="Hello", + provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}, + usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15), + ), + {"response_headers": {"nvcf-status": "fulfilled"}}, ), - ) - - result = await iorails.generate_async(_USER, options={}) + (LLMResponse(content="Hello", usage=UsageInfo(input_tokens=3, output_tokens=4, total_tokens=7)), None), + (LLMResponse(content="Hello"), None), + ], + ids=["provider-metadata-passthrough", "usage-only-none", "nothing-none"], + ) + async def test_llm_metadata_is_provider_metadata_only(self, iorails, response, expected): + """llm_metadata is the main-call provider_metadata verbatim (else None); token usage is never grafted under ``usage``.""" + result = await _generate_structured(iorails, response) assert isinstance(result, GenerationResponse) - assert result.llm_metadata == {"response_headers": {"nvcf-status": "fulfilled"}} - - @pytest.mark.asyncio - async def test_usage_alone_does_not_populate_metadata(self, iorails): - """With usage but no provider_metadata, ``llm_metadata`` is None — usage moved to log.""" - _stub_safe_rails(iorails) - _stub_model( - iorails, - LLMResponse(content="Hello", usage=UsageInfo(input_tokens=3, output_tokens=4, total_tokens=7)), - ) - - result = await iorails.generate_async(_USER, options={}) - - assert isinstance(result, GenerationResponse) - assert result.llm_metadata is None - - @pytest.mark.asyncio - async def test_no_metadata_or_usage_is_none(self, iorails): - """No provider_metadata and no usage leaves ``llm_metadata`` as None.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) - - result = await iorails.generate_async(_USER, options={}) - - assert isinstance(result, GenerationResponse) - assert result.llm_metadata is None + assert result.llm_metadata == expected class TestLLMOutput: @@ -265,38 +220,33 @@ class TestLLMOutput: @pytest.mark.asyncio async def test_llm_output_none_even_when_requested(self, iorails): """``options={"llm_output": True}`` is accepted but the field stays None.""" - _stub_safe_rails(iorails) - _stub_model( - iorails, - LLMResponse(content="Hello", provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}), - ) - - result = await iorails.generate_async(_USER, options={"llm_output": True}) + response = LLMResponse(content="Hello", provider_metadata={"response_headers": {"nvcf-status": "fulfilled"}}) + result = await _generate_structured(iorails, response, options={"llm_output": True}) assert isinstance(result, GenerationResponse) assert result.llm_output is None class TestUnsupportedOptionGuards: - """Colang-coupled options IORails cannot honor raise ValueError.""" - - @pytest.mark.asyncio - async def test_output_vars_true_raises(self, iorails): - """``output_vars=True`` requests Colang context IORails has no access to.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) - - with pytest.raises(ValueError): - await iorails.generate_async(_USER, options={"output_vars": True}) + """Colang-coupled options IORails cannot honor raise.""" @pytest.mark.asyncio - async def test_output_vars_list_raises(self, iorails): - """A list of ``output_vars`` also raises.""" + @pytest.mark.parametrize( + "gen_kwargs", + [ + {"options": {"output_vars": True}}, + {"options": {"output_vars": ["relevant_chunks"]}}, + {"options": {}, "state": {"conversation": []}}, + ], + ids=["output_vars-true", "output_vars-list", "state-arg"], + ) + async def test_colang_state_option_raises_value_error(self, iorails, gen_kwargs): + """``output_vars`` (any form) and ``state`` need Colang runtime context IORails lacks — ValueError.""" _stub_safe_rails(iorails) _stub_model(iorails, LLMResponse(content="Hello")) with pytest.raises(ValueError): - await iorails.generate_async(_USER, options={"output_vars": ["relevant_chunks"]}) + await iorails.generate_async(_USER, **gen_kwargs) @pytest.mark.asyncio async def test_colang_only_log_flag_raises(self, iorails): @@ -312,25 +262,23 @@ async def test_colang_only_log_flag_raises(self, iorails): with pytest.raises(NotImplementedError): await iorails.generate_async(_USER, options={"log": {"internal_events": True}}) - @pytest.mark.asyncio - async def test_state_argument_raises(self, iorails): - """A ``state`` argument is unsupported for the stateless IORails engine.""" - _stub_safe_rails(iorails) - _stub_model(iorails, LLMResponse(content="Hello")) - - with pytest.raises(ValueError): - await iorails.generate_async(_USER, options={}, state={"conversation": []}) - class TestBlockedStructuredResponse: """A blocked request returns the refusal in ``response`` with the other fields empty.""" @pytest.mark.asyncio - async def test_input_block_returns_refusal_response(self, iorails): - """Input-rail block yields a GenerationResponse whose response is the refusal message.""" - iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=False, reason="unsafe")) - iorails.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) - _stub_model(iorails, LLMResponse(content="unused")) + @pytest.mark.parametrize( + "input_safe, output_safe", [(False, True), (True, False)], ids=["input-block", "output-block"] + ) + async def test_block_returns_refusal_response(self, iorails, input_safe, output_safe): + """A block at either the input or output rail yields a GenerationResponse whose response is the refusal and whose other fields are empty.""" + iorails.rails_manager.is_input_safe = AsyncMock( + return_value=RailResult(is_safe=input_safe, reason=None if input_safe else "unsafe") + ) + iorails.rails_manager.is_output_safe = AsyncMock( + return_value=RailResult(is_safe=output_safe, reason=None if output_safe else "unsafe") + ) + _stub_model(iorails, LLMResponse(content="bad answer")) result = await iorails.generate_async(_USER, options={}) @@ -340,18 +288,6 @@ async def test_input_block_returns_refusal_response(self, iorails): assert result.reasoning_content is None assert result.llm_metadata is None - @pytest.mark.asyncio - async def test_output_block_returns_refusal_response(self, iorails): - """Output-rail block yields a GenerationResponse whose response is the refusal message.""" - 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=False, reason="unsafe")) - _stub_model(iorails, LLMResponse(content="bad answer")) - - result = await iorails.generate_async(_USER, options={}) - - assert isinstance(result, GenerationResponse) - assert result.response == [{"role": "assistant", "content": REFUSAL_MESSAGE}] - class TestBarePathUnchanged: """The optionless bare-dict path keeps its existing behavior.""" @@ -385,25 +321,17 @@ def test_sync_generate_with_options_returns_generation_response(self, iorails_sy class TestResponseContentForCapture: """`_response_content_for_capture` extracts assistant text from either return shape.""" - def test_generation_response_list(self): - """A structured (list) response yields the last assistant message's content.""" - result = GenerationResponse(response=[{"role": "assistant", "content": "hi there"}]) - assert _response_content_for_capture(result) == "hi there" - - def test_generation_response_str(self): - """A structured (str) response is returned as-is.""" - result = GenerationResponse(response="hi there") - assert _response_content_for_capture(result) == "hi there" - - def test_generation_response_empty_list(self): - """An empty response list has no content to capture.""" - assert _response_content_for_capture(GenerationResponse(response=[])) is None - - def test_generation_response_non_str_content(self): - """A non-string message content is not captured.""" - result = GenerationResponse(response=[{"role": "assistant", "content": None}]) - assert _response_content_for_capture(result) is None - - def test_bare_message_dict(self): - """The bare-dict return path reads content off the message directly.""" - assert _response_content_for_capture({"role": "assistant", "content": "hi there"}) == "hi there" + @pytest.mark.parametrize( + "result, expected", + [ + (GenerationResponse(response=[{"role": "assistant", "content": "hi there"}]), "hi there"), + (GenerationResponse(response="hi there"), "hi there"), + (GenerationResponse(response=[]), None), + (GenerationResponse(response=[{"role": "assistant", "content": None}]), None), + ({"role": "assistant", "content": "hi there"}, "hi there"), + ], + ids=["structured-list", "structured-str", "empty-list", "non-str-content", "bare-dict"], + ) + def test_capture_content(self, result, expected): + """Assistant content is pulled from structured list/str responses and from the bare-dict path; non-str/absent content yields None.""" + assert _response_content_for_capture(result) == expected From cdee6c3ddc9c787c5bd06e7612ea5665a633fddc Mon Sep 17 00:00:00 2001 From: tgasser-nv <200644301+tgasser-nv@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:40:35 -0500 Subject: [PATCH 11/11] Revert change to nemoguards config main model to meta/llama-3.3-70b-instruct --- examples/configs/nemoguards/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/configs/nemoguards/config.yml b/examples/configs/nemoguards/config.yml index 36af0ac8a9..878f8061a2 100644 --- a/examples/configs/nemoguards/config.yml +++ b/examples/configs/nemoguards/config.yml @@ -1,7 +1,7 @@ models: - type: main engine: nim - model: nvidia/nemotron-3-ultra-550b-a55b + model: meta/llama-3.3-70b-instruct - type: content_safety engine: nim