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/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/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.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/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index 546f66361c..9dd3825815 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,56 @@ 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 + 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 + prompt: Optional[str] = None + completion: 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. @@ -84,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 4b8df06120..aed5bb5ef8 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -37,8 +37,10 @@ from nemoguardrails.guardrails.guardrails_types import ( LLMMessage, LLMMessages, + RailCallRecord, RailDirection, get_request_id, + serialize_prompt, truncate, ) from nemoguardrails.guardrails.rails_manager import RailsManager @@ -60,10 +62,21 @@ 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 GenerationOptions, RailsResult, RailStatus, RailType +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 @@ -170,6 +183,225 @@ def _build_assistant_message(content: str, tool_calls: Optional[list[ToolCall]]) } +def _build_llm_metadata(response: LLMResponse) -> Optional[dict]: + """Return the main call's ``provider_metadata`` verbatim (LLMRails mirror). + + 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. + """ + return dict(response.provider_metadata or {}) or None + + +def _make_generation_record( + response: LLMResponse, + started_at: Optional[float], + 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. ``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", + 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, + prompt=prompt, + completion=response.content, + 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/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: + """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, + 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", + ) + + +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), + ) + + +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). + """ + 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, + log: Optional[GenerationLog] = None, +) -> 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. ``log`` is set when the caller requested + log details via ``options.log``. + """ + 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 + if log is not None: + result.log = log + return result + + +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 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]: + """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.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") + 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.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]: """Normalize the request ``options`` argument into a ``GenerationOptions`` or None.""" if isinstance(options, GenerationOptions): @@ -464,7 +696,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) -> LLMMessage: + 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 @@ -484,11 +721,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) -> LLMMessage: + 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 @@ -508,13 +750,18 @@ async def generate_async(self, messages: LLMMessages, **kwargs) -> LLMMessage: 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) -> LLMMessage: + 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 → @@ -525,7 +772,7 @@ async def _run_generate(self, messages: LLMMessages, **kwargs) -> LLMMessage: 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) @@ -533,7 +780,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 @@ -559,13 +806,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 - ) -> LLMMessage: + 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_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 @@ -573,25 +830,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 {"role": "assistant", "content": REFUSAL_MESSAGE} + 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 {"role": "assistant", "content": REFUSAL_MESSAGE} + 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)) @@ -608,11 +880,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 {"role": "assistant", "content": REFUSAL_MESSAGE} + 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 @@ -623,25 +896,39 @@ 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 {"role": "assistant", "content": REFUSAL_MESSAGE} + return _blocked_return() - # TODO: Support returning GenerationResponse `reasoning_content` to match LLMRails - # For now, embed the reasoning on the content with think-tags + if has_generation_options: + 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 + # reasoning_content field and keeps the message content clean. if reasoning_content: response_text = f"{reasoning_content}\n" + response_text 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: @@ -649,7 +936,16 @@ 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.time() + t0 = time.monotonic() + response = await self.engine_registry.model_call("main", messages, **llm_kwargs) + duration = time.monotonic() - t0 + finished_at = time.time() + if records_out is not None: + provider = self.engine_registry.provider_name("main") + 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( self, @@ -659,6 +955,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) @@ -668,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 + 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): @@ -701,6 +1003,9 @@ 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, + 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) @@ -713,6 +1018,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) @@ -738,6 +1045,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) @@ -751,6 +1060,11 @@ 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 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, 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 5d4a74ade1..761955c6a6 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 @@ -31,12 +34,13 @@ LLMMessages, RailResult, get_request_id, + serialize_prompt, truncate, ) 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 +48,40 @@ log = logging.getLogger(__name__) +@dataclass(frozen=True) +class RailLLMCall: + """Usage/model/timing captured for the call a rail made, for GenerationLog. + + 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] + prompt: Optional[str] + completion: 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 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) + return call + + class RailAction(ABC): """Base class for all IORails rail actions. @@ -81,6 +119,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 +194,32 @@ 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.time() + t0 = time.monotonic() + response = await self.engine_registry.model_call(model_type, messages, **kwargs) + 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), + prompt=serialize_prompt(messages), + completion=response.content, + started_at=started_at, + finished_at=finished_at, + duration=duration, + ) + ) + return response async def _get_api_response( self, @@ -164,8 +227,31 @@ 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, + prompt=None, + completion=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 d296289ff1..e6b5dd6e75 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -37,16 +37,17 @@ 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, 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 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: @@ -79,6 +80,39 @@ _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} + # 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=action_name, + return_value=verdict, + 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, + 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, + ) + + class RailsManager: """Orchestrates input and output safety checks for IORails. @@ -304,8 +338,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 = 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),)) 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 +360,9 @@ 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). + 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( @@ -346,6 +385,9 @@ 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). + 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] @@ -368,14 +410,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 +435,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 +457,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_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, 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_log.py b/tests/guardrails/test_iorails_generation_log.py new file mode 100644 index 0000000000..aedd07c2d6 --- /dev/null +++ b/tests/guardrails/test_iorails_generation_log.py @@ -0,0 +1,330 @@ +# 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, _activated_rail +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 + + +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} + + 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 + + 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 new file mode 100644 index 0000000000..9dfd979569 --- /dev/null +++ b/tests/guardrails/test_iorails_generation_log_capture.py @@ -0,0 +1,194 @@ +# 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"}) + + +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).""" + 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, + request_id="req-cs-input", + ) + ) + + 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.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": []} + + +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.""" + 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}}) + + 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": []} + # 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 + 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.""" + iorails.engine_registry.model_call = AsyncMock(side_effect=_cs_and_main_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.""" + 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 new file mode 100644 index 0000000000..2586ee9a55 --- /dev/null +++ b/tests/guardrails/test_iorails_generation_response.py @@ -0,0 +1,337 @@ +# 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), ``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 + +import pytest +import pytest_asyncio + +from nemoguardrails.guardrails.guardrails_types import RailResult +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 +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"}] + +_WEATHER_TOOL_CALL = ToolCall( + id="call_1", + type="function", + function=ToolCallFunction(name="get_weather", arguments={"city": "SF"}), +) + + +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) + + +class TestStructuredResponseTrigger: + """``options`` presence decides GenerationResponse vs. bare dict.""" + + @pytest.mark.asyncio + @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) + + @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}]``.""" + result = await _generate_structured(iorails, LLMResponse(content="Hello there")) + + assert isinstance(result, GenerationResponse) + assert result.response == [{"role": "assistant", "content": "Hello there"}] + + +class TestReasoningContent: + """Reasoning goes to ``reasoning_content`` with clean content (no inline ).""" + + @pytest.mark.asyncio + @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" + 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.""" + result = await _generate_structured(iorails, LLMResponse(content="plain answer")) + + 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.""" + result = await _generate_structured(iorails, LLMResponse(content="", tool_calls=[_WEATHER_TOOL_CALL])) + + 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.""" + result = await _generate_structured(iorails, LLMResponse(content="", tool_calls=[_WEATHER_TOOL_CALL])) + + 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.""" + result = await _generate_structured(iorails, LLMResponse(content="Hello")) + + assert isinstance(result, GenerationResponse) + assert result.tool_calls is None + + +class TestLLMMetadata: + """``llm_metadata`` is the main-call ``provider_metadata`` verbatim; usage lives in ``log``.""" + + @pytest.mark.asyncio + @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"}}, + ), + (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 == expected + + +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.""" + 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.""" + + @pytest.mark.asyncio + @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, **gen_kwargs) + + @pytest.mark.asyncio + 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": {"internal_events": True}}) + + +class TestBlockedStructuredResponse: + """A blocked request returns the refusal in ``response`` with the other fields empty.""" + + @pytest.mark.asyncio + @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={}) + + 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 + + +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) + + +class TestResponseContentForCapture: + """`_response_content_for_capture` extracts assistant text from either return shape.""" + + @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 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_rails_manager.py b/tests/guardrails/test_rails_manager.py index f3957d0697..7bcd92f17b 100644 --- a/tests/guardrails/test_rails_manager.py +++ b/tests/guardrails/test_rails_manager.py @@ -31,8 +31,8 @@ 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.rails_manager import RailsManager +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 from nemoguardrails.tracing.constants import GuardrailsAttributes @@ -971,3 +971,57 @@ 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" + + +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: " 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 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):