diff --git a/docs/configure-rails/guardrail-catalog/tool-calling.mdx b/docs/configure-rails/guardrail-catalog/tool-calling.mdx index bbd94c677f..58028acf36 100644 --- a/docs/configure-rails/guardrail-catalog/tool-calling.mdx +++ b/docs/configure-rails/guardrail-catalog/tool-calling.mdx @@ -280,12 +280,14 @@ models: ## Tool Call Validation -The `tool call validation` rail inspects every tool call in the model's response and blocks the request on the first violation. +The `tool call validation` rail inspects every tool call in the model's response and blocks the request when any call breaks a rule. - **Allowlist.** The call must name a tool you declared. IORails blocks a call to a tool that is not in the declared set, for example `tool call 'delete_database' is not an allowed tool`. - **Argument schema.** The call's arguments must validate against the tool's declared JSON Schema (the `function.parameters` block). - IORails blocks arguments that violate the schema, for example `arguments for tool 'get_weather' do not match its schema: 'city' is a required property`. + IORails blocks arguments that violate the schema, for example `arguments for tool 'get_weather' do not match its schema: 'required' failed at '/city'`. + The reason names the failing argument as the schema declares it, and the JSON Schema keyword. + It never quotes an argument's value or a key the model chose, which appears as `*`. - **No-argument tools.** A function tool that declares no parameters must receive no arguments. IORails blocks any supplied argument. - **Hosted tools.** IORails allowlists a hosted or server-side tool that only `type` identifies (for example a built-in web search) by type, and does not schema-validate its arguments because the provider owns the call shape. @@ -343,6 +345,65 @@ response = await rails.generate_async( ) ``` +## Check Tool Calls and Results Without Generation + +An agent harness that calls its own model can use the same rails to validate tool traffic without IORails generating a response. +Run a check with `RailType.TOOL_CALL` on the conversation that ends in the model's tool calls, and with `RailType.TOOL_RESULT` after appending the tool results: + +```python +from nemoguardrails import Guardrails, RailsConfig +from nemoguardrails.rails.llm.options import RailStatus, RailType + +rails = Guardrails(RailsConfig.from_path("./config"), require_iorails=True) + +call_check = await rails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=tools) +if call_check.status == RailStatus.BLOCKED: + ... # Drop or re-request the calls named in call_check.tool_violations. + +messages.append({"role": "tool", "tool_call_id": "call_1", "content": "18C"}) +result_check = await rails.check_async(messages, rail_types=[RailType.TOOL_RESULT]) +``` + +A `tool_call` check validates the tool calls on the last assistant message against `tools`, which replace the tools declared on the main model. +Only a `tool_call` check reads `tools`, so passing them to any other check raises `InvalidCheckRequestError`, or returns HTTP 422 from the server. +The allowlist and argument schemas are enforced for every call by `tool call validation`, and for the calls a `per_tool` rail checks by that rail; `per_tool` rails check only the tools they are listed under. +A `tool_result` check validates every `tool` message against the calls it answers. +Both run the global validator and any `per_tool` rails, and neither calls the model. +A blocked check reports each failing call or result in `tool_violations`, with a `violation_type` to switch on. +For the full behavior, the violation types, and the engine requirements, refer to [Checking Messages Against Rails: Tool Rails](/run-guardrailed-inference/using-python-apis/check-messages#tool-rails-iorails-only). + +The server exposes the same check through `/v1/checks`. +Run the server with `NEMO_GUARDRAILS_IORAILS_ENGINE=1` so the configuration is served by IORails; on `LLMRails`, the tool rail types return HTTP 422. +Pass the tools in the request's top-level `tools` field: + +```json +{ + "model": "meta/llama-3.3-70b-instruct", + "messages": [ + {"role": "user", "content": "What's the weather in Paris?"}, + {"role": "assistant", "content": null, "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "delete_files", "arguments": "{}"}} + ]} + ], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}], + "guardrails": {"config_id": "tools", "rail_types": ["tool_call"]} +} +``` + +The call names a tool that is not declared, so the response is: + +```json +{ + "status": "blocked", + "content": "I'm sorry, I can't respond to that.", + "rail": "tool call validation", + "reason": "tool call 'delete_files' is not an allowed tool", + "tool_violations": [ + {"kind": "tool_call", "violation_type": "tool_not_allowed", "reason": "tool call 'delete_files' is not an allowed tool", "tool_call_id": "call_1", "tool_name": "delete_files", "index": 0} + ] +} +``` + ## Streaming Tool-calling rails work with `stream_async`. diff --git a/docs/reference/engine-feature-support.mdx b/docs/reference/engine-feature-support.mdx index 3d4df14b64..f400fd7acb 100644 --- a/docs/reference/engine-feature-support.mdx +++ b/docs/reference/engine-feature-support.mdx @@ -147,11 +147,13 @@ With `require_iorails=True`, `Guardrails` raises `ValueError` instead of falling | Feature | LLMRails | IORails | Notes | |---|:---:|:---:|---| | Tool-call passthrough | ✓ | ✓ | | -| Tool-call validation rail | ✓ | ✓ | IORails flow: `tool call validation` | -| Tool-result validation rail | ✓ | ✓ | IORails flow: `tool result validation` | +| Tool-call validation rail | ✗ | ✓ | IORails flow: `tool call validation` | +| Tool-result validation rail | ✗ | ✓ | IORails flow: `tool result validation` | +| Tool rails in `check` / `check_async` | ✗ | ✓ | `RailType.TOOL_CALL` and `RailType.TOOL_RESULT`. LLMRails raises `RailTypeNotSupportedError` | -Both engines support passing model tool calls through to the caller and validating tool calls and tool results. -`LLMRails` handles these through the Colang runtime and tool rails. +Both engines support passing model tool calls through to the caller. +`LLMRails` runs the Colang flows you configure under `rails.tool_output` and `rails.tool_input` through the Colang runtime. +The built-in `tool call validation` and `tool result validation` validators run on `IORails` only. `IORails` validates tool calls and tool results through directional flows: `tool call validation` on the tool-output rail and `tool result validation` on the tool-input rail. Tool calls are returned in the OpenAI-style `tool_calls` field of the response message. @@ -163,7 +165,7 @@ Tool calls are returned in the OpenAI-style `tool_calls` field of the response m | `generate` / `generate_async` | ✓ | ✓ | | | `stream_async` | ✓ | ✓ | | | Event-based API (`generate_events` / `process_events`) | ✓ | ✗ | Requires the Colang runtime | -| `check` / `check_async` (rails-only validation) | ✓ | ✓ | Input and output rails only on both engines. Neither runs tool rails | +| `check` / `check_async` (rails-only validation) | ✓ | ✓ | Input and output rails on both engines. Only IORails runs tool rails, when they are named in `rail_types` | | `GenerationOptions` | ✓ | ◐ | IORails supports rail toggles, `llm_params`, and the `log` options that do not need Colang. `output_vars` raises. The separate `state` keyword argument also raises. | | `GenerationResponse` (structured response object) | ✓ | ✓ | Returned by both engines when `options` is passed | | `GenerationLog` `internal_events` / `colang_history` | ✓ | ✗ | Colang runtime only. IORails raises `NotImplementedError`. | @@ -187,7 +189,9 @@ For the field-by-field comparison, refer to [Generation Options: IORails and LLM The event-based API and `explain()` are not available on `IORails`. On the `Guardrails` facade, these raise `NotImplementedError` when `IORails` is the active engine. -Both engines accept the same `check` arguments and return the same `RailsResult`, and neither runs tool rails from a check. +Both engines accept `messages` and `rail_types` and return the same `RailsResult`. +Only `IORails` runs tool rails from a check, when `rail_types` names `tool_call` or `tool_result`. +`LLMRails` raises `RailTypeNotSupportedError` for those rail types and for the `tools` argument. `LLMRails` runs the check through the Colang runtime with the main model call disabled. `IORails` runs the rails directly through its rails manager and reuses the `generate_async` admission queue, request span, and request metrics. The two differ in what a blocked check reports and in how a direction with no content to check is handled. diff --git a/docs/reference/rail-engine-support.mdx b/docs/reference/rail-engine-support.mdx index e6e1a90312..db96b2e6a0 100644 --- a/docs/reference/rail-engine-support.mdx +++ b/docs/reference/rail-engine-support.mdx @@ -182,8 +182,8 @@ Tool rails are not manifest surfaces. | Flow | Section | LLMRails | IORails | Notes | |---|---|:---:|:---:|---| -| `tool call validation` | `rails.tool_output.flows` | ✓ | ✓ | Validates model-emitted tool calls | -| `tool result validation` | `rails.tool_input.flows` | ✓ | ✓ | Validates application-returned tool results | +| `tool call validation` | `rails.tool_output.flows` | ✗ | ✓ | Validates model-emitted tool calls, in generation and in a `tool_call` check | +| `tool result validation` | `rails.tool_input.flows` | ✗ | ✓ | Validates application-returned tool results, in generation and in a `tool_result` check | For configuration and behavior, refer to [Tool Calling](/configure-guardrails/guardrail-catalog/tool-calling). diff --git a/docs/run-rails/using-python-apis/check-messages.mdx b/docs/run-rails/using-python-apis/check-messages.mdx index 1d79bba315..1abb1d44df 100644 --- a/docs/run-rails/using-python-apis/check-messages.mdx +++ b/docs/run-rails/using-python-apis/check-messages.mdx @@ -3,8 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 title: "Checking Messages Against Rails" sidebar-title: "Check Messages" -description: "Validate messages against input and output rails using check_async and check methods on the LLMRails and IORails engines." -keywords: ["check_async", "check", "RailsResult", "RailStatus", "RailType", "input rails", "output rails", "validation", "IORails", "LLMRails"] +description: "Validate messages against input, output, and tool rails using check_async and check methods on the LLMRails and IORails engines." +keywords: ["check_async", "check", "RailsResult", "RailStatus", "RailType", "input rails", "output rails", "tool rails", "ToolViolation", "validation", "IORails", "LLMRails"] content: type: "reference" --- @@ -13,7 +13,9 @@ The `check_async()` and `check()` methods validate messages against input and ou Both the `LLMRails` and `IORails` engines implement these methods. The `Guardrails` facade forwards each call to the selected engine. -The method signatures and returned `RailsResult` type are the same on both engines. +Both engines take `messages` and `rail_types` and return a `RailsResult`. +`IORails` also runs tool-call and tool-result rails in a check, and takes a `tools` argument for them. +For more information, refer to [Tool Rails (IORails Only)](#tool-rails-iorails-only). For engine-specific rail selection and execution behavior, refer to [Engine Differences](#engine-differences). ## Method Signatures @@ -28,6 +30,8 @@ The primary asynchronous method for checking messages against rails. async def check_async( messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, # IORails and the Guardrails facade only ) -> RailsResult ``` @@ -39,6 +43,8 @@ Synchronous wrapper around `check_async()`. def check( messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, # IORails and the Guardrails facade only ) -> RailsResult ``` @@ -48,6 +54,7 @@ def check( |-----------|------|-------------| | `messages` | `List[dict]` | List of message dictionaries with `role` and `content` fields | | `rail_types` | `Optional[List[RailType]]` | Optional list of rail types to run. When provided, overrides automatic detection based on message roles. | +| `tools` | `Optional[List[dict]]` | IORails only; `LLMRails` raises `RailTypeNotSupportedError`. The tools a `tool_call` check validates calls against, in the main model's wire format. Only a `tool_call` check reads them, so passing `tools` without `RailType.TOOL_CALL` in `rail_types` raises `InvalidCheckRequestError`. Refer to [Tool Rails (IORails Only)](#tool-rails-iorails-only). | **Returns:** `RailsResult` object containing validation results. @@ -95,6 +102,8 @@ result = await rails.check_async( |-------|-------------| | `RailType.INPUT` | Run input rails | | `RailType.OUTPUT` | Run output rails | +| `RailType.TOOL_CALL` | Run tool-call rails on the last assistant message's tool calls (IORails only) | +| `RailType.TOOL_RESULT` | Run tool-result rails on the `tool` messages (IORails only) | On either engine, passing an empty list (`rail_types=[]`) runs no rails and returns a passed result. @@ -107,12 +116,83 @@ input flows will raise an error instead of silently returning a passed result. -### Tool Rails Are Not Included +### Tool Rails (IORails Only) -Neither engine runs tool-call or tool-result rails from `check()` or `check_async()`. -`RailType` covers input and output rails only, so a configuration with tool rails still validates only its input and output flows here. -To validate tool traffic, run it through `generate_async()`. -For more information, refer to [Tool Calling](/configure-guardrails/guardrail-catalog/tool-calling). +`IORails` runs tool rails in a check when you name their rail types in `rail_types`. +This lets an agent harness validate the tool calls its model produced, and the tool results it is about to send back, without generating a response. + +| Rail type | What the check validates | +|---|---| +| `RailType.TOOL_CALL` | The tool calls on the last assistant message. With `tool call validation` configured, each call must name a declared tool, and its arguments must match that tool's JSON Schema. | +| `RailType.TOOL_RESULT` | Every `tool` message. Each result must link to exactly one earlier tool call by `tool_call_id`, must not name a different tool than that call, and must carry well-formed content. | + +A tool check runs the rails that generation runs for the same direction: the global validator (`tool call validation` or `tool result validation`) and any `per_tool` rails. +The global validator enforces the declared-tool allowlist and argument schemas for every call. +A `per_tool` rail on tool calls enforces them too, but only for the calls it checks, before its own check runs. +A `per_tool` rail checks only the tools it is listed under, so with `per_tool` rails alone, a call to any other tool, declared or not, passes. +A tool check never calls the main model. + +Pass the tools your model was offered in `tools`, in the main model's wire format. +For the `openai` and `nim` engines, that is the OpenAI Chat Completions tool shape: + +```python +from nemoguardrails import Guardrails, RailsConfig +from nemoguardrails.rails.llm.options import RailStatus, RailType + +config = RailsConfig.from_path("path/to/config") +rails = Guardrails(config, use_iorails=True, require_iorails=True) + +weather_tool = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +messages = [ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"}} + ], + }, +] + +result = await rails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[weather_tool]) +if result.status == RailStatus.BLOCKED: + for violation in result.tool_violations: + print(f"{violation.violation_type.value}: call {violation.tool_call_id} at index {violation.index}") +``` + +Keep the following behavior in mind: + +- Tool rails run only when you name them. + When `rail_types` is omitted, a check runs input and output rails as described in [Automatic Detection](#automatic-detection-default). + It logs an INFO message when the configuration has tool rails and the messages carry tool calls or tool results. +- `tools` replaces the tools declared in the main model's `parameters.tools`, schemas included; it does not add to them. + When you omit `tools`, the configured tools are the allowlist. + `tools=[]` declares no tools, so `tool call validation` blocks every call. +- Only a `tool_call` check reads `tools`. + Passing `tools` to any other check, including one that leaves `rail_types` to automatic detection, raises `InvalidCheckRequestError` rather than ignoring them. +- A `tool_call` check validates only the last assistant message. + Earlier calls were validated when they were produced, and a `tool_result` check validates how results link to them. +- A tool check needs a `main` model, because its engine selects the wire format. + Without one, the check raises `RailTypeNotConfiguredError`. +- A tool rail type counts as configured when its configuration section lists `flows` or `per_tool` rails: `rails.tool_output` for `tool_call`, and `rails.tool_input` for `tool_result`. + Otherwise the check raises `RailTypeNotConfiguredError`. +- When you name several rail types, they run in the order generation runs them: `tool_result`, `input`, `tool_call`, then `output`. + The check stops at the first block. +- Tool checks fail closed. + A call whose arguments are not a JSON object blocks, and so does a legacy `function_call` on the last assistant message. + The other calls on that message are still validated, so `tool_violations` lists every failing call in call order, not only the one that could not be parsed. +- Legacy `role: "function"` messages are not validated by the tool-result rails, in checks or in generation. +- The configuration sections, and the `rail_type` recorded for tool rails in a generation's log, keep the names `tool_output` and `tool_input`. +- `LLMRails` does not run tool rails in a check. + Requesting a tool rail type, or passing `tools`, raises `RailTypeNotSupportedError` rather than ignoring the request. + +For configuration and the validation rules, refer to [Tool Calling](/configure-guardrails/guardrail-catalog/tool-calling). ## RailsResult @@ -124,6 +204,7 @@ The `RailsResult` object contains the outcome of the rails check. | `content` | `str` | The final content after rails processing | | `rail` | `Optional[str]` | Name of the rail that blocked the content (only when `BLOCKED`) | | `reason` | `Optional[str]` | Why the rail blocked the content (IORails only; `None` on `LLMRails` and when not `BLOCKED`) | +| `tool_violations` | `Optional[List[ToolViolation]]` | The tool calls or tool results a tool rail blocked (IORails only; `None` unless a tool rail blocked). Refer to [Tool Violations](#tool-violations) | ### RailStatus Enum @@ -135,15 +216,67 @@ The `RailStatus` enum represents the three possible outcomes of a rails check. | `MODIFIED` | Content was modified by rails but not blocked | | `BLOCKED` | Content was blocked by a rail | +### Tool Violations + +When a tool rail blocks, `tool_violations` lists each tool call or tool result that failed, so an agent harness can drop, repair, or re-request that item. +The global validators check every item, so one check can report several violations, and `reason` repeats the first one's reason. +When a tool rail fails rather than deciding, the check returns the internal-error message, and `reason` is that rail's failure even when an earlier item has its own violation. +Each tool call or tool result reports at most one violation: the first check it fails. +Switch on `violation_type`, which is a closed set, and never on the `reason` text, which can change between releases. + +| Field | Type | Description | +|---|---|---| +| `kind` | `str` | `tool_call` or `tool_result` | +| `violation_type` | `ToolViolationType` | What the call or result violated. Refer to the following table | +| `reason` | `str` | Why it was blocked. Never contains argument values or tool-result content. A tool name or call id it quotes is cut to 64 characters, followed by `...` | +| `tool_call_id` | `Optional[str]` | The call's id, or the id a result links to. Always whole, so you can match the violation to its call | +| `tool_name` | `Optional[str]` | The tool called. For a result, taken from the call it links to, never from the result's own `name`. A name longer than 64 characters is cut to 64, followed by `...` | +| `index` | `Optional[int]` | For a call, its position in the last assistant message's `tool_calls`. For a result, the position of its message in `messages` | +| `argument_path` | `Optional[str]` | For `arguments_invalid`, a JSON Pointer to the failing argument, such as `/city`. A key the schema does not declare appears as `*`. An empty string points at the arguments object itself | +| `schema_keyword` | `Optional[str]` | For `arguments_invalid`, the JSON Schema keyword that failed, such as `required` | +| `rail` | `Optional[str]` | For `per_tool_rail` and `rail_failed`, the rail that blocked or failed. Also set when a `per_tool` rail's own allowlist or schema check blocked a call, naming that rail | + +| `violation_type` | Kind | Meaning | +|---|---|---| +| `tool_not_allowed` | Call | The call names a tool that is not declared | +| `arguments_invalid` | Call | The arguments do not match the tool's JSON Schema | +| `unexpected_arguments` | Call | The tool declares no parameters, but the call passes arguments | +| `invalid_tool_schema` | Call | The tool's declared schema is not valid JSON Schema | +| `malformed_arguments` | Call | The arguments are not a JSON object | +| `malformed_tool_call` | Call | The tool-call entry is not an object, has no `function` object, or has an `id`, `type`, or function name that is not a string. An empty function name counts as malformed | +| `legacy_function_call` | Call | The last assistant message carries a legacy `function_call`, which tool rails do not validate | +| `invalid_toolset` | Call | The declared tools cannot be used, for example because a tool is declared twice | +| `missing_call_id` | Result | The result has no `tool_call_id` | +| `unknown_call_id` | Result | No earlier call in the same turn has the result's `tool_call_id` | +| `duplicate_result` | Result | An earlier result already answers the same call | +| `duplicate_prior_call_id` | Result | Several calls in a turn share an id, so that turn's results cannot be linked and are not checked further | +| `name_mismatch` | Result | The result names a different tool than the call it answers | +| `malformed_content` | Result | The content is neither a string nor a list of content blocks | +| `unlinkable_result` | Result | A `per_tool` rail applies, but the result cannot be linked to exactly one call | +| `extraction_failed` | Either | The messages could not be read in the main model's wire format | +| `rail_failed` | Either | A tool rail raised an error instead of deciding. The check returns the internal-error message | +| `per_tool_rail` | Either | A `per_tool` rail blocked the call or result. When a `per_tool` rail blocks a call for an undeclared tool or schema-invalid arguments, the violation carries that underlying type instead, such as `tool_not_allowed` or `arguments_invalid`, so a call is categorized the same way whichever rails are configured | + +An `argument_path` names only what the schema declares: property names from its `properties`, and array indices. +For a missing required property, it names that property from the schema. +A key the model chose appears as `*`, for example a key inside a map-shaped argument such as `/contacts/*`. +An unexpected property rejected by `additionalProperties: false` is reported at the object that holds it, without its name. +So neither `argument_path` nor `reason` contains an argument value or a key the model wrote. +For the same reason, `unexpected_arguments` reports how many arguments the call supplied, not their names. + ### Which Content Is Reported The `content` field reports the direction the caller asked about. A check that runs output rails reports the assistant text. A check that runs input rails only reports the user text. +A `tool_call` check reports the assistant text, unless input rails run in the same check, in which case it reports the user text so that an input rewrite stays visible. +A `tool_result` check reports the user text. + Both engines compare the reported text against its original value to decide between `PASSED` and `MODIFIED`. When a check runs both directions and only the user message is rewritten, the status is `PASSED`, because the reported assistant text is unchanged. `MODIFIED` names no rail: a rewrite is not a rail blocking the request, and several rails can contribute to the final text. +Tool rails never rewrite, so a check that runs only tool rails is `PASSED` or `BLOCKED`. ## Usage Examples @@ -229,14 +362,14 @@ For more information, refer to [Engine Feature Support](/reference/engine-featur ## Engine Differences -The API surface is identical on both engines, but each engine executes a check differently: +Both engines take `messages` and `rail_types`, but each engine executes a check differently: | Behavior | LLMRails | IORails | |---|---|---| | How the check runs | Through the Colang runtime, as a generation with dialog rails and the unrequested rail types disabled | Directly through the engine's rails manager, with no Colang runtime | | Rails that can run | Every configured input and output flow, including community, third-party, and custom-action rails | The input and output rails IORails supports. A configuration outside that set routes to `LLMRails` | | Main LLM call | Not made | Not made | -| Tool rails | Not run | Not run | +| Tool rails | Not run. Requesting a tool rail type, or passing `tools`, raises `RailTypeNotSupportedError` | Run when named in `rail_types`. Refer to [Tool Rails (IORails Only)](#tool-rails-iorails-only) | | `BLOCKED` content | The refusal message defined by the configuration's Colang flow | A fixed refusal message that the configuration does not change | | `BLOCKED` reason | Not reported. `reason` is `None` | The blocking rail's own explanation, or the rail's name when it gives none | | Rail that fails to execute | Handled by the Colang runtime's internal-error path | `BLOCKED` with an internal-error message that is distinct from the refusal message | @@ -290,6 +423,7 @@ When your application requires that a rail actually ran, assert that the message For more information, refer to [Span Reference](/observability/tracing/span-reference#iorails). - Metrics record the check in `guardrails.requests`, `guardrails.requests.active`, `guardrails.request.duration`, and `guardrails.requests.errors`. A block also updates `guardrails.requests.blocked` with the `rail.type` label. + A tool-result block counts as an `Input` block and a tool-call block as an `Output` block, as in generation. A check rejected by a full admission queue increments `guardrails.nonstream.rejections`. For more information, refer to [Metric Reference](/observability/metrics/reference). - Content capture records the checked messages in `guardrails.request.input` and the returned content in `guardrails.request.output` on the check's request span. @@ -304,9 +438,17 @@ The server resolves that engine through the top-level `LLMRails` import. Setting `NEMO_GUARDRAILS_IORAILS_ENGINE=1` runs compatible configurations on `IORails` and the rest on `LLMRails`. Requesting a rail type with no configured flows returns HTTP 422. +`guardrails.rail_types` accepts `tool_call` and `tool_result` on a configuration that `IORails` serves. +On a configuration that `LLMRails` serves, they return HTTP 422. +The request's top-level `tools` field sets the tools a `tool_call` check validates against. +On a configuration that `IORails` serves, a request that carries `tools` without `tool_call` in `guardrails.rail_types` returns HTTP 422. +On a configuration that `LLMRails` serves, a request that carries `tools` also returns HTTP 422. + The response carries `status`, `content`, and, on a block, `rail`. -A block on `IORails` also carries `reason`. -The response omits a field whose value is `None`, so a passed check and an `LLMRails` block carry no `reason`. +A block on `IORails` also carries `reason`, and a tool rail block also carries `tool_violations`. +When the request sets `guardrails.context`, the server prepends a context message, and each violation's `index` still counts positions in the request's own `messages`. +The response omits a field whose value is `None`, so a passed check and an `LLMRails` block carry no `reason` and no `tool_violations`. +For a request and response example, refer to [Tool Calling: Check Tool Calls and Results Without Generation](/configure-guardrails/guardrail-catalog/tool-calling#check-tool-calls-and-results-without-generation). diff --git a/docs/run-rails/using-python-apis/core-classes.mdx b/docs/run-rails/using-python-apis/core-classes.mdx index 4ba27f7337..ef3b2e4bab 100644 --- a/docs/run-rails/using-python-apis/core-classes.mdx +++ b/docs/run-rails/using-python-apis/core-classes.mdx @@ -353,6 +353,15 @@ result = await rails.check_async( The `LLMRails` and `IORails` engines both implement these methods, and the `Guardrails` facade forwards the call to the engine it selected. +On `IORails`, a check can also validate tool traffic without generation. +Name `RailType.TOOL_CALL` or `RailType.TOOL_RESULT` in `rail_types`, and pass the tools the model was offered in `tools`: + +```python +result = await rails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[weather_tool]) +``` + +`LLMRails` does not run tool rails in a check, and raises `RailTypeNotSupportedError` for these rail types. + For detailed usage, method signatures, engine differences, and examples, refer to [Check Messages](/run-guardrailed-inference/using-python-apis/check-messages). --- diff --git a/nemoguardrails/exceptions.py b/nemoguardrails/exceptions.py index 487bac312d..bcf396db0b 100644 --- a/nemoguardrails/exceptions.py +++ b/nemoguardrails/exceptions.py @@ -19,6 +19,7 @@ "ConfigurationError", "InvalidModelConfigurationError", "InvalidRailsConfigurationError", + "InvalidCheckRequestError", "InvalidStateError", "LLMCallException", "LLMClientError", @@ -35,6 +36,7 @@ "StreamingCapacityExceededError", "StreamingNotSupportedError", "RailTypeNotConfiguredError", + "RailTypeNotSupportedError", ] @@ -77,6 +79,18 @@ class RailTypeNotConfiguredError(InvalidRailsConfigurationError): pass +class RailTypeNotSupportedError(InvalidRailsConfigurationError): + """Raised when an explicitly requested rail type cannot run on the serving engine.""" + + pass + + +class InvalidCheckRequestError(ValueError): + """Raised when a check's arguments contradict each other, such as `tools` without a `tool_call` check.""" + + pass + + class InvalidStateError(ValueError): """Raised when a caller-supplied `state` argument is not valid public input. diff --git a/nemoguardrails/guardrails/actions/tool_call_action.py b/nemoguardrails/guardrails/actions/tool_call_action.py index 42ce16e034..9aa0da4930 100644 --- a/nemoguardrails/guardrails/actions/tool_call_action.py +++ b/nemoguardrails/guardrails/actions/tool_call_action.py @@ -27,8 +27,10 @@ from typing import TYPE_CHECKING, List from nemoguardrails.actions.rail_outcome import RailOutcome -from nemoguardrails.guardrails.tool_rail_action import ToolRailAction +from nemoguardrails.guardrails.guardrails_types import quoted_identity +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction, violations_outcome from nemoguardrails.guardrails.tool_schema import validate_arguments +from nemoguardrails.rails.llm.options import ToolViolation, ToolViolationType if TYPE_CHECKING: from nemoguardrails.guardrails.tool_schema import Toolset @@ -45,15 +47,39 @@ async def run(self, toolset: "Toolset", tool_calls: List["ToolCall"]) -> RailOut return self._guarded(lambda: self._validate(toolset, tool_calls)) def _validate(self, toolset: "Toolset", tool_calls: List["ToolCall"]) -> RailOutcome: - """Allowlist each call by name, then validate its arguments against the tool schema.""" - for call in tool_calls: - # Hosted/server tools (e.g. web_search) have no function name; fall back to - # call.type, mirroring Tool.key = name or type used when indexing the toolset. - name = call.function.name or call.type - tool = toolset.get(name) - if tool is None: - return RailOutcome.block(reason=f"tool call '{name}' is not an allowed tool") - block_reason = validate_arguments(tool, call.function.arguments) - if block_reason is not None: - return RailOutcome.block(reason=block_reason) - return RailOutcome.allow() + """Allowlist each call by name, then validate its arguments, reporting every call that fails.""" + violations: list[ToolViolation] = [] + for index, call in enumerate(tool_calls): + violation = self._call_violation(toolset, index, call) + if violation is not None: + violations.append(violation) + return violations_outcome(violations) + + def _call_violation(self, toolset: "Toolset", index: int, call: "ToolCall") -> "ToolViolation | None": + """The first check *call* fails, as a violation at its *index*, or None when it passes.""" + # Hosted/server tools (e.g. web_search) have no function name; fall back to + # call.type, mirroring Tool.key = name or type used when indexing the toolset. + name = call.function.name or call.type + tool = toolset.get(name) + if tool is None: + return ToolViolation( + kind="tool_call", + violation_type=ToolViolationType.TOOL_NOT_ALLOWED, + reason=f"tool call '{quoted_identity(name)}' is not an allowed tool", + tool_call_id=call.id or None, + tool_name=name, + index=index, + ) + arguments = validate_arguments(tool, call.function.arguments) + if arguments is None: + return None + return ToolViolation( + kind="tool_call", + violation_type=arguments.violation_type, + reason=arguments.reason, + tool_call_id=call.id or None, + tool_name=name, + index=index, + argument_path=arguments.argument_path, + schema_keyword=arguments.schema_keyword, + ) diff --git a/nemoguardrails/guardrails/actions/tool_result_action.py b/nemoguardrails/guardrails/actions/tool_result_action.py index cd602085da..e1986781bd 100644 --- a/nemoguardrails/guardrails/actions/tool_result_action.py +++ b/nemoguardrails/guardrails/actions/tool_result_action.py @@ -29,7 +29,9 @@ from typing import TYPE_CHECKING, List from nemoguardrails.actions.rail_outcome import RailOutcome -from nemoguardrails.guardrails.tool_rail_action import ToolRailAction +from nemoguardrails.guardrails.guardrails_types import quoted_identity +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction, violations_outcome +from nemoguardrails.rails.llm.options import ToolViolation, ToolViolationType if TYPE_CHECKING: from nemoguardrails.guardrails.tool_schema import ToolResult @@ -47,6 +49,33 @@ def _is_well_formed_content(content: object) -> bool: return isinstance(content, list) and all(isinstance(block, dict) for block in content) +def _names_a_different_tool(result: "ToolResult", prior: "ToolCall") -> bool: + """Whether the result's own name contradicts the tool its linked call named.""" + if not prior.function.name: + return False + if not result.name: + # An OpenAI tool message carries no name; its call_id already binds it to exactly one call. + return False + return result.name != prior.function.name + + +def _result_violation( + result: "ToolResult", + violation_type: ToolViolationType, + reason: str, + prior: "ToolCall | None" = None, +) -> ToolViolation: + """A violation for *result*, naming the tool from its linked call, never from the result itself.""" + return ToolViolation( + kind="tool_result", + violation_type=violation_type, + reason=reason, + tool_call_id=result.call_id or None, + tool_name=(prior.function.name or prior.type) if prior is not None else None, + index=result.message_index, + ) + + class ToolResultRailAction(ToolRailAction): """Check incoming tool results link to a prior call and are structurally well-formed.""" @@ -57,93 +86,83 @@ async def run(self, tool_results: List["ToolResult"], prior_calls: List["ToolCal return self._guarded(lambda: self._validate(tool_results, prior_calls)) def _validate(self, tool_results: List["ToolResult"], prior_calls: List["ToolCall"]) -> RailOutcome: - """Check call_id linkage, name consistency, and content shape for each result.""" - calls_by_id = self._validate_prior_calls(prior_calls) - if isinstance(calls_by_id, RailOutcome): - return calls_by_id - return self._validate_results(tool_results, calls_by_id) - - def _validate_prior_calls(self, prior_calls: List["ToolCall"]) -> "RailOutcome | dict[str, ToolCall]": - """Build a call_id index from prior_calls; return a blocking RailOutcome on duplicate IDs.""" - calls_by_id: dict[str, "ToolCall"] = {} + """Check call_id linkage, name consistency and content shape, reporting every result that fails.""" + prior_violations = self._duplicate_prior_call_violations(prior_calls) + if prior_violations: + # No result can be linked to an ambiguous call, so checking them would only add noise. + return violations_outcome(prior_violations) + calls_by_id = {call.id: call for call in prior_calls if call.id} + return violations_outcome(self._result_violations(tool_results, calls_by_id)) + + def _duplicate_prior_call_violations(self, prior_calls: List["ToolCall"]) -> list[ToolViolation]: + """One violation per call id that several prior calls share, in first-seen order.""" + seen: set[str] = set() + duplicated: list[str] = [] for call in prior_calls: if not call.id: continue - if call.id in calls_by_id: - return RailOutcome.block( - reason=f"duplicate prior tool call id '{call.id}' makes tool-result linkage ambiguous", - ) - calls_by_id[call.id] = call - return calls_by_id - - def _validate_results(self, tool_results: List["ToolResult"], calls_by_id: "dict[str, ToolCall]") -> RailOutcome: - """Check each result links to a prior call with a consistent name and well-formed content.""" - outcome = self._validate_tool_result_ids(tool_results) - if outcome is not None: - return outcome - - for result in tool_results: - outcome = self._validate_result_call_id(result, calls_by_id) - if outcome is not None: - return outcome - - prior = calls_by_id[result.call_id] # ty: ignore[invalid-argument-type] - outcome = self._validate_result_name(result, prior) - if outcome is not None: - return outcome - - outcome = self._validate_result_content(result) - if outcome is not None: - return outcome - - return RailOutcome.allow() - - def _validate_tool_result_ids(self, tool_results: List["ToolResult"]) -> "RailOutcome | None": - """Return a blocking RailOutcome if any call_id appears more than once in the result list.""" - seen: set[str] = set() + if call.id in seen and call.id not in duplicated: + duplicated.append(call.id) + seen.add(call.id) + return [ + ToolViolation( + kind="tool_result", + violation_type=ToolViolationType.DUPLICATE_PRIOR_CALL_ID, + reason=f"duplicate prior tool call id '{quoted_identity(call_id)}' makes tool-result linkage ambiguous", + tool_call_id=call_id, + ) + for call_id in duplicated + ] + + def _result_violations( + self, tool_results: List["ToolResult"], calls_by_id: "dict[str, ToolCall]" + ) -> list[ToolViolation]: + """The first check each result fails, in result order.""" + violations: list[ToolViolation] = [] + seen_ids: set[str] = set() for result in tool_results: - if not result.call_id: - continue - if result.call_id in seen: - return RailOutcome.block( - reason=f"duplicate tool result for call_id '{result.call_id}': each tool call must have exactly one result", - ) - seen.add(result.call_id) - return None - - def _validate_result_call_id( - self, result: "ToolResult", calls_by_id: "dict[str, ToolCall]" - ) -> "RailOutcome | None": - """Return a blocking RailOutcome if the result is missing a call_id or it has no prior call.""" + violation = self._first_failed_check(result, calls_by_id, seen_ids) + if violation is not None: + violations.append(violation) + if result.call_id: + seen_ids.add(result.call_id) + return violations + + def _first_failed_check( + self, result: "ToolResult", calls_by_id: "dict[str, ToolCall]", seen_ids: set[str] + ) -> "ToolViolation | None": + """The first check *result* fails: its call_id, a duplicate, linkage, name, then content.""" call_id = result.call_id if not call_id: - return RailOutcome.block(reason="tool result is missing a call_id") - if calls_by_id.get(call_id) is None: - return RailOutcome.block( - reason=f"tool result for call_id '{call_id}' does not correspond to a prior tool call", + return _result_violation(result, ToolViolationType.MISSING_CALL_ID, "tool result is missing a call_id") + prior = calls_by_id.get(call_id) + quoted_call_id = quoted_identity(call_id) + if call_id in seen_ids: + return _result_violation( + result, + ToolViolationType.DUPLICATE_RESULT, + f"duplicate tool result for call_id '{quoted_call_id}': each tool call must have exactly one result", + prior, + ) + if prior is None: + return _result_violation( + result, + ToolViolationType.UNKNOWN_CALL_ID, + f"tool result for call_id '{quoted_call_id}' does not correspond to a prior tool call", + ) + if _names_a_different_tool(result, prior): + return _result_violation( + result, + ToolViolationType.NAME_MISMATCH, + f"tool result name '{quoted_identity(result.name)}' does not match the called tool " + f"'{quoted_identity(prior.function.name)}' for call_id '{quoted_call_id}'", + prior, ) - return None - - def _validate_result_name(self, result: "ToolResult", prior: "ToolCall") -> "RailOutcome | None": - """Return a blocking RailOutcome if the result names a different tool than its linked call.""" - if not prior.function.name: - return None - if result.name == prior.function.name: - return None - if not result.name: - # An OpenAI tool message carries no name; its call_id already binds it to exactly one call. - return None - return RailOutcome.block( - reason=( - f"tool result name '{result.name}' does not match the called tool " - f"'{prior.function.name}' for call_id '{result.call_id}'" - ), - ) - - def _validate_result_content(self, result: "ToolResult") -> "RailOutcome | None": - """Return a blocking RailOutcome if the result content is not a string or list of dicts.""" if result.content is not None and not _is_well_formed_content(result.content): - return RailOutcome.block( - reason=f"tool result for call_id '{result.call_id}' has malformed content", + return _result_violation( + result, + ToolViolationType.MALFORMED_CONTENT, + f"tool result for call_id '{quoted_call_id}' has malformed content", + prior, ) return None diff --git a/nemoguardrails/guardrails/engine_registry.py b/nemoguardrails/guardrails/engine_registry.py index 22b9473b19..0106c2e15c 100644 --- a/nemoguardrails/guardrails/engine_registry.py +++ b/nemoguardrails/guardrails/engine_registry.py @@ -22,7 +22,7 @@ from nemoguardrails.guardrails.base_engine import BaseEngine from nemoguardrails.guardrails.guardrails_types import get_request_id, truncate from nemoguardrails.guardrails.model_engine import ModelEngine -from nemoguardrails.guardrails.tool_schema import ToolExchange, ToolResult, Toolset +from nemoguardrails.guardrails.tool_schema import LatestToolCall, ToolExchange, ToolResult, Toolset from nemoguardrails.rails.llm.config import Model from nemoguardrails.types import LLMModel, LLMResponse, LLMResponseChunk @@ -244,6 +244,11 @@ def extract_tool_exchanges(self, model_type: str, messages: list[dict]) -> list[ engine = self._get_engine(model_type, ModelEngine) return engine.extract_tool_exchanges(messages) + def extract_latest_tool_calls(self, model_type: str, messages: list[dict]) -> list[LatestToolCall]: + """Strictly parse the last assistant message's tool calls with the named model engine.""" + engine = self._get_engine(model_type, ModelEngine) + return engine.extract_latest_tool_calls(messages) + async def __aenter__(self): """Async context manager entry: start all engine clients.""" await self.start() diff --git a/nemoguardrails/guardrails/guardrails.py b/nemoguardrails/guardrails/guardrails.py index 8459ae977f..1c552f7c3d 100644 --- a/nemoguardrails/guardrails/guardrails.py +++ b/nemoguardrails/guardrails/guardrails.py @@ -353,22 +353,26 @@ async def check_async( self, messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, ) -> RailsResult: """Run rails on messages based on their content (asynchronous). - Supported by both LLMRails and IORails. + Supported by both LLMRails and IORails; tool rail types and ``tools`` by IORails only. """ await self._ensure_started() - return await self.rails_engine.check_async(messages, rail_types=rail_types) + return await self.rails_engine.check_async(messages, rail_types=rail_types, tools=tools) def check( self, messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, ) -> RailsResult: """Synchronous version of check_async. - Supported by both LLMRails and IORails. + Supported by both LLMRails and IORails; tool rail types and ``tools`` by IORails only. """ - return self.rails_engine.check(messages, rail_types=rail_types) + return self.rails_engine.check(messages, rail_types=rail_types, tools=tools) def register_action(self, action: Callable, name: Optional[str] = None) -> Self: """Register a custom action for the rails configuration. diff --git a/nemoguardrails/guardrails/guardrails_types.py b/nemoguardrails/guardrails/guardrails_types.py index fb226d2c3c..a7c102e5bb 100644 --- a/nemoguardrails/guardrails/guardrails_types.py +++ b/nemoguardrails/guardrails/guardrails_types.py @@ -22,6 +22,7 @@ from typing import Any, Optional, TypeAlias from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.rails.llm.options import MAX_QUOTED_IDENTITY_LENGTH, ToolViolation from nemoguardrails.types import LLMResponse, UsageInfo # LLMMessage can contain role/content, plus optional tool_calls / tool_call_id / name; content may be None @@ -101,7 +102,8 @@ class RailResult: truth: ``is_safe``, ``reason`` and ``return_value`` are derived views of it rather than a second copy that could drift. What this type adds is the aggregation ``RailOutcome`` has no concept of, because it belongs to running *many* rails: - which one blocked (``triggered_rail``) and what every rail did (``records``). + which one blocked (``triggered_rail``), which tool calls or results it blocked + (``tool_violations``) and what every rail did (``records``). ``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``. @@ -115,6 +117,7 @@ class RailResult: outcome: RailOutcome triggered_rail: str | None = None + tool_violations: tuple[ToolViolation, ...] = () records: tuple[RailCallRecord, ...] = field(default=(), compare=False) __hash__ = None @@ -170,12 +173,14 @@ def block( reason: str | None = None, metadata: Mapping[str, Any] | None = None, triggered_rail: str | None = None, + tool_violations: tuple[ToolViolation, ...] = (), records: tuple[RailCallRecord, ...] = (), ) -> "RailResult": - """A result that stops the content. Only a block names a triggering rail.""" + """A result that stops the content. Only a block names a triggering rail or tool violations.""" return cls( RailOutcome.block(reason=reason, metadata=metadata), triggered_rail=triggered_rail, + tool_violations=tool_violations, records=records, ) @@ -239,6 +244,11 @@ def truncate(text: object, max_len: int | None = None) -> str: return s[:limit] + "..." +def quoted_identity(value: object) -> str: + """A tool name or call id as a reason quotes it: cut to ``MAX_QUOTED_IDENTITY_LENGTH`` characters.""" + return truncate(value, MAX_QUOTED_IDENTITY_LENGTH) + + def serialize_prompt(messages: list[dict]) -> str: """Render a chat message list to a role-labeled string for GenerationLog's ``prompt``. diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index b413c41935..df625f71cd 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -33,6 +33,7 @@ from nemoguardrails.actions.rail_outcome import TransformTarget from nemoguardrails.base_guardrails import BaseGuardrails from nemoguardrails.exceptions import ( + InvalidCheckRequestError, NonStreamingWorkQueueFullError, RailTypeNotConfiguredError, StreamingCapacityExceededError, @@ -90,6 +91,7 @@ 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 ( + TOOL_RAIL_TYPES, ActivatedRail, ExecutedAction, GenerationLog, @@ -547,6 +549,49 @@ def _get_last_content_by_role(messages: list[dict], role: str) -> str: return "" +# TODO Rename tool_output -> tool_call and tool_input -> tool_result throughout. Use this mapping in the meantime +_RAIL_TYPE_CONFIG_SECTION = {RailType.TOOL_CALL: "tool_output", RailType.TOOL_RESULT: "tool_input"} + + +def _rail_type_configured(config: RailsConfig, rail_type: RailType) -> bool: + """Whether *config* has rails for *rail_type*; for a tool rail type, per-tool rails count too.""" + section = getattr(config.rails, _RAIL_TYPE_CONFIG_SECTION.get(rail_type, rail_type.value)) + if rail_type in TOOL_RAIL_TYPES: + # A per_tool entry with no flows runs nothing, so it does not count. + return bool(section.flows) or any(section.per_tool.values()) + return bool(section.flows) + + +def _reject_unread_tools(rail_types: Optional[list[RailType]], tools: Optional[list[dict]]) -> None: + """Raise when *tools* is passed to a check that runs no tool_call rails, which would leave it unread.""" + # Refused rather than ignored, so a caller who expected its tool calls validated learns they were not. + if tools is None: + return + if rail_types is not None and RailType.TOOL_CALL in rail_types: + return + raise InvalidCheckRequestError( + "tools is read only by a tool_call check; include tool_call in rail_types or leave tools out" + ) + + +def _has_tool_traffic(messages: LLMMessages) -> bool: + """Whether *messages* carry a tool result or an assistant turn that calls tools.""" + for message in messages: + if message.get("role") == "tool": + return True + if message.get("role") == "assistant" and message.get("tool_calls"): + return True + return False + + +def _reports_assistant_content(rails_to_run: list[str]) -> bool: + """Whether a check reports the assistant turn's text rather than the user's.""" + if "output" in rails_to_run: + return True + # A tool-call check reports the turn it checked, unless input rails run, whose rewrite must stay visible. + return "tool_call" in rails_to_run and "input" not in rails_to_run + + def _blocked_message(result: RailResult) -> str: """The text a blocked turn returns, which says whether the rail broke or fired.""" if result.failed: @@ -561,6 +606,7 @@ def _blocked_check_result(result: RailResult) -> RailsResult: content=_blocked_message(result), rail=result.triggered_rail, reason=client_reason(result), + tool_violations=list(result.tool_violations) or None, ) @@ -1331,7 +1377,13 @@ def _record_generation(timed: TimedLLMResponse) -> None: _record_generation(timed) return _GeneratedTurn(response=timed.response, messages=messages) - def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = None) -> RailsResult: + def check( + self, + messages: LLMMessages, + rail_types: Optional[list[RailType]] = None, + *, + tools: Optional[list[dict]] = None, + ) -> RailsResult: """Synchronous version of ``check_async``. Mirrors ``generate``: spins up a short-lived IORails engine with tracing @@ -1340,7 +1392,8 @@ def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = No Raises: RailTypeNotConfiguredError: If a requested rail type has no - configured flows. + configured rails, or a tool rail type is requested on a config + with no ``main`` model. """ if check_sync_call_from_async_loop(): raise RuntimeError( @@ -1356,31 +1409,51 @@ def check(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = No async def _run_sync_iorails(): """Spin up a short-lived IORails engine for one synchronous check call.""" async with IORails(sync_config, _report_usage=False) as iorails_engine: - return await iorails_engine.check_async(messages, rail_types=rail_types) + return await iorails_engine.check_async(messages, rail_types=rail_types, tools=tools) return asyncio.run(_run_sync_iorails()) - async def check_async(self, messages: LLMMessages, rail_types: Optional[list[RailType]] = None) -> RailsResult: - """Run input and/or output rails on messages without main-LLM generation. + async def check_async( + self, + messages: LLMMessages, + rail_types: Optional[list[RailType]] = None, + *, + tools: Optional[list[dict]] = None, + ) -> RailsResult: + """Run input, output and tool rails on messages without main-LLM generation. When ``rail_types`` is None the rails to run are auto-detected from the message roles (user-only -> input, assistant-only -> output, both -> - input and output). When provided, exactly the named rail types run; an - empty list (``[]``) runs no rails and returns PASSED. + input and output); tool rails never run unless named. When provided, + exactly the named rail types run, in the order generation runs them: + tool_result, input, tool_call, output. An empty list (``[]``) runs no + rails and returns PASSED. + + ``tools`` declares the tools a ``tool_call`` check validates calls + against, in the main model's wire format. It replaces the tools declared + on the main model's parameters, rather than adding to them; ``[]`` + declares none. Submitted through the same admission queue as ``generate_async`` so the check path shares non-streaming concurrency limits, request metrics, and the per-request trace span. Raises: + InvalidCheckRequestError: If ``tools`` is given but ``rail_types`` + does not include ``RailType.TOOL_CALL``. RailTypeNotConfiguredError: If a requested rail type has no - configured flows. + configured rails, or a tool rail type is requested on a config + with no ``main`` model. """ + # Before queueing: a request the check cannot serve takes no queue slot and is not logged as a failed check. + _reject_unread_tools(rail_types, tools) + if rail_types is not None: + self._validate_requested_rail_types(rail_types) await self.start() metrics_ctx = request_metrics() if self._metrics_enabled else nullcontext() with metrics_ctx: try: - return await self._generate_async_queue.submit(self._run_check, messages, rail_types) + return await self._generate_async_queue.submit(self._run_check, messages, rail_types, tools) except asyncio.QueueFull as e: if self._metrics_enabled: record_nonstream_rejected() @@ -1388,13 +1461,15 @@ async def check_async(self, messages: LLMMessages, rail_types: Optional[list[Rai # this apart from a full streaming semaphore. raise NonStreamingWorkQueueFullError(*e.args) from e - async def _run_check(self, messages: LLMMessages, rail_types: Optional[list[RailType]]) -> RailsResult: + async def _run_check( + self, messages: LLMMessages, rail_types: Optional[list[RailType]], tools: Optional[list[dict]] + ) -> RailsResult: """Queue-worker entry for ``check_async``: wrap the rails in a request span.""" tracer = self._tracer if self._tracing_enabled else None with traced_request(tracer) as (request_span, req_id): t0 = time.monotonic() try: - result = await self._do_check(messages, rail_types, req_id) + result = await self._do_check(messages, rail_types, req_id, tools=tools) except Exception: elapsed_ms = (time.monotonic() - t0) * 1000 log.error("[%s] check failed time=%.1fms", req_id, elapsed_ms, exc_info=True) @@ -1415,17 +1490,17 @@ async def _do_check( messages: LLMMessages, rail_types: Optional[list[RailType]], req_id: str, + *, + tools: Optional[list[dict]] = None, ) -> RailsResult: - """Core check pipeline: run the requested input/output rails on messages.""" + """Core check pipeline: run the requested input, output and tool rails on messages.""" log.info("[%s] check called", req_id) log.debug("[%s] check messages=%s", req_id, truncate(messages)) if rail_types is not None: - for rt in rail_types: - if not getattr(self.config.rails, rt.value).flows: - raise RailTypeNotConfiguredError(f"Requested rail type '{rt.value}' has no configured rails.") rails_to_run = [rail_type.value for rail_type in rail_types] else: + self._log_unrequested_tool_rails(messages, req_id) determined = _determine_rails_from_messages(messages) if determined is None: last = messages[-1].get("content") if messages else "" @@ -1435,13 +1510,18 @@ async def _do_check( # Which direction's text the caller gets back, and so which rewrites it can observe: with # output rails in play the answer is the response, and an input rewrite is internal to the # check. Matches how LLMRails decides what ``check`` reports. - reports_output = "output" in rails_to_run + reports_output = _reports_assistant_content(rails_to_run) if reports_output: pass_content = _get_last_content_by_role(messages, "assistant") else: pass_content = _get_last_content_by_role(messages, "user") original_content = pass_content + if "tool_result" in rails_to_run: + blocked = await self._check_tool_results(messages, req_id) + if blocked is not None: + return blocked + if "input" in rails_to_run: user_content = _get_last_content_by_role(messages, "user") # Skip when there is no user content: the content-safety action requires @@ -1463,7 +1543,12 @@ async def _do_check( else: log.info("[%s] Input rails requested but no user content to check; skipping", req_id) - if reports_output: + if "tool_call" in rails_to_run: + blocked = await self._check_latest_tool_calls(messages, tools, req_id) + if blocked is not None: + return blocked + + if "output" in rails_to_run: bot_response = _get_last_content_by_role(messages, "assistant") # Skip when there is no assistant content: the content-safety action requires # bot_response and would otherwise raise, surfacing a false BLOCK. @@ -1488,6 +1573,52 @@ async def _do_check( return RailsResult(status=RailStatus.MODIFIED, content=pass_content) return RailsResult(status=RailStatus.PASSED, content=pass_content) + def _validate_requested_rail_types(self, rail_types: list[RailType]) -> None: + """Raise when a requested rail type has no rails to run, or a tool rail type has no main model.""" + for rail_type in rail_types: + if not _rail_type_configured(self.config, rail_type): + raise RailTypeNotConfiguredError(f"Requested rail type '{rail_type.value}' has no configured rails.") + if rail_type in TOOL_RAIL_TYPES and "main" not in self.engine_registry.llms: + raise RailTypeNotConfiguredError( + f"Requested rail type '{rail_type.value}' needs a `main` model: its engine selects the " + "tool wire format and supplies config-declared tools." + ) + + def _log_unrequested_tool_rails(self, messages: LLMMessages, req_id: str) -> None: + """Say at INFO when auto-detection skips tool rails that could check this conversation.""" + tool_rails_configured = any(_rail_type_configured(self.config, rail_type) for rail_type in TOOL_RAIL_TYPES) + if tool_rails_configured and _has_tool_traffic(messages): + log.info( + "[%s] tool rails are configured but were not requested; " + "name tool_call or tool_result in rail_types to run them", + req_id, + ) + + async def _check_tool_results(self, messages: LLMMessages, req_id: str) -> Optional[RailsResult]: + """Run the tool-result rails, returning the blocked check result or None to continue.""" + log.info("[%s] Running tool result rails", req_id) + result = await self.rails_manager.are_tool_results_safe(messages) + if result.is_safe: + return None + log.info("[%s] Tool results blocked: %s", req_id, display_reason(result)) + if self._metrics_enabled: + record_request_blocked(RailDirection.INPUT) + return _blocked_check_result(result) + + async def _check_latest_tool_calls( + self, messages: LLMMessages, tools: Optional[list[dict]], req_id: str + ) -> Optional[RailsResult]: + """Run the tool-call rails on the last assistant turn, returning the blocked check result or None.""" + log.info("[%s] Running tool call rails", req_id) + llm_params = {"tools": tools} if tools is not None else None + result = await self.rails_manager.are_latest_tool_calls_safe(messages, llm_params) + if result.is_safe: + return None + log.info("[%s] Tool calls blocked: %s", req_id, display_reason(result)) + if self._metrics_enabled: + record_request_blocked(RailDirection.OUTPUT) + return _blocked_check_result(result) + def _validate_streaming_with_output_rails(self) -> None: """Raise if output rails exist but streaming is not enabled for them.""" if len(self.config.rails.output.flows) > 0 and not self._has_streaming_output_rails: diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index 29e0460753..de4d92e499 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -48,16 +48,24 @@ safe_read_body, ) from nemoguardrails.guardrails.base_engine import BaseEngine -from nemoguardrails.guardrails.guardrails_types import LLMMessages, get_request_id, truncate +from nemoguardrails.guardrails.guardrails_types import LLMMessages, get_request_id, quoted_identity, truncate from nemoguardrails.guardrails.telemetry import ( llm_call_span, set_llm_call_content, set_llm_request_attributes, set_llm_response_attributes, ) -from nemoguardrails.guardrails.tool_schema import Tool, ToolExchange, ToolResult, Toolset +from nemoguardrails.guardrails.tool_schema import ( + LatestToolCall, + Tool, + ToolCallExtractionError, + ToolExchange, + ToolResult, + Toolset, +) from nemoguardrails.llm.clients._errors import ErrorContext, raise_for_sse_error, raise_for_status from nemoguardrails.rails.llm.config import Model +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.tracing.constants import ( llm_operation_duration, record_time_per_output_chunk, @@ -430,7 +438,7 @@ def _parse_tools_nim(tools: list) -> list[Tool]: } -def _tool_result_from_message(message: dict) -> ToolResult: +def _tool_result_from_message(message: dict, message_index: int | None = None) -> ToolResult: """Normalize one OpenAI Chat Completions ``role:"tool"`` message into a ``ToolResult``. This shape has no error flag, so ``is_error`` is always ``False``. @@ -439,6 +447,7 @@ def _tool_result_from_message(message: dict) -> ToolResult: call_id=message.get("tool_call_id"), name=message.get("name"), content=message.get("content"), + message_index=message_index, ) @@ -491,7 +500,8 @@ def _extract_tool_exchanges_openai(messages: LLMMessages) -> list[ToolExchange]: # it to None so the next tool result starts a fresh exchange instead of attaching to # a closed turn. open_exchange: ToolExchange | None = None - for message in messages: + # Every message counts toward the index, so a violation points into the caller's own list. + for message_index, message in enumerate(messages): if not isinstance(message, dict): continue role = message.get("role") @@ -499,7 +509,7 @@ def _extract_tool_exchanges_openai(messages: LLMMessages) -> list[ToolExchange]: if open_exchange is None: open_exchange = ToolExchange(calls=[], results=[]) exchanges.append(open_exchange) - open_exchange.results.append(_tool_result_from_message(message)) + open_exchange.results.append(_tool_result_from_message(message, message_index)) elif role == "assistant" and message.get("tool_calls"): open_exchange = ToolExchange(calls=_tool_calls_from_message(message), results=[]) exchanges.append(open_exchange) @@ -526,6 +536,127 @@ def _extract_tool_exchanges_nim(messages: LLMMessages) -> list[ToolExchange]: } +def _latest_assistant_message(messages: LLMMessages) -> dict | None: + """The last assistant message in *messages*, or None when there is none.""" + for message in reversed(messages): + if isinstance(message, dict) and message.get("role") == "assistant": + return message + return None + + +def _string_or_none(value: object) -> str | None: + """*value* when it is a string, else None, so an error's identity fields hold only strings.""" + return value if isinstance(value, str) else None + + +def _describe_tool_call(index: int, tool_call_id: str | None) -> str: + """Name a tool call in an error message by its id, else by its position.""" + if tool_call_id: + return f"tool call '{quoted_identity(tool_call_id)}'" + return f"tool call at index {index}" + + +def _legacy_function_call_error(function_call: object) -> ToolCallExtractionError: + """The error for a legacy ``function_call``, which no tool rail validates.""" + name = function_call.get("name") if isinstance(function_call, dict) else None + return ToolCallExtractionError( + "legacy function_call is not validated by tool rails", + violation_type=ToolViolationType.LEGACY_FUNCTION_CALL, + tool_name=_string_or_none(name), + ) + + +def _malformed_tool_call_field(entry: dict, function: dict) -> str | None: + """The first identity field of a wire tool call that is not a string, or None when all of them are.""" + tool_call_id = entry.get("id") + if tool_call_id is not None and not isinstance(tool_call_id, str): + return "id" + if not isinstance(entry.get("type", "function"), str): + return "type" + name = function.get("name") + if not isinstance(name, str) or not name: + return "function name" + return None + + +def _parse_latest_tool_call(entry: object, index: int) -> ToolCall: + """Strictly parse one wire tool call, raising a ``ToolCallExtractionError`` that omits its arguments.""" + if not isinstance(entry, dict): + raise ToolCallExtractionError( + f"tool call at index {index} is not an object", + violation_type=ToolViolationType.MALFORMED_TOOL_CALL, + index=index, + ) + tool_call_id = _string_or_none(entry.get("id")) + function = entry.get("function") + # A call without a function object is not the Chat Completions shape; ChatMessage.from_dict would + # read it as the legacy flat shape, which ignores an `arguments` key and validates `{}` instead. + if not isinstance(function, dict): + raise ToolCallExtractionError( + f"{_describe_tool_call(index, tool_call_id)} does not carry a function object", + violation_type=ToolViolationType.MALFORMED_TOOL_CALL, + index=index, + tool_call_id=tool_call_id, + ) + malformed_field = _malformed_tool_call_field(entry, function) + if malformed_field is not None: + # ChatMessage.from_dict does not type these, and a non-string name breaks the per-tool rail lookup. + raise ToolCallExtractionError( + f"{_describe_tool_call(index, tool_call_id)} has a malformed {malformed_field}", + violation_type=ToolViolationType.MALFORMED_TOOL_CALL, + index=index, + tool_call_id=tool_call_id, + tool_name=_string_or_none(function.get("name")) or None, + ) + # A completed turn's "" arguments are not JSON, so from_dict rejects them, as non-streaming generation does. + try: + message = ChatMessage.from_dict({"role": "assistant", "tool_calls": [entry]}) + except ValueError: + # from_dict's message quotes the raw arguments, so it is dropped rather than chained. + raise ToolCallExtractionError( + f"{_describe_tool_call(index, tool_call_id)} has malformed arguments", + violation_type=ToolViolationType.MALFORMED_ARGUMENTS, + index=index, + tool_call_id=tool_call_id, + tool_name=_string_or_none(function.get("name")), + ) from None + return cast(list[ToolCall], message.tool_calls)[0] + + +def _parsed_or_error(entry: object, index: int) -> LatestToolCall: + """The parsed call, or the error that keeps it from being validated, so one bad call doesn't stop the rest.""" + try: + return _parse_latest_tool_call(entry, index) + except ToolCallExtractionError as error: + return error + + +def _latest_tool_calls_openai(messages: LLMMessages) -> list[LatestToolCall]: + """Strictly parse the last assistant message's tool calls; an unparsable call's error takes its place.""" + # Only that message is parsed: ChatMessage rejects roles such as the server's "context". + message = _latest_assistant_message(messages) + if message is None: + return [] + # Checked first: a function_call beside tool_calls would otherwise go unvalidated. + if message.get("function_call"): + raise _legacy_function_call_error(message["function_call"]) + raw_calls = message.get("tool_calls") + if not raw_calls: + return [] + return [_parsed_or_error(entry, index) for index, entry in enumerate(raw_calls)] + + +def _latest_tool_calls_nim(messages: LLMMessages) -> list[LatestToolCall]: + """Parse NIM's latest tool calls. NIM uses the OpenAI Chat Completions shape.""" + return _latest_tool_calls_openai(messages) + + +_LATEST_TOOL_CALL_EXTRACTORS = { + "openai": _latest_tool_calls_openai, + "nim": _latest_tool_calls_nim, +} + + class ModelEngineError(Exception): """Raised when a model engine call fails.""" @@ -1331,3 +1462,8 @@ def extract_tool_exchanges(self, messages: LLMMessages) -> list[ToolExchange]: """ extractor = _TOOL_EXCHANGE_EXTRACTORS.get(self.model_config.engine, _extract_tool_exchanges_openai) return extractor(messages) + + def extract_latest_tool_calls(self, messages: LLMMessages) -> list[LatestToolCall]: + """Strictly parse the last assistant message's tool calls in the wire shape of the model's engine.""" + extractor = _LATEST_TOOL_CALL_EXTRACTORS.get(self.model_config.engine, _latest_tool_calls_openai) + return extractor(messages) diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index 3fb5aa5c4e..79f0b39630 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -20,7 +20,7 @@ import warnings from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union +from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget from nemoguardrails.guardrails.actions.tool_call_action import ToolCallRailAction @@ -31,18 +31,30 @@ RailCallRecord, RailDirection, RailResult, + client_reason, display_reason, get_request_id, last_user_content, rewrite_user_message, ) 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 Tool, ToolExchange, ToolResult, Toolset +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction, reported_violations, violations_outcome +from nemoguardrails.guardrails.tool_schema import ( + ArgumentsViolation, + DuplicateToolError, + LatestToolCall, + Tool, + ToolCallExtractionError, + ToolExchange, + ToolResult, + Toolset, + schema_gate_violation, +) from nemoguardrails.http.runtime import create_http_client from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.manifests import RailDirection as SurfaceDirection from nemoguardrails.rails.llm.config import _get_flow_model, _get_flow_name +from nemoguardrails.rails.llm.options import ToolViolation, ToolViolationType from nemoguardrails.types import ToolCall, UsageInfo if TYPE_CHECKING: @@ -100,6 +112,8 @@ class _PerToolCheck: tool_call: ToolCall tool_result: Optional[ToolResult] = None tool_definition: Optional[Tool] = None + # Where a violation points: the call's position in tool_calls, or the result's message index. + index: Optional[int] = None def _transforms_first(rewriting: Sequence[str], flows: Sequence[str]) -> list[str]: @@ -123,6 +137,90 @@ def _tool_rail_result(outcome: RailOutcome, flow: str) -> RailResult: return RailResult(outcome) +def _blocked_before_rails(violation: ToolViolation, records: tuple[RailCallRecord, ...] = ()) -> RailResult: + """A block the manager makes itself, before any tool rail runs, stating the violation's reason.""" + return RailResult.block(reason=violation.reason, tool_violations=(violation,), records=records) + + +def _extraction_violation(error: ToolCallExtractionError) -> ToolViolation: + """The violation for a tool call the extractor could not parse, with its content-free message as the reason.""" + return ToolViolation( + kind="tool_call", + violation_type=error.violation_type, + reason=f"tool call extraction failed: {error}", + tool_call_id=error.tool_call_id, + tool_name=error.tool_name, + index=error.index, + ) + + +def _at_call_position(violation: ToolViolation, positions: list[int]) -> ToolViolation: + """*violation* re-indexed from the validated subset of calls to the call's position in ``tool_calls``.""" + if violation.index is None: + return violation + return violation.model_copy(update={"index": positions[violation.index]}) + + +def _call_order(violation: ToolViolation) -> tuple[int, int]: + """Sort key putting violations about the whole turn first, then the rest by call position.""" + if violation.index is None: + return (0, 0) + return (1, violation.index) + + +def _global_tool_rail_violations( + result: RailResult, kind: Literal["tool_call", "tool_result"], flow: str +) -> tuple[ToolViolation, ...]: + """What a blocking global tool validator reported, or one ``rail_failed`` when it broke instead of deciding.""" + if result.failed: + return ( + ToolViolation( + kind=kind, + violation_type=ToolViolationType.RAIL_FAILED, + reason=client_reason(result), + rail=_get_flow_name(flow) or flow, + ), + ) + return reported_violations(result.outcome) + + +def _tool_parsing_failure_reason(error: Exception) -> str: + """The client-facing reason the declared tools failed to parse, quoting only a message known to be safe.""" + if isinstance(error, DuplicateToolError): + return f"tool parsing failed: {error}" + return "tool parsing failed" + + +def _per_tool_violation( + result: RailResult, direction: SurfaceDirection, tool_call: ToolCall, index: Optional[int] +) -> ToolViolation: + """The violation for the call or result a per-tool rail blocked or broke on, naming that rail.""" + category = _per_tool_block_category(result) + return ToolViolation( + kind="tool_call" if direction == SurfaceDirection.TOOL_OUTPUT else "tool_result", + violation_type=category.violation_type, + reason=client_reason(result), + tool_call_id=tool_call.id or None, + tool_name=tool_call.function.name or tool_call.type, + index=index, + argument_path=category.argument_path, + schema_keyword=category.schema_keyword, + rail=result.triggered_rail, + ) + + +def _per_tool_block_category(result: RailResult) -> ArgumentsViolation: + """The violation type, and any argument path, for a per-tool rail's block or failure.""" + # A rail that raised did not decide, so it is reported as rail_failed, as the global validators are. + if result.failed: + return ArgumentsViolation(violation_type=ToolViolationType.RAIL_FAILED, reason=client_reason(result)) + # The rail's own allowlist or schema gate blocked, so the call is categorized as the global validator would. + gate_violation = schema_gate_violation(result.outcome) + if gate_violation is not None: + return gate_violation + return ArgumentsViolation(violation_type=ToolViolationType.PER_TOOL_RAIL, reason=client_reason(result)) + + def _rewritten_text(outcome: RailOutcome, direction: RailDirection, flow: str) -> str: """Return what a rail rewrote this direction's variable to, ignoring any it also named. @@ -451,20 +549,22 @@ async def are_tool_calls_safe( """ if not tool_calls: return RailResult.allow() - - active = self._enabled_flows(list(self._tool_call_actions), enabled) - enabled_per_tool_flows = { - flow for flows in self.per_tool_call_flows.values() for flow in self._enabled_flows(flows, enabled) - } - if not active and not enabled_per_tool_flows: + if not self._any_tool_call_rail_enabled(enabled): return RailResult.allow() + active = self._enabled_flows(list(self._tool_call_actions), enabled) global_result = RailResult.allow() try: toolset = self.engine_registry.parse_tools(model_type, llm_params) except Exception as e: log.warning("[%s] tool parsing failed; blocking tool calls: %s", get_request_id(), e) - return RailResult.block(reason=f"tool parsing failed: {e}") + return _blocked_before_rails( + ToolViolation( + kind="tool_call", + violation_type=ToolViolationType.INVALID_TOOLSET, + reason=_tool_parsing_failure_reason(e), + ) + ) if active: rails = {flow: self._run_tool_call_rail(flow, tool_calls, toolset) for flow in active} @@ -478,7 +578,9 @@ async def are_tool_calls_safe( tool_definition = toolset.get(tool_name) for flow in self._enabled_flows(self.per_tool_call_flows.get(tool_name, []), enabled): checks.append( - _PerToolCheck(f"{index}:{flow}", flow, tool_call=tool_call, tool_definition=tool_definition) + _PerToolCheck( + f"{index}:{flow}", flow, tool_call=tool_call, tool_definition=tool_definition, index=index + ) ) if not checks: return global_result @@ -491,6 +593,7 @@ async def are_tool_calls_safe( tool_call=check.tool_call, tool_result=check.tool_result, tool_definition=check.tool_definition, + index=check.index, ) for check in checks } @@ -504,6 +607,76 @@ async def are_tool_calls_safe( return replace(per_tool_result, records=combined_records) return RailResult.allow(records=combined_records) + async def are_latest_tool_calls_safe( + self, + messages: list[dict], + llm_params: Optional[dict], + *, + enabled: Union[bool, list[str]] = True, + model_type: str = "main", + ) -> RailResult: + """Validate the tool calls on the last assistant message of *messages*, as a check with no model response.""" + if not self._any_tool_call_rail_enabled(enabled): + return RailResult.allow() + try: + extracted = self.engine_registry.extract_latest_tool_calls(model_type, messages) + except ToolCallExtractionError as e: + log.warning("[%s] tool call extraction failed; blocking: %s", get_request_id(), e) + return _blocked_before_rails(_extraction_violation(e)) + except Exception as e: + # Only a ToolCallExtractionError is written to omit argument text, so this one stays in the log. + log.warning("[%s] tool call extraction failed unexpectedly; blocking: %s", get_request_id(), e) + return _blocked_before_rails( + ToolViolation( + kind="tool_call", + violation_type=ToolViolationType.EXTRACTION_FAILED, + reason="tool call extraction failed", + ) + ) + if any(isinstance(entry, ToolCallExtractionError) for entry in extracted): + return await self._blocked_with_every_call_violation( + extracted, llm_params, enabled=enabled, model_type=model_type + ) + tool_calls = [entry for entry in extracted if isinstance(entry, ToolCall)] + if not tool_calls: + log.debug("[%s] the latest assistant message carries no tool calls", get_request_id()) + return await self.are_tool_calls_safe(tool_calls, llm_params, enabled=enabled, model_type=model_type) + + async def _blocked_with_every_call_violation( + self, + extracted: list[LatestToolCall], + llm_params: Optional[dict], + *, + enabled: Union[bool, list[str]], + model_type: str, + ) -> RailResult: + """Block a turn with unparsable calls, still validating the calls that parsed so they are reported too.""" + failures = [entry for entry in extracted if isinstance(entry, ToolCallExtractionError)] + # One line per check, however many calls fail, so a large request cannot flood the log. + log.warning( + "[%s] tool call extraction failed for %d of %d tool calls; blocking. First: %s", + get_request_id(), + len(failures), + len(extracted), + failures[0], + ) + violations = [_extraction_violation(failure) for failure in failures] + parsed = [(position, entry) for position, entry in enumerate(extracted) if isinstance(entry, ToolCall)] + rails_result = await self.are_tool_calls_safe( + [call for _, call in parsed], llm_params, enabled=enabled, model_type=model_type + ) + positions = [position for position, _ in parsed] + violations.extend(_at_call_position(violation, positions) for violation in rails_result.tool_violations) + violations.sort(key=_call_order) + # A rail that raised keeps this a failure, so the check answers with its internal-error message, not a refusal, + # and with the failed rail's reason rather than an earlier call's. + if rails_result.failed: + outcome = RailOutcome.failure(reason=client_reason(rails_result)) + else: + outcome = RailOutcome.block(reason=violations[0].reason) + # No triggered_rail: the block is the manager's, as for any call it cannot parse. + return RailResult(outcome, tool_violations=tuple(violations), records=rails_result.records) + async def are_tool_results_safe( self, messages: list[dict], @@ -528,8 +701,15 @@ async def are_tool_results_safe( try: exchanges = self.engine_registry.extract_tool_exchanges(model_type, messages) except Exception as e: + # The exception's text is not known to omit message content, so it stays in the log. log.warning("[%s] tool exchange extraction failed; blocking: %s", get_request_id(), e) - return RailResult.block(reason=f"tool exchange extraction failed: {e}") + return _blocked_before_rails( + ToolViolation( + kind="tool_result", + violation_type=ToolViolationType.EXTRACTION_FAILED, + reason="tool exchange extraction failed", + ) + ) if not any(exchange.results for exchange in exchanges): return RailResult.allow() @@ -552,13 +732,28 @@ async def are_tool_results_safe( if matched_call is None: # Fail closed: a per-tool policy is enabled but this result's identity is # unverifiable, so which policy applies can't be known. - return RailResult.block( - reason="tool result cannot be linked to exactly one prior call; per-tool policy cannot be verified", + return _blocked_before_rails( + ToolViolation( + kind="tool_result", + violation_type=ToolViolationType.UNLINKABLE_RESULT, + reason="tool result cannot be linked to exactly one prior call; " + "per-tool policy cannot be verified", + tool_call_id=tool_result.call_id or None, + index=tool_result.message_index, + ), records=global_result.records, ) tool_name = matched_call.function.name or matched_call.type for flow in self._enabled_flows(self.per_tool_result_flows.get(tool_name, []), enabled): - checks.append(_PerToolCheck(f"{index}:{flow}", flow, tool_call=matched_call, tool_result=tool_result)) + checks.append( + _PerToolCheck( + f"{index}:{flow}", + flow, + tool_call=matched_call, + tool_result=tool_result, + index=tool_result.message_index, + ) + ) if not checks: return global_result @@ -570,6 +765,7 @@ async def are_tool_results_safe( tool_call=check.tool_call, tool_result=check.tool_result, tool_definition=check.tool_definition, + index=check.index, ) for check in checks } @@ -590,6 +786,12 @@ def _flows_to_run( active = self._enabled_flows(configured, enabled) return _transforms_first(_rewriting_flows(self._rails, direction, active), active) + def _any_tool_call_rail_enabled(self, enabled: Union[bool, list[str]]) -> bool: + """Whether the toggle leaves any global or per-tool tool-call rail to run.""" + if self._enabled_flows(list(self._tool_call_actions), enabled): + return True + return any(self._enabled_flows(flows, enabled) for flows in self.per_tool_call_flows.values()) + @staticmethod def _enabled_flows(configured: list[str], enabled: Union[bool, list[str]]) -> list[str]: """Resolve the per-request toggle into the configured flows to run: True all, False none.""" @@ -647,6 +849,12 @@ 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 = _tool_rail_result(await self._tool_call_actions[flow].run(toolset, tool_calls), flow) + if not result.is_safe: + result = replace( + result, + triggered_rail=_get_flow_name(flow) or flow, + tool_violations=_global_tool_rail_violations(result, "tool_call", flow), + ) result = replace(result, records=(_rail_call_record(flow, "tool_output", result),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: @@ -660,17 +868,18 @@ async def _run_tool_call_rail(self, flow: str, tool_calls: list[ToolCall], tools async def _run_tool_result_rail(self, flow: str, exchanges: list[ToolExchange]) -> RailResult: """Validate each turn's results against that turn's calls, wrapped in an INPUT rail span. - Each exchange is validated independently so ``call_id`` linkage stays turn-local; - the first unsafe exchange short-circuits. + Each exchange is validated independently so ``call_id`` linkage stays turn-local, and + every exchange is checked so the result reports the violations of all of them. """ action = self._tool_result_actions[flow] with rail_span(self._tracer, flow, RailDirection.INPUT) as span: - outcome = RailOutcome.allow() - for exchange in exchanges: - outcome = await action.run(exchange.results, exchange.calls) - if outcome.is_blocked: - break - result = _tool_rail_result(outcome, flow) + result = _tool_rail_result(await self._outcome_across_exchanges(action, exchanges), flow) + if not result.is_safe: + result = replace( + result, + triggered_rail=_get_flow_name(flow) or flow, + tool_violations=_global_tool_rail_violations(result, "tool_result", flow), + ) result = replace(result, records=(_rail_call_record(flow, "tool_input", result),)) mark_rail_stop(span, result.is_safe) if self._content_capture_enabled: @@ -686,6 +895,19 @@ async def _run_tool_result_rail(self, flow: str, exchanges: list[ToolExchange]) ) return result + @staticmethod + async def _outcome_across_exchanges(action: ToolResultRailAction, exchanges: list[ToolExchange]) -> RailOutcome: + """Run the result validator on every turn and merge the violations each one reports.""" + violations: list[ToolViolation] = [] + for exchange in exchanges: + outcome = await action.run(exchange.results, exchange.calls) + reported = reported_violations(outcome) + if outcome.is_blocked and not reported: + # A block with nothing to merge, such as a validator that broke, stands as it is. + return outcome + violations.extend(reported) + return violations_outcome(violations) + async def _run_per_tool_rail( self, direction: SurfaceDirection, @@ -693,6 +915,8 @@ async def _run_per_tool_rail( tool_call: ToolCall, tool_result: Optional[ToolResult] = None, tool_definition: Optional[Tool] = None, + # Where the checked item sits: the call's position in tool_calls, or the result's message position. + index: Optional[int] = None, ) -> RailResult: """Dispatch one per-tool rail, passing the resolved ToolCall/ToolResult/Tool through.""" tool_name = tool_call.function.name or tool_call.type @@ -705,6 +929,7 @@ async def _run_per_tool_rail( result = _rail_result(rail_execution.outcome) if not result.is_safe: result = replace(result, triggered_rail=_get_flow_name(flow) or flow) + result = replace(result, tool_violations=(_per_tool_violation(result, direction, tool_call, index),)) record = _rail_call_record(flow, rail_type, result, rail_execution.llm_calls, tool_name=tool_name) result = replace(result, records=(record,)) mark_rail_stop(span, result.is_safe) @@ -802,6 +1027,7 @@ async def _run_per_tool_rails_sequential( tool_call=check.tool_call, tool_result=check.tool_result, tool_definition=check.tool_definition, + index=check.index, ) collected.extend(result.records) log.debug("[%s] %s flow %s result %s", req_id, direction.value, check.flow, result) diff --git a/nemoguardrails/guardrails/tool_rail_action.py b/nemoguardrails/guardrails/tool_rail_action.py index 70027df56f..e38a0b739a 100644 --- a/nemoguardrails/guardrails/tool_rail_action.py +++ b/nemoguardrails/guardrails/tool_rail_action.py @@ -22,17 +22,37 @@ from __future__ import annotations import logging +from collections.abc import Sequence from typing import TYPE_CHECKING, Callable, Optional from nemoguardrails.actions.rail_outcome import RailOutcome from nemoguardrails.guardrails.rail_guard import rail_error_outcome from nemoguardrails.guardrails.telemetry import action_span +from nemoguardrails.rails.llm.options import ToolViolation if TYPE_CHECKING: from opentelemetry.trace import Tracer log = logging.getLogger(__name__) +# Metadata is the only channel out of an action, so a tool rail's violations travel in it. +TOOL_VIOLATIONS_KEY = "tool_violations" + + +def violations_outcome(violations: Sequence[ToolViolation]) -> RailOutcome: + """Allow when nothing was violated, else block once, with the first violation's reason.""" + if not violations: + return RailOutcome.allow() + return RailOutcome.block( + reason=violations[0].reason, + metadata={TOOL_VIOLATIONS_KEY: [violation.model_dump(mode="json") for violation in violations]}, + ) + + +def reported_violations(outcome: RailOutcome) -> tuple[ToolViolation, ...]: + """The violations a tool rail put in *outcome*, read back from its metadata.""" + return tuple(ToolViolation.model_validate(violation) for violation in outcome.metadata.get(TOOL_VIOLATIONS_KEY, ())) + class ToolRailAction: """Base for the local, model-free tool-calling rails (tool-call and tool-result).""" diff --git a/nemoguardrails/guardrails/tool_schema.py b/nemoguardrails/guardrails/tool_schema.py index 363dc25195..c3cf650334 100644 --- a/nemoguardrails/guardrails/tool_schema.py +++ b/nemoguardrails/guardrails/tool_schema.py @@ -31,16 +31,24 @@ """ import functools +import json +import logging import re from collections.abc import Iterable from dataclasses import dataclass from typing import Any, Callable, NamedTuple import jsonschema +from jsonschema.exceptions import best_match +from jsonschema.validators import validator_for from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.guardrails.guardrails_types import quoted_identity +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.types import ToolCall +log = logging.getLogger(__name__) + @dataclass(frozen=True, slots=True) class Tool: @@ -66,6 +74,10 @@ def key(self) -> str: return self.name or self.type +class DuplicateToolError(ValueError): + """A toolset declares the same tool key twice; the message names only that key.""" + + class Toolset: """The set of tools declared on a request, indexed by tool key. @@ -89,7 +101,7 @@ def __init__(self, tools: Iterable[Tool] | None = None) -> None: if not tool.key: continue if tool.key in self._by_key: - raise ValueError(f"duplicate tool '{tool.key}' in toolset") + raise DuplicateToolError(f"duplicate tool '{tool.key}' in toolset") self._by_key[tool.key] = tool def get(self, key: str) -> Tool | None: @@ -118,6 +130,8 @@ class ToolResult: name: str | None = None content: str | list[dict] | None = None is_error: bool = False + # Position of the result's message in the conversation, so a violation can point at it. + message_index: int | None = None def to_dict(self) -> dict: return { @@ -135,6 +149,29 @@ class ToolExchange(NamedTuple): results: list[ToolResult] +class ToolCallExtractionError(ValueError): + """A tool call that cannot be validated, described without its argument text.""" + + def __init__( + self, + message: str, + *, + violation_type: ToolViolationType, + index: int | None = None, + tool_call_id: str | None = None, + tool_name: str | None = None, + ) -> None: + super().__init__(message) + self.violation_type = violation_type + self.index = index + self.tool_call_id = tool_call_id + self.tool_name = tool_name + + +# One entry per tool call on a checked turn: the parsed call, or the error that keeps it from being validated. +LatestToolCall = ToolCall | ToolCallExtractionError + + def _schema_accepts_no_arguments(schema: dict) -> bool: """Whether an argument schema declares no way to supply arguments. @@ -170,8 +207,46 @@ def _schema_rejects_argument(schema: dict, argument_name: str) -> bool: return not any(keyword in schema for keyword in ("anyOf", "oneOf", "allOf", "$ref")) -def _no_arguments_reason(tool: Tool, arguments: dict) -> str | None: - """Block reason for a tool that accepts no arguments, or ``None`` when the call is allowed. +@dataclass(frozen=True, slots=True) +class ArgumentsViolation: + """Why a call's arguments fail its tool's declaration, stated without any argument value.""" + + violation_type: ToolViolationType + reason: str + argument_path: str | None = None + schema_keyword: str | None = None + + +# Where tool_output_validation records the allowlist or schema violation it blocked on, so a per-tool rail's +# block is reported with the same type and argument path the global validator gives the same call. +SCHEMA_GATE_VIOLATION_KEY = "tool_schema_violation" + + +def _schema_gate_block(violation: ArgumentsViolation) -> RailOutcome: + """A block stating *violation*'s reason and carrying its type and argument path in the metadata.""" + recorded = { + "violation_type": violation.violation_type.value, + "argument_path": violation.argument_path, + "schema_keyword": violation.schema_keyword, + } + return RailOutcome.block(reason=violation.reason, metadata={SCHEMA_GATE_VIOLATION_KEY: recorded}) + + +def schema_gate_violation(outcome: RailOutcome) -> ArgumentsViolation | None: + """The violation tool_output_validation blocked *outcome* on, or None when the rail itself decided.""" + recorded = outcome.metadata.get(SCHEMA_GATE_VIOLATION_KEY) + if not isinstance(recorded, dict): + return None + return ArgumentsViolation( + violation_type=ToolViolationType(recorded["violation_type"]), + reason=outcome.reason or "", + argument_path=recorded.get("argument_path"), + schema_keyword=recorded.get("schema_keyword"), + ) + + +def _no_arguments_violation(tool: Tool, arguments: dict) -> ArgumentsViolation | None: + """The violation for a tool that accepts no arguments, or ``None`` when the call is allowed. Hosted/server tools (no ``name``) are allowlist-only -- the provider owns the call shape -- so arguments are accepted here. A function tool that declares no parameters @@ -180,32 +255,152 @@ def _no_arguments_reason(tool: Tool, arguments: dict) -> str | None: if tool.name is None: return None if arguments: - return f"tool '{tool.key}' accepts no arguments but the call supplied: {sorted(arguments)}" + # A count, not the names: the model chose them, so they can carry data. + noun = "argument" if len(arguments) == 1 else "arguments" + return ArgumentsViolation( + violation_type=ToolViolationType.UNEXPECTED_ARGUMENTS, + reason=f"tool '{tool.key}' accepts no arguments but the call supplied {len(arguments)} {noun}", + ) + return None + + +# Stands in for a path segment the model chose, such as a key in a map-shaped argument. +_UNDECLARED_KEY = "*" + + +def _declared_segments(schema: Any, segments: Iterable[object]) -> list[str]: + """The path's segments, keeping array indices and property names the schema declares, masking the rest.""" + # Only the direct `properties` count as declared, so a name reached through allOf or $ref is + # masked too: that shows less than it could, never text the model wrote. + declared: list[str] = [] + current = schema + for segment in segments: + properties = _declared_properties(current) + if isinstance(segment, int): + declared.append(str(segment)) + current = _array_items(current) + elif isinstance(segment, str) and segment in properties: + declared.append(segment) + current = properties[segment] + else: + declared.append(_UNDECLARED_KEY) + current = {} + return declared + + +def _declared_properties(schema: Any) -> dict[str, Any]: + """The ``properties`` mapping *schema* declares, or an empty one.""" + properties = schema.get("properties") if isinstance(schema, dict) else None + return properties if isinstance(properties, dict) else {} + + +def _array_items(schema: Any) -> Any: + """The ``items`` subschema *schema* declares for array elements, or an empty one.""" + return schema.get("items", {}) if isinstance(schema, dict) else {} + + +def _json_pointer(segments: Iterable[str]) -> str: + """An RFC 6901 pointer into the arguments; the empty string is the arguments object itself.""" + return "".join("/" + segment.replace("~", "~0").replace("/", "~1") for segment in segments) + + +def _first_missing_name(names: list[str], instance: dict) -> str | None: + """The first of *names* the object lacks; the names come from the schema.""" + return next((name for name in names if name not in instance), None) + + +def _first_missing_dependency(dependencies: dict[str, list[str]], instance: dict) -> str | None: + """The first ``dependentRequired`` name the object lacks; the names come from the schema.""" + triggered = (names for trigger, names in dependencies.items() if trigger in instance) + return next((name for names in triggered for name in names if name not in instance), None) + + +def _missing_property(error: jsonschema.ValidationError) -> str | None: + """The schema-named property a ``required`` or ``dependentRequired`` failure is about, else None.""" + # These keywords are reported at the object missing the property. Draft 3 writes `required: true` + # on the property itself, which jsonschema already reports at that property, so the type checks + # below leave its path alone. + value, instance = error.validator_value, error.instance + if not isinstance(instance, dict): + return None + if error.validator == "required" and isinstance(value, list): + return _first_missing_name(value, instance) + if error.validator == "dependentRequired" and isinstance(value, dict): + return _first_missing_dependency(value, instance) return None -def validate_arguments(tool: Tool, arguments: dict) -> str | None: +def _argument_path(error: jsonschema.ValidationError, schema: dict) -> str: + """Where in the arguments *error* failed, in schema-declared names only.""" + segments = _declared_segments(schema, error.absolute_path) + missing = _missing_property(error) + if missing is not None: + segments.append(missing) + return _json_pointer(segments) + + +def _schema_mismatch_violation(tool: Tool, error: jsonschema.ValidationError) -> ArgumentsViolation: + """The violation for arguments that fail the schema, stated by keyword and path, never by value.""" + # jsonschema's own message quotes the failing value, so it is not used. + keyword = str(error.validator) + path = _argument_path(error, tool.arguments_schema or {}) + location = f"'{path}'" if path else "the top level" + return ArgumentsViolation( + violation_type=ToolViolationType.ARGUMENTS_INVALID, + reason=f"arguments for tool '{tool.key}' do not match its schema: '{keyword}' failed at {location}", + argument_path=path, + schema_keyword=keyword, + ) + + +def validate_arguments(tool: Tool, arguments: dict) -> ArgumentsViolation | None: """Validate model-supplied tool-call arguments against the tool's schema. - Returns ``None`` when the arguments are valid. Returns a human-readable reason when + Returns ``None`` when the arguments are valid. Returns an ``ArgumentsViolation`` when the arguments violate the schema, when the declared schema itself is not valid JSON Schema (e.g. a non-JSON-Schema dialect reaching this validator before its engine adapter normalizes it), or when a function tool that declares no parameters is called with arguments. """ if tool.arguments_schema is None: - return _no_arguments_reason(tool, arguments) + return _no_arguments_violation(tool, arguments) try: - jsonschema.validate(instance=arguments, schema=tool.arguments_schema) - except jsonschema.ValidationError as exc: - return f"arguments for tool '{tool.key}' do not match its schema: {exc.message}" + schema_json = json.dumps(tool.arguments_schema, sort_keys=True) + except (TypeError, ValueError) as exc: + return _invalid_schema_violation(tool, str(exc)) + try: + validator = _checked_validator(schema_json) except jsonschema.SchemaError as exc: - return f"declared schema for tool '{tool.key}' is not valid JSON Schema: {exc.message}" + return _invalid_schema_violation(tool, exc.message) + # best_match picks the error jsonschema.validate would have raised. + error = best_match(validator.iter_errors(arguments)) + if error is not None: + return _schema_mismatch_violation(tool, error) if _schema_accepts_no_arguments(tool.arguments_schema): - return _no_arguments_reason(tool, arguments) + return _no_arguments_violation(tool, arguments) return None +@functools.lru_cache(maxsize=256) +def _checked_validator(schema_json: str) -> Any: + """A validator for one argument schema, checked against its metaschema once rather than on every call.""" + # Built from the JSON text, not the caller's dict, so a later change to that dict cannot reach the cache. + schema = json.loads(schema_json) + validator_class = validator_for(schema) + validator_class.check_schema(schema) + return validator_class(schema) + + +def _invalid_schema_violation(tool: Tool, detail: str) -> ArgumentsViolation: + """The violation for a declared schema that is not JSON Schema, with the detail kept to the log.""" + # The detail can quote the declared schema, which may be operator config, so it stays in the log. + log.warning("declared schema for tool '%s' is not valid JSON Schema: %s", tool.key, detail) + return ArgumentsViolation( + violation_type=ToolViolationType.INVALID_TOOL_SCHEMA, + reason=f"declared schema for tool '{tool.key}' is not valid JSON Schema", + ) + + def tool_output_validation(func: Callable[..., Any]) -> Callable[..., Any]: """Validate a TOOL_OUTPUT action's arguments against the tool's schema before it runs. @@ -222,10 +417,15 @@ async def wrapper(*args: Any, **kwargs: Any) -> RailOutcome: tool_definition: Tool | None = kwargs["tool_definition"] name = tool_call.function.name or tool_call.type if tool_definition is None: - return RailOutcome.block(reason=f"tool call '{name}' is not an allowed tool") - reason = validate_arguments(tool_definition, tool_call.function.arguments) - if reason is not None: - return RailOutcome.block(reason=reason) + return _schema_gate_block( + ArgumentsViolation( + violation_type=ToolViolationType.TOOL_NOT_ALLOWED, + reason=f"tool call '{quoted_identity(name)}' is not an allowed tool", + ) + ) + violation = validate_arguments(tool_definition, tool_call.function.arguments) + if violation is not None: + return _schema_gate_block(violation) argument_name = kwargs.get("argument_name") if argument_name is not None: diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 581773aa32..14d6fcd82b 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -72,6 +72,7 @@ InvalidRailsConfigurationError, InvalidStateError, RailTypeNotConfiguredError, + RailTypeNotSupportedError, StreamingNotSupportedError, ) from nemoguardrails.kb.kb import KnowledgeBase @@ -93,6 +94,7 @@ RailsConfig, ) from nemoguardrails.rails.llm.options import ( + TOOL_RAIL_TYPES, GenerationLog, GenerationOptions, GenerationResponse, @@ -1634,6 +1636,8 @@ async def check_async( self, messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, ) -> RailsResult: """Run rails on messages based on their content (asynchronous). @@ -1654,6 +1658,8 @@ async def check_async( rail_types: Optional list of rail types to run, e.g. ``[RailType.INPUT]`` or ``[RailType.OUTPUT]``. When provided, overrides automatic detection. + tools: Not supported. Only a tool-call check reads tools, and + only the IORails engine runs one. Returns: RailsResult containing: @@ -1680,10 +1686,15 @@ async def check_async( result = await rails.check_async(messages, rail_types=[RailType.INPUT]) Raises: + RailTypeNotSupportedError: If ``tools`` is given, or a requested rail + type is ``tool_call`` or ``tool_result``, which only the IORails + engine runs. RailTypeNotConfiguredError: If a requested rail type has no configured flows. """ + _reject_tools(tools) if rail_types is not None: + _reject_tool_rail_types(rail_types) for rt in rail_types: if not getattr(self.config.rails, rt.value).flows: raise RailTypeNotConfiguredError(f"Requested rail type '{rt.value}' has no configured rails.") @@ -1723,6 +1734,8 @@ def check( self, messages: List[dict], rail_types: Optional[List[RailType]] = None, + *, + tools: Optional[List[dict]] = None, ) -> RailsResult: """Run rails on messages based on their content (synchronous). @@ -1731,11 +1744,15 @@ def check( Args: messages: List of message dicts with 'role' and 'content' fields. rail_types: Optional list of rail types to run. See check_async() for details. + tools: Not supported. See check_async() for details. Returns: RailsResult containing status, content, and optional blocking rail name. Raises: + RailTypeNotSupportedError: If ``tools`` is given, or a requested rail + type is ``tool_call`` or ``tool_result``, which only the IORails + engine runs. RailTypeNotConfiguredError: If a requested rail type has no configured flows. """ @@ -1745,7 +1762,7 @@ def check( ) loop = get_or_create_event_loop() - return loop.run_until_complete(self.check_async(messages, rail_types=rail_types)) + return loop.run_until_complete(self.check_async(messages, rail_types=rail_types, tools=tools)) def register_action(self, action: Callable, name: Optional[str] = None) -> Self: """Register a custom action for the rails configuration.""" @@ -2107,6 +2124,27 @@ def _prepare_params( yield chunk +def _reject_tools(tools: Optional[List[dict]]) -> None: + """Raise when a check passes tools, which only an IORails tool-call check reads.""" + # Refused rather than dropped, so a caller relying on tool validation learns it did not run. + if tools is not None: + raise RailTypeNotSupportedError( + "LLMRails check() does not run tool rails, so it does not take tools; " + "tool_call checks run on the IORails engine only." + ) + + +def _reject_tool_rail_types(rail_types: List[RailType]) -> None: + """Raise when any requested rail type is a tool rail type, which LLMRails check() cannot run.""" + # Checked before the configured-flows loop: without it a tool rail type raises AttributeError + # there, or would pass silently once the config grows a matching section. + requested_rail_type = sorted({rail_type.value for rail_type in rail_types if rail_type in TOOL_RAIL_TYPES}) + if requested_rail_type: + raise RailTypeNotSupportedError( + f"LLMRails supports input and output rails only, not {', '.join(requested_rail_type)}" + ) + + def _determine_rails_from_messages(messages: List[dict]) -> Optional[dict]: roles = {msg.get("role") for msg in reversed(messages)} has_user = "user" in roles diff --git a/nemoguardrails/rails/llm/options.py b/nemoguardrails/rails/llm/options.py index 2d9a80aeb6..83e3532caa 100644 --- a/nemoguardrails/rails/llm/options.py +++ b/nemoguardrails/rails/llm/options.py @@ -78,9 +78,9 @@ """ from enum import Enum -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from nemoguardrails.logging.explain import LLMCallInfo @@ -88,6 +88,12 @@ class RailType(str, Enum): INPUT = "input" OUTPUT = "output" + # Tool rail types run on the IORails engine only; LLMRails check() rejects them. + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + + +TOOL_RAIL_TYPES = frozenset({RailType.TOOL_CALL, RailType.TOOL_RESULT}) class RailStatus(str, Enum): @@ -104,6 +110,85 @@ class RailsResult(BaseModel): default=None, description="Why the rail blocked the content, when the engine reports it (IORails).", ) + tool_violations: Optional[List["ToolViolation"]] = Field( + default=None, + description="The tool calls or tool results a tool rail blocked (IORails). None unless a tool rail blocked.", + ) + + +class ToolViolationType(str, Enum): + """What a blocked tool call or tool result violated; callers switch on this, never on the reason text.""" + + # Tool calls + TOOL_NOT_ALLOWED = "tool_not_allowed" + ARGUMENTS_INVALID = "arguments_invalid" + UNEXPECTED_ARGUMENTS = "unexpected_arguments" + INVALID_TOOL_SCHEMA = "invalid_tool_schema" + MALFORMED_ARGUMENTS = "malformed_arguments" + MALFORMED_TOOL_CALL = "malformed_tool_call" + LEGACY_FUNCTION_CALL = "legacy_function_call" + INVALID_TOOLSET = "invalid_toolset" + # Tool results + MISSING_CALL_ID = "missing_call_id" + UNKNOWN_CALL_ID = "unknown_call_id" + DUPLICATE_RESULT = "duplicate_result" + DUPLICATE_PRIOR_CALL_ID = "duplicate_prior_call_id" + NAME_MISMATCH = "name_mismatch" + MALFORMED_CONTENT = "malformed_content" + UNLINKABLE_RESULT = "unlinkable_result" + # Either direction + EXTRACTION_FAILED = "extraction_failed" + RAIL_FAILED = "rail_failed" + PER_TOOL_RAIL = "per_tool_rail" + + +# A tool name or call id is text the model or the caller chose, so a reason or violation keeps at most this much of it. +MAX_QUOTED_IDENTITY_LENGTH = 64 + + +class ToolViolation(BaseModel): + """One tool call or tool result a tool rail blocked, identified so a caller can act on it.""" + + model_config = ConfigDict(frozen=True) + + kind: Literal["tool_call", "tool_result"] = Field(description="Whether a tool call or a tool result was blocked.") + violation_type: ToolViolationType = Field( + description="What the call or result violated. Switch on this, never on `reason`." + ) + reason: str = Field(description="Why it was blocked. Never contains argument values or tool-result content.") + tool_call_id: Optional[str] = Field(default=None, description="The call's id, or the id the result links to.") + tool_name: Optional[str] = Field( + default=None, description="The tool called. For a result, taken from the call it links to." + ) + index: Optional[int] = Field( + default=None, + description="For a call, its position in the last assistant message's `tool_calls`. " + "For a result, its message's position in `messages`.", + ) + argument_path: Optional[str] = Field( + default=None, description="JSON Pointer to the failing argument (`arguments_invalid` only)." + ) + schema_keyword: Optional[str] = Field( + default=None, description="The JSON Schema keyword that failed (`arguments_invalid` only)." + ) + rail: Optional[str] = Field( + default=None, description="The rail that blocked or failed (`per_tool_rail` and `rail_failed` only)." + ) + + @field_validator("tool_call_id", "tool_name", mode="before") + @classmethod + def _string_identity_or_none(cls, value: Any) -> Optional[str]: + """Store a non-string id or name as None, so a malformed message cannot break building its violation.""" + return value if isinstance(value, str) else None + + @field_validator("tool_name") + @classmethod + def _capped_tool_name(cls, value: Optional[str]) -> Optional[str]: + """Cut a tool name to ``MAX_QUOTED_IDENTITY_LENGTH`` characters, since an undeclared one is model-made.""" + # The call id stays whole: it is the identity a harness matches the violation on. + if value is None or len(value) <= MAX_QUOTED_IDENTITY_LENGTH: + return value + return value[:MAX_QUOTED_IDENTITY_LENGTH] + "..." class GenerationLogOptions(BaseModel): diff --git a/nemoguardrails/server/api.py b/nemoguardrails/server/api.py index 5e5680d088..4f95a9d234 100644 --- a/nemoguardrails/server/api.py +++ b/nemoguardrails/server/api.py @@ -36,11 +36,13 @@ from nemoguardrails import Guardrails, LLMRails, RailsConfig, utils from nemoguardrails.exceptions import ( + InvalidCheckRequestError, InvalidModelConfigurationError, InvalidStateError, LLMCallException, NonStreamingWorkQueueFullError, RailTypeNotConfiguredError, + RailTypeNotSupportedError, StreamingCapacityExceededError, StreamingNotSupportedError, ) @@ -51,17 +53,19 @@ from nemoguardrails.llm.clients._errors import build_error_payload, normalize_error_status from nemoguardrails.llm.models.initializer import ModelInitializationError from nemoguardrails.rails.llm.config import Model -from nemoguardrails.rails.llm.options import GenerationResponse, RailStatus +from nemoguardrails.rails.llm.options import GenerationResponse, RailStatus, ToolViolation from nemoguardrails.server.datastore.datastore import DataStore from nemoguardrails.server.exception_handlers import ( bad_request_error_handler, http_exception_handler, internal_error_handler, + invalid_check_request_error_handler, invalid_state_error_handler, llm_call_exception_handler, model_initialization_error_handler, queue_full_error_handler, rail_type_not_configured_error_handler, + rail_type_not_supported_error_handler, streaming_capacity_error_handler, validation_error_handler, ) @@ -224,6 +228,8 @@ async def _server_lifespan(app: GuardrailsApp): (ModelInitializationError, model_initialization_error_handler), (StreamingNotSupportedError, bad_request_error_handler), (RailTypeNotConfiguredError, rail_type_not_configured_error_handler), + (RailTypeNotSupportedError, rail_type_not_supported_error_handler), + (InvalidCheckRequestError, invalid_check_request_error_handler), (InvalidStateError, invalid_state_error_handler), (RequestValidationError, validation_error_handler), (StarletteHTTPException, http_exception_handler), @@ -843,6 +849,23 @@ def _map_rail_status(status: RailStatus) -> str: return status.value +def _index_past_prepended_context(violation: ToolViolation) -> ToolViolation: + """*violation* with a tool result's message index counted in the request's own messages.""" + # A call's index is its position in tool_calls, which the prepended message does not move. + if violation.kind == "tool_result" and violation.index is not None: + return violation.model_copy(update={"index": violation.index - 1}) + return violation + + +def _violations_in_request_positions( + violations: Optional[List[ToolViolation]], context_prepended: bool +) -> Optional[List[ToolViolation]]: + """The tool violations with each index pointing into the messages the caller sent.""" + if violations is None or not context_prepended: + return violations + return [_index_past_prepended_context(violation) for violation in violations] + + @app.post( "/v1/checks", response_model=GuardrailCheckResponse, @@ -851,7 +874,9 @@ def _map_rail_status(status: RailStatus) -> str: async def guardrail_check(body: GuardrailCheckRequest, request: Request): """Guardrail check request. - Returns 422 when ``rail_types`` includes a type with no configured flows. + Returns 422 when ``rail_types`` includes a type with no configured flows, or when + the serving engine cannot run a tool check it was asked for: a tool rail type, or + ``tools``, on a config the IORails engine does not serve. """ api_request_headers.set(request.headers) @@ -880,16 +905,18 @@ async def guardrail_check(body: GuardrailCheckRequest, request: Request): ) messages = list(body.messages) - if body.guardrails.context: + context_prepended = bool(body.guardrails.context) + if context_prepended: messages.insert(0, {"role": "context", "content": body.guardrails.context}) - result = await llm_rails.check_async(messages=messages, rail_types=body.guardrails.rail_types) + result = await llm_rails.check_async(messages=messages, rail_types=body.guardrails.rail_types, tools=body.tools) return GuardrailCheckResponse( status=_map_rail_status(result.status), content=result.content, rail=result.rail, reason=result.reason, + tool_violations=_violations_in_request_positions(result.tool_violations, context_prepended), ) diff --git a/nemoguardrails/server/exception_handlers.py b/nemoguardrails/server/exception_handlers.py index 3cd42fff6d..5e84e0656b 100644 --- a/nemoguardrails/server/exception_handlers.py +++ b/nemoguardrails/server/exception_handlers.py @@ -24,10 +24,12 @@ from starlette.responses import JSONResponse, Response from nemoguardrails.exceptions import ( + InvalidCheckRequestError, InvalidStateError, LLMCallException, LLMRateLimitError, RailTypeNotConfiguredError, + RailTypeNotSupportedError, StreamingCapacityExceededError, StreamingNotSupportedError, ) @@ -126,6 +128,16 @@ async def rail_type_not_configured_error_handler(request: Request, exc: RailType return _error_response(422, str(exc)) +async def rail_type_not_supported_error_handler(request: Request, exc: RailTypeNotSupportedError) -> Response: + log.warning("Rail type not supported: %s", exc) + return _error_response(422, str(exc)) + + +async def invalid_check_request_error_handler(request: Request, exc: InvalidCheckRequestError) -> Response: + log.warning("Invalid check request: %s", exc) + return _error_response(422, str(exc)) + + async def invalid_state_error_handler(request: Request, exc: InvalidStateError) -> Response: log.warning("Invalid state: %s", exc) return _error_response(422, str(exc)) diff --git a/nemoguardrails/server/schemas/openai.py b/nemoguardrails/server/schemas/openai.py index 0cbf58aaf1..eb7994a47e 100644 --- a/nemoguardrails/server/schemas/openai.py +++ b/nemoguardrails/server/schemas/openai.py @@ -30,7 +30,7 @@ model_validator, ) -from nemoguardrails.rails.llm.options import GenerationOptions, RailType +from nemoguardrails.rails.llm.options import GenerationOptions, RailType, ToolViolation class GuardrailsDataOutput(BaseModel): @@ -383,13 +383,19 @@ class GuardrailCheckDataInput(GuardrailsDataInput): rail_types: Optional[List[RailType]] = Field( default=None, - description="Rail types to run. When omitted, auto-detected from message roles.", + description="Rail types to run. When omitted, input and output rails are auto-detected from message roles. " + "`tool_call` and `tool_result` run only when named, and only on the IORails engine.", ) class GuardrailCheckRequest(OpenAIChatCompletionRequest): """Request body for the /v1/checks endpoint.""" + tools: Optional[list[dict]] = Field( + default=None, + description="Tools that `tool_call` rails validate calls against: the allowlist and each tool's argument " + "schema. Only the IORails engine runs tool checks; on LLMRails a request with `tools` returns 422.", + ) guardrails: GuardrailCheckDataInput = Field( default_factory=GuardrailCheckDataInput, description="Guardrails specific options for the request.", @@ -405,3 +411,8 @@ class GuardrailCheckResponse(BaseModel): reason: Optional[str] = Field( default=None, description="Why the rail blocked the content, if the engine reports it." ) + tool_violations: Optional[List[ToolViolation]] = Field( + default=None, + description="The tool calls or tool results a tool rail blocked, each with a `violation_type` to switch on. " + "Present only when a tool rail blocked (IORails).", + ) diff --git a/tests/guardrails/test_engine_registry.py b/tests/guardrails/test_engine_registry.py index aa285f5eac..86c9a14d2c 100644 --- a/tests/guardrails/test_engine_registry.py +++ b/tests/guardrails/test_engine_registry.py @@ -38,6 +38,7 @@ from nemoguardrails.types import LLMModel, LLMResponse, LLMResponseChunk, UsageInfo from tests.guardrails.metric_helpers import collect_histogram_sum, collect_metric_points from tests.guardrails.test_data import NEMOGUARDS_CONFIG +from tests.guardrails.tool_helpers import tool_call_turn, wire_tool_call @pytest.fixture @@ -1408,6 +1409,17 @@ def test_extract_tool_exchanges_delegates_to_model_engine(self, manager): assert [(c.id, c.function.name, c.function.arguments) for c in calls] == [("c1", "get_weather", {"city": "X"})] assert [r.call_id for r in results] == ["c1"] + def test_extract_latest_tool_calls_delegates_to_model_engine(self, manager): + """The registry parses the last assistant message's tool calls through the named model engine.""" + messages = tool_call_turn(wire_tool_call(arguments='{"city": "X"}', call_id="c1")) + calls = manager.extract_latest_tool_calls("main", messages) + assert [(c.id, c.function.name, c.function.arguments) for c in calls] == [("c1", "get_weather", {"city": "X"})] + + def test_extract_latest_tool_calls_unknown_engine_raises_keyerror(self, manager): + """An unconfigured model type raises KeyError rather than guessing a wire format.""" + with pytest.raises(KeyError): + manager.extract_latest_tool_calls("nonexistent", []) + def test_parse_tools_unknown_engine_raises_keyerror(self, manager): with pytest.raises(KeyError): manager.parse_tools("nonexistent", {"tools": self._TOOLS}) diff --git a/tests/guardrails/test_guardrails.py b/tests/guardrails/test_guardrails.py index c35b87d199..6d9e33baa1 100644 --- a/tests/guardrails/test_guardrails.py +++ b/tests/guardrails/test_guardrails.py @@ -26,6 +26,7 @@ class correctly delegates method calls with properly formatted parameters. import pytest from nemoguardrails import Guardrails +from nemoguardrails.exceptions import RailTypeNotSupportedError from nemoguardrails.guardrails.compiled_rail import RailCompilationError, _is_installed, unservable_reason from nemoguardrails.guardrails.iorails import ( REFUSAL_MESSAGE, @@ -43,6 +44,9 @@ class correctly delegates method calls with properly formatted parameters. from nemoguardrails.types import LLMResponse from tests.guardrails.async_helpers import JAILBREAK_NIM_URL, mock_jailbreak_nim, mock_rail_model from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG, NEMOGUARDS_CONFIG, TOPIC_SAFETY_CONFIG +from tests.utils import check_entry_points, run_check + +_TOOLS = [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}] # Valid IORails input/output rails for has_only_iorails_flows tests _IORAILS_BASE_RAILS = { @@ -233,7 +237,7 @@ async def mock_stream(): async def test_check_delegates_to_iorails( self, mock_iorails_init, mock_start, mock_stop, _content_safety_rails_config ): - """check / check_async delegate to the IORails engine instead of raising.""" + """check / check_async delegate to the IORails engine, forwarding rail_types and tools.""" async with Guardrails(config=_content_safety_rails_config, verbose=False, use_iorails=True) as guardrails: assert isinstance(guardrails.rails_engine, IORails) @@ -243,13 +247,15 @@ async def test_check_delegates_to_iorails( messages = [{"role": "user", "content": "hello"}] - result = await guardrails.check_async(messages, rail_types=[RailType.INPUT]) + result = await guardrails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=_TOOLS) assert result is expected - guardrails.rails_engine.check_async.assert_awaited_once_with(messages, rail_types=[RailType.INPUT]) + guardrails.rails_engine.check_async.assert_awaited_once_with( + messages, rail_types=[RailType.TOOL_CALL], tools=_TOOLS + ) sync_result = guardrails.check(messages) assert sync_result is expected - guardrails.rails_engine.check.assert_called_once_with(messages, rail_types=None) + guardrails.rails_engine.check.assert_called_once_with(messages, rail_types=None, tools=None) @pytest.mark.asyncio @patch.object(LLMRails, "__init__", return_value=None) @@ -1543,7 +1549,7 @@ def test_check_delegates(self, mock_llmrails_class, _nemoguards_rails_config): messages = [{"role": "user", "content": "hi"}] result = guardrails.check(messages, rail_types=[RailType.INPUT]) - mock_llmrails_instance.check.assert_called_once_with(messages, rail_types=[RailType.INPUT]) + mock_llmrails_instance.check.assert_called_once_with(messages, rail_types=[RailType.INPUT], tools=None) assert result is sentinel @pytest.mark.asyncio @@ -1560,7 +1566,7 @@ async def test_check_async_delegates(self, mock_llmrails_class, _nemoguards_rail messages = [{"role": "user", "content": "hi"}] result = await guardrails.check_async(messages, rail_types=[RailType.OUTPUT]) - mock_llmrails_instance.check_async.assert_called_once_with(messages, rail_types=[RailType.OUTPUT]) + mock_llmrails_instance.check_async.assert_called_once_with(messages, rail_types=[RailType.OUTPUT], tools=None) assert result is sentinel @patch("nemoguardrails.guardrails.guardrails.LLMRails") @@ -2300,3 +2306,53 @@ def test_no_surface_is_refused_for_being_out_of_scope(self): def test_the_refusal_reason_names_the_limitation(self, flow, direction, expected): """Each class of refusal reports why, in wording a config author can act on.""" assert self._surface_reason(flow, direction) == expected + + +@pytest.fixture +def llmrails_guardrails(): + """A Guardrails facade serving a Colang input-rail config on LLMRails.""" + config = RailsConfig.from_content( + """ + define flow input rail + if $user_message == "block" + bot refuse to respond + stop + """, + """ + rails: + input: + flows: + - input rail + """, + ) + return Guardrails(config=config, use_iorails=False) + + +class TestCheckToolRailsOnLLMRails: + """On LLMRails, the facade fails loud for tool checks rather than ignoring what LLMRails cannot run.""" + + @check_entry_points + @pytest.mark.parametrize("tools", [None, _TOOLS], ids=["without_tools", "with_tools"]) + def test_tool_rail_types_raise_not_supported(self, llmrails_guardrails, entry_point, tools): + """Tool rail types raise ``RailTypeNotSupportedError`` from either entry point, with or without ``tools``.""" + with pytest.raises(RailTypeNotSupportedError): + run_check( + llmrails_guardrails, + entry_point, + [{"role": "user", "content": "hi"}], + rail_types=[RailType.TOOL_CALL, RailType.TOOL_RESULT], + tools=tools, + ) + + @check_entry_points + @pytest.mark.parametrize("rail_types", [None, [RailType.INPUT]], ids=["auto_detected", "input"]) + def test_tools_alone_raise_not_supported(self, llmrails_guardrails, entry_point, rail_types): + """``tools`` without a tool rail type raises ``RailTypeNotSupportedError`` from either entry point.""" + with pytest.raises(RailTypeNotSupportedError, match="tools"): + run_check( + llmrails_guardrails, + entry_point, + [{"role": "user", "content": "hi"}], + rail_types=rail_types, + tools=_TOOLS, + ) diff --git a/tests/guardrails/test_iorails_check.py b/tests/guardrails/test_iorails_check.py index 91e5512eb6..2f753aadd5 100644 --- a/tests/guardrails/test_iorails_check.py +++ b/tests/guardrails/test_iorails_check.py @@ -22,13 +22,13 @@ import asyncio import logging -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio -from nemoguardrails.exceptions import RailTypeNotConfiguredError -from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.exceptions import InvalidCheckRequestError, RailTypeNotConfiguredError +from nemoguardrails.guardrails.guardrails_types import RailDirection, RailResult from nemoguardrails.guardrails.iorails import ( INTERNAL_ERROR_MESSAGE, REFUSAL_MESSAGE, @@ -38,9 +38,27 @@ ) from nemoguardrails.guardrails.rail_guard import rail_error_outcome from nemoguardrails.rails.llm.config import RailsConfig -from nemoguardrails.rails.llm.options import RailStatus, RailType +from nemoguardrails.rails.llm.options import RailStatus, RailType, ToolViolationType +from tests.guardrails.async_helpers import started_iorails from tests.guardrails.rail_stubs import bot_message_rewrite, rail_failure, user_message_rewrite from tests.guardrails.test_data import NEMOGUARDS_CONFIG, TOPIC_SAFETY_CONFIG +from tests.guardrails.test_tool_rails_iorails import ( + CONFIG_TOOLS_CONFIG, + TOOL_CONFIG, + WEATHER_TOOL, + _inject_forbidden_transport, +) +from tests.guardrails.tool_helpers import ( + TOOL_CALL_QUESTION, + UNREAD_TOOLS_MESSAGE, + call_violation, + make_tool_conversation, + malformed_prior_tool_call_messages, + multi_turn_reused_call_id_messages, + result_violation, + tool_call_turn, + wire_tool_call, +) SAFE = RailResult.allow() @@ -852,3 +870,503 @@ async def test_auto_detect_skips_validation(self, input_only_engine): [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] ) assert result.status == RailStatus.PASSED + + +_SECRET_VALUE = "SECRET-VALUE" + +# The Nemoguard input and output rails plus the global tool validators, for checks that mix rail families. +TOOL_AND_IO_CONFIG = { + **NEMOGUARDS_CONFIG, + "rails": { + **NEMOGUARDS_CONFIG["rails"], + "tool_output": {"flows": ["tool call validation"]}, + "tool_input": {"flows": ["tool result validation"]}, + }, +} + +PER_TOOL_ONLY_CONFIG = { + **TOOL_CONFIG, + "rails": { + "config": {"regex_detection": {"tool_output": {"run_sql": {"patterns": [r"DROP\s+TABLE"]}}}}, + "tool_output": {"per_tool": {"run_sql": ["regex check tool output"]}}, + }, +} + +_CLOSED_RUN_SQL = { + "type": "function", + "function": { + "name": "run_sql", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "additionalProperties": False}, + }, +} + + +@pytest_asyncio.fixture +async def tool_iorails(): + """A started IORails with only the global tool validators, whose model must never be called.""" + async with started_iorails(TOOL_CONFIG) as engine: + _inject_forbidden_transport(engine) + yield engine + + +@pytest_asyncio.fixture +async def config_tools_iorails(): + """Like ``tool_iorails``, with ``get_weather`` declared on the main model's parameters.""" + async with started_iorails(CONFIG_TOOLS_CONFIG) as engine: + _inject_forbidden_transport(engine) + yield engine + + +@pytest_asyncio.fixture +async def configured_tool_iorails(request): + """A started IORails for the config a test passes indirectly, whose model must never be called.""" + async with started_iorails(request.param) as engine: + _inject_forbidden_transport(engine) + yield engine + + +@pytest_asyncio.fixture +async def tool_and_io_iorails(): + """A started IORails with input, output and tool rails; the input and output rails are stubbed per test.""" + async with started_iorails(TOOL_AND_IO_CONFIG) as engine: + _inject_forbidden_transport(engine) + yield engine + + +_GET_TIME_TOOL = {"type": "function", "function": {"name": "get_time", "parameters": {"type": "object"}}} + +_FINAL_TEXT_TURN = [*make_tool_conversation(), {"role": "assistant", "content": "It's 18C in Paris."}] + +# Case id -> (engine config, messages, rail type, tools) for a tool check that finds nothing to block. +_PASSING_TOOL_CHECKS = { + "declared_call_with_valid_arguments": ( + TOOL_CONFIG, + tool_call_turn(wire_tool_call()), + RailType.TOOL_CALL, + [WEATHER_TOOL], + ), + "config_declared_tools_when_tools_is_none": ( + CONFIG_TOOLS_CONFIG, + tool_call_turn(wire_tool_call()), + RailType.TOOL_CALL, + None, + ), + "earlier_malformed_turn_not_rechecked": ( + TOOL_CONFIG, + malformed_prior_tool_call_messages(), + RailType.TOOL_CALL, + [WEATHER_TOOL], + ), + "final_text_turn_has_no_calls": (TOOL_CONFIG, _FINAL_TEXT_TURN, RailType.TOOL_CALL, []), + "nameless_result_linked_by_call_id": ( + TOOL_CONFIG, + make_tool_conversation(result_name=None), + RailType.TOOL_RESULT, + None, + ), + "call_ids_reused_across_turns": (TOOL_CONFIG, multi_turn_reused_call_id_messages(), RailType.TOOL_RESULT, None), +} + +_PARIS = '{"city": "Paris"}' + +# Case id -> (engine config, last call's arguments, tools, (violation type, argument path)). +_BLOCKING_TOOL_CALL_CHECKS = { + "schema_invalid_arguments": (TOOL_CONFIG, "{}", [WEATHER_TOOL], ("arguments_invalid", "/city")), + "empty_string_arguments_are_malformed": (TOOL_CONFIG, "", [WEATHER_TOOL], ("malformed_arguments", None)), + "empty_tools_block_every_call": (CONFIG_TOOLS_CONFIG, _PARIS, [], ("tool_not_allowed", None)), + "request_tools_replace_config_tools": (CONFIG_TOOLS_CONFIG, _PARIS, [_GET_TIME_TOOL], ("tool_not_allowed", None)), + "duplicate_tool_definitions": (TOOL_CONFIG, _PARIS, [WEATHER_TOOL, WEATHER_TOOL], ("invalid_toolset", None)), +} + + +class TestCheckToolRailsPass: + """Tool checks that find nothing to block.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("configured_tool_iorails", "messages", "rail_type", "tools"), + list(_PASSING_TOOL_CHECKS.values()), + ids=list(_PASSING_TOOL_CHECKS), + indirect=["configured_tool_iorails"], + ) + async def test_tool_check_passes(self, configured_tool_iorails, messages, rail_type, tools): + """A declared, well-formed call or a linked result passes with no tool violations.""" + result = await configured_tool_iorails.check_async(messages, rail_types=[rail_type], tools=tools) + + assert (result.status, result.tool_violations) == (RailStatus.PASSED, None) + + +class TestCheckToolCalls: + """``rail_types=[RailType.TOOL_CALL]`` validates the last assistant message's tool calls.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("configured_tool_iorails", "arguments", "tools", "expected"), + list(_BLOCKING_TOOL_CALL_CHECKS.values()), + ids=list(_BLOCKING_TOOL_CALL_CHECKS), + indirect=["configured_tool_iorails"], + ) + async def test_tool_call_check_blocks_with_one_violation(self, configured_tool_iorails, arguments, tools, expected): + """A call the declared tools or its schema reject blocks with one violation of the expected type.""" + violation_type, argument_path = expected + + result = await configured_tool_iorails.check_async( + tool_call_turn(wire_tool_call(arguments=arguments)), rail_types=[RailType.TOOL_CALL], tools=tools + ) + + assert result.status == RailStatus.BLOCKED + assert [(v.violation_type, v.argument_path) for v in result.tool_violations] == [ + (ToolViolationType(violation_type), argument_path) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("configured_tool_iorails", "rail"), + [(TOOL_CONFIG, "tool call validation"), (PER_TOOL_ONLY_CONFIG, "regex check tool output")], + ids=["global_validator", "per_tool_rail_only"], + indirect=["configured_tool_iorails"], + ) + @pytest.mark.parametrize( + ("arguments", "tools", "expected"), + [ + ('{"query": "SELECT 1"}', [], (ToolViolationType.TOOL_NOT_ALLOWED, None, None)), + ('{"query": 5}', [_CLOSED_RUN_SQL], (ToolViolationType.ARGUMENTS_INVALID, "/query", "type")), + ], + ids=["undeclared_tool", "schema_invalid_arguments"], + ) + async def test_a_bad_call_is_categorized_alike_whichever_rails_check_it( + self, configured_tool_iorails, rail, arguments, tools, expected + ): + """The global validator and a per-tool rail's gate report the same violation; only the rail differs.""" + result = await configured_tool_iorails.check_async( + tool_call_turn(wire_tool_call("run_sql", arguments)), rail_types=[RailType.TOOL_CALL], tools=tools + ) + + assert [(v.violation_type, v.argument_path, v.schema_keyword) for v in result.tool_violations] == [expected] + assert [(v.index, v.tool_call_id) for v in result.tool_violations] == [(0, "call_1")] + assert result.rail == rail + + @pytest.mark.asyncio + async def test_undeclared_call_blocks(self, tool_iorails): + """An undeclared call blocks, naming the rail, the reason, and the call in a ``tool_not_allowed`` violation.""" + result = await tool_iorails.check_async( + tool_call_turn(wire_tool_call("delete_files", "{}")), rail_types=[RailType.TOOL_CALL], tools=[WEATHER_TOOL] + ) + + assert (result.status, result.content) == (RailStatus.BLOCKED, REFUSAL_MESSAGE) + assert result.rail == "tool call validation" + assert result.reason == "tool call 'delete_files' is not an allowed tool" + assert result.tool_violations == [ + call_violation( + "tool_not_allowed", + "tool call 'delete_files' is not an allowed tool", + tool_call_id="call_1", + tool_name="delete_files", + index=0, + ) + ] + + @pytest.mark.asyncio + async def test_malformed_arguments_block_without_hiding_later_calls(self, tool_iorails): + """Malformed arguments block with ``malformed_arguments``, never quoted, and the later calls are still checked.""" + messages = tool_call_turn( + wire_tool_call(arguments='{"city": "SECRET-VALUE"'), wire_tool_call("delete_files", "{}", "call_2") + ) + + result = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[WEATHER_TOOL]) + + assert (result.status, result.rail) == (RailStatus.BLOCKED, None) + assert result.reason == "tool call extraction failed: tool call 'call_1' has malformed arguments" + assert [(v.violation_type, v.index, v.tool_call_id) for v in result.tool_violations] == [ + (ToolViolationType.MALFORMED_ARGUMENTS, 0, "call_1"), + (ToolViolationType.TOOL_NOT_ALLOWED, 1, "call_2"), + ] + assert _SECRET_VALUE not in result.model_dump_json() + + +class TestCheckToolResults: + """``rail_types=[RailType.TOOL_RESULT]`` validates every tool result against the calls it links to.""" + + @pytest.mark.asyncio + async def test_unlinked_result_blocks(self, tool_iorails): + """A result for an unknown call id blocks, and its violation points at the tool message's position.""" + messages = make_tool_conversation(result_call_id="call_999", result_name=None) + + result = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_RESULT]) + + assert (result.status, result.rail) == (RailStatus.BLOCKED, "tool result validation") + assert result.reason == "tool result for call_id 'call_999' does not correspond to a prior tool call" + assert result.tool_violations == [ + result_violation( + "unknown_call_id", + "tool result for call_id 'call_999' does not correspond to a prior tool call", + tool_call_id="call_999", + index=2, + ) + ] + + +class TestCheckToolRoundTrip: + """A harness checks its model's tool calls, runs the tools, then checks the results, all without generation.""" + + @pytest.mark.asyncio + async def test_harness_round_trip_passes_both_checks(self, tool_iorails): + """A declared call passes the tool-call check, and the nameless result the harness appends passes next.""" + messages = tool_call_turn(wire_tool_call()) + + call_check = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[WEATHER_TOOL]) + messages.append({"role": "tool", "tool_call_id": "call_1", "content": "18C"}) + result_check = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_RESULT]) + + assert (call_check.status, result_check.status) == (RailStatus.PASSED, RailStatus.PASSED) + + @pytest.mark.asyncio + async def test_result_for_an_unknown_call_blocks_the_second_check(self, tool_iorails): + """A passing tool-call check does not vouch for a result that answers a different call id.""" + messages = tool_call_turn(wire_tool_call()) + + call_check = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[WEATHER_TOOL]) + messages.append({"role": "tool", "tool_call_id": "call_2", "content": "18C"}) + result_check = await tool_iorails.check_async(messages, rail_types=[RailType.TOOL_RESULT]) + + assert (call_check.status, result_check.status) == (RailStatus.PASSED, RailStatus.BLOCKED) + + +def _record_calls(engine: IORails, *, input_result: RailResult = SAFE) -> MagicMock: + """Spy on each rail family's manager entry point; the returned mock's calls give the order they ran in.""" + manager = engine.rails_manager + order = MagicMock() + spies = { + "are_tool_results_safe": ("tool_result", AsyncMock(wraps=manager.are_tool_results_safe)), + "is_input_safe": ("input", AsyncMock(return_value=input_result)), + "are_latest_tool_calls_safe": ("tool_call", AsyncMock(wraps=manager.are_latest_tool_calls_safe)), + "is_output_safe": ("output", AsyncMock(return_value=SAFE)), + } + for method, (family, spy) in spies.items(): + order.attach_mock(spy, family) + setattr(manager, method, spy) + return order + + +def _families_run(order: MagicMock) -> list[str]: + """The rail families a check ran, in order, from the spy ``_record_calls`` returned.""" + return [name for name, _args, _kwargs in order.mock_calls] + + +# A user turn, an assistant turn with text and a declared call, and the call's result: every family has work. +FULL_TURN = [ + *tool_call_turn(wire_tool_call(), content="Let me check."), + {"role": "tool", "tool_call_id": "call_1", "content": "18C"}, +] +ALL_RAIL_TYPES = [RailType.INPUT, RailType.OUTPUT, RailType.TOOL_CALL, RailType.TOOL_RESULT] + + +class TestCheckToolRailOrchestration: + """How a check sequences tool rails with input and output rails, and what it reports.""" + + @pytest.mark.asyncio + async def test_rail_families_run_as_generation_orders_them(self, tool_and_io_iorails): + """Tool results, input, tool calls, then output: the order generation runs them in.""" + order = _record_calls(tool_and_io_iorails) + + result = await tool_and_io_iorails.check_async(FULL_TURN, rail_types=ALL_RAIL_TYPES, tools=[WEATHER_TOOL]) + + assert result.status == RailStatus.PASSED + assert _families_run(order) == ["tool_result", "input", "tool_call", "output"] + + @pytest.mark.asyncio + async def test_a_tool_call_block_stops_before_output_rails(self, tool_and_io_iorails): + """A tool-call block ends the check, so the output rails never run.""" + order = _record_calls(tool_and_io_iorails) + + result = await tool_and_io_iorails.check_async(FULL_TURN, rail_types=ALL_RAIL_TYPES, tools=[]) + + assert result.rail == "tool call validation" + assert _families_run(order) == ["tool_result", "input", "tool_call"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rail_type", "tools", "expected_content"), + [(RailType.TOOL_CALL, [WEATHER_TOOL], "Let me check."), (RailType.TOOL_RESULT, None, TOOL_CALL_QUESTION)], + ids=["tool_call_reports_assistant_text", "tool_result_reports_user_text"], + ) + async def test_tool_check_reports_the_text_of_its_direction( + self, tool_and_io_iorails, rail_type, tools, expected_content + ): + """A tool-call check reports the assistant text, as an output check does; a tool-result check the user text.""" + result = await tool_and_io_iorails.check_async(FULL_TURN, rail_types=[rail_type], tools=tools) + + assert (result.status, result.content) == (RailStatus.PASSED, expected_content) + + @pytest.mark.asyncio + async def test_input_rewrite_stays_visible_with_a_tool_call_check(self, tool_and_io_iorails): + """With ``[INPUT, TOOL_CALL]``, an input rewrite is what the check reports, as MODIFIED.""" + _record_calls(tool_and_io_iorails, input_result=user_message_rewrite("What's the weather in ?")) + + result = await tool_and_io_iorails.check_async( + FULL_TURN, rail_types=[RailType.INPUT, RailType.TOOL_CALL], tools=[WEATHER_TOOL] + ) + + assert (result.status, result.content) == (RailStatus.MODIFIED, "What's the weather in ?") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rail_types", "tools", "direction"), + [([RailType.TOOL_RESULT], None, RailDirection.INPUT), ([RailType.TOOL_CALL], [], RailDirection.OUTPUT)], + ids=["tool_result", "tool_call"], + ) + async def test_a_tool_block_records_the_directional_block_metric(self, tool_iorails, rail_types, tools, direction): + """A tool-result block counts as an input block and a tool-call block as an output block.""" + tool_iorails._metrics_enabled = True + messages = make_tool_conversation(result_call_id="call_999") + + with patch("nemoguardrails.guardrails.iorails.record_request_blocked") as record_blocked: + result = await tool_iorails.check_async(messages, rail_types=rail_types, tools=tools) + + assert result.status == RailStatus.BLOCKED + record_blocked.assert_called_once_with(direction) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rail_types", "blocked_rail"), + [([RailType.INPUT], "input"), ([RailType.OUTPUT], "output")], + ) + async def test_input_or_output_block_has_no_tool_violations(self, iorails, rail_types, blocked_rail): + """Only a tool rail block reports tool violations.""" + blocking = _unsafe(f"content safety check {blocked_rail}") + _mock_rails(iorails, input_result=blocking, output_result=blocking) + + result = await iorails.check_async(CONVERSATION, rail_types=rail_types) + + assert result.status == RailStatus.BLOCKED + assert result.tool_violations is None + + +class TestCheckToolRailConfiguration: + """Which configs can run a tool check, and what happens when rail types are left to auto-detection.""" + + @pytest.mark.asyncio + async def test_unconfigured_tool_rail_type_raises(self, iorails): + """A tool rail type whose config section has no rails raises ``RailTypeNotConfiguredError``.""" + with pytest.raises(RailTypeNotConfiguredError, match="rail type 'tool_call' has no configured rails"): + await iorails.check_async(tool_call_turn(wire_tool_call()), rail_types=[RailType.TOOL_CALL]) + + @pytest.mark.asyncio + async def test_per_tool_rails_alone_count_as_configured(self): + """A section with only ``per_tool`` rails runs them, and a per-tool block names the call it blocked.""" + messages = tool_call_turn(wire_tool_call("run_sql", '{"query": "DROP TABLE users"}')) + run_sql = { + "type": "function", + "function": {"name": "run_sql", "parameters": {"type": "object", "additionalProperties": True}}, + } + + async with started_iorails(PER_TOOL_ONLY_CONFIG) as engine: + result = await engine.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[run_sql]) + + assert result.rail == "regex check tool output" + assert [(v.violation_type, v.rail, v.index) for v in result.tool_violations] == [ + (ToolViolationType.PER_TOOL_RAIL, "regex check tool output", 0) + ] + + @pytest.mark.asyncio + async def test_per_tool_rails_alone_leave_unlisted_tools_unchecked(self): + """Without the global validator only tools with per-tool rails are checked, so an undeclared tool passes.""" + messages = tool_call_turn(wire_tool_call("delete_files", "{}")) + + async with started_iorails(PER_TOOL_ONLY_CONFIG) as engine: + result = await engine.check_async(messages, rail_types=[RailType.TOOL_CALL], tools=[]) + + assert result.status == RailStatus.PASSED + + @pytest.mark.asyncio + async def test_per_tool_rails_with_no_flows_are_not_configured(self): + """A ``per_tool`` map whose lists are all empty runs nothing, so the tool rail type counts as unconfigured.""" + config = {**TOOL_CONFIG, "rails": {"tool_output": {"per_tool": {"run_sql": []}}}} + + async with started_iorails(config) as engine: + with pytest.raises(RailTypeNotConfiguredError, match="rail type 'tool_call' has no configured rails"): + await engine.check_async(tool_call_turn(wire_tool_call()), rail_types=[RailType.TOOL_CALL]) + + @pytest.mark.asyncio + async def test_tool_rail_type_without_a_main_model_raises(self): + """Without a ``main`` model there is no wire format to read, so a tool check refuses to run.""" + config = {"models": [], "rails": TOOL_CONFIG["rails"]} + + async with started_iorails(config) as engine: + with pytest.raises(RailTypeNotConfiguredError, match="needs a `main` model"): + await engine.check_async(tool_call_turn(wire_tool_call()), rail_types=[RailType.TOOL_CALL]) + + @pytest.mark.asyncio + async def test_auto_detection_runs_input_and_output_only_and_logs_a_hint(self, tool_and_io_iorails, caplog): + """With ``rail_types`` omitted, tool traffic runs only input and output rails, with an INFO hint.""" + order = _record_calls(tool_and_io_iorails) + + with caplog.at_level(logging.INFO, logger="nemoguardrails.guardrails.iorails"): + result = await tool_and_io_iorails.check_async(FULL_TURN) + + assert result.status == RailStatus.PASSED + assert _families_run(order) == ["input", "output"] + assert "tool rails are configured but were not requested" in caplog.text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("messages", "hint_count"), + [ + ([{"role": "user", "content": "hi"}, {"role": "tool", "tool_call_id": "call_1", "content": "18C"}], 1), + ([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}], 0), + ], + ids=["tool_result_only", "no_tool_traffic"], + ) + async def test_auto_detection_hints_only_at_tool_traffic(self, tool_and_io_iorails, caplog, messages, hint_count): + """The INFO hint appears when the messages carry a tool result, and not when they carry no tool traffic.""" + _record_calls(tool_and_io_iorails) + + with caplog.at_level(logging.INFO, logger="nemoguardrails.guardrails.iorails"): + await tool_and_io_iorails.check_async(messages) + + assert caplog.text.count("tool rails are configured but were not requested") == hint_count + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rail_types", + [None, [], [RailType.INPUT], [RailType.OUTPUT, RailType.TOOL_RESULT]], + ids=["auto_detection", "empty", "input", "output_and_tool_result"], + ) + async def test_tools_without_a_tool_call_check_raise(self, tool_and_io_iorails, rail_types): + """Only a tool_call check reads ``tools``, so any other check given them raises before a rail runs.""" + order = _record_calls(tool_and_io_iorails) + + with pytest.raises(InvalidCheckRequestError, match=f"^{UNREAD_TOOLS_MESSAGE}$"): + await tool_and_io_iorails.check_async(FULL_TURN, rail_types=rail_types, tools=[WEATHER_TOOL]) + + assert _families_run(order) == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rail_types", "tools", "error"), + [ + ([RailType.INPUT], [WEATHER_TOOL], InvalidCheckRequestError), + ([RailType.TOOL_CALL], None, RailTypeNotConfiguredError), + ], + ids=["unread_tools", "unconfigured_rail_type"], + ) + async def test_request_error_is_raised_before_queueing(self, iorails, caplog, rail_types, tools, error): + """A request the check cannot serve raises before it is queued, so it takes no queue slot and logs no ERROR.""" + with patch.object(iorails._generate_async_queue, "submit") as submit, caplog.at_level(logging.WARNING): + with pytest.raises(error): + await iorails.check_async(CONVERSATION, rail_types=rail_types, tools=tools) + + submit.assert_not_called() + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] == [] + + def test_sync_check_forwards_tools(self): + """Sync ``check`` hands ``tools`` to the engine it spins up.""" + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + engine = IORails(RailsConfig.from_content(config=TOOL_CONFIG)) + + with patch("nemoguardrails.guardrails.iorails.IORails", return_value=engine): + result = engine.check( + tool_call_turn(wire_tool_call()), rail_types=[RailType.TOOL_CALL], tools=[WEATHER_TOOL] + ) + + assert result.status == RailStatus.PASSED diff --git a/tests/guardrails/test_model_engine.py b/tests/guardrails/test_model_engine.py index fabc0077b4..bb135188b2 100644 --- a/tests/guardrails/test_model_engine.py +++ b/tests/guardrails/test_model_engine.py @@ -45,10 +45,11 @@ _parse_chat_completion, _parse_chat_completion_chunk, ) -from nemoguardrails.guardrails.tool_schema import Toolset +from nemoguardrails.guardrails.tool_schema import ToolCallExtractionError, Toolset from nemoguardrails.llm.call import llm_call from nemoguardrails.llm.clients._errors import client_facing_message from nemoguardrails.rails.llm.config import Model +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.types import ( ChatMessage, LLMModel, @@ -59,7 +60,13 @@ ToolCallFunction, UsageInfo, ) -from tests.guardrails.tool_helpers import make_tool_conversation, multi_turn_reused_call_id_messages +from tests.guardrails.tool_helpers import ( + assistant_tool_calls, + make_tool_conversation, + malformed_prior_tool_call_messages, + multi_turn_reused_call_id_messages, + wire_tool_call, +) def _make_model( @@ -2894,6 +2901,13 @@ def test_skips_non_dict_messages(self): engine = ModelEngine(_make_model(engine="openai")) assert self._ids(engine.extract_tool_exchanges(["garbage", *_TOOL_MESSAGES])) == [(["call_1"], ["call_1"])] + def test_records_each_results_message_index(self): + """Each result records its message's position in the conversation, counting non-dict and context messages.""" + engine = ModelEngine(_make_model(engine="openai")) + messages = ["garbage", {"role": "context", "content": {"user_id": "u1"}}, *multi_turn_reused_call_id_messages()] + exchanges = engine.extract_tool_exchanges(messages) + assert [result.message_index for _calls, results in exchanges for result in results] == [4, 8] + def test_malformed_arguments_degrade_to_empty_for_linkage(self): """A historical call with malformed argument JSON degrades to empty arguments (id/name preserved) instead of aborting extraction.""" @@ -2921,6 +2935,169 @@ def test_unknown_engine_falls_back_to_openai_extractor(self): assert self._ids(engine.extract_tool_exchanges(_TOOL_MESSAGES)) == [(["call_1"], ["call_1"])] +_SECRET_VALUE = "SECRET-VALUE" + + +def _with_second_call_arguments(arguments: Any) -> dict: + """An assistant turn whose second call, ``call_2``, carries *arguments*.""" + return assistant_tool_calls(wire_tool_call(), wire_tool_call(call_id="call_2", arguments=arguments)) + + +_LEGACY_CALL = {"name": "get_weather", "arguments": "{}"} +_FLAT_CALL = {"id": "call_1", "type": "function", "name": "run_sql", "arguments": '{"q": "SECRET-VALUE"}'} + +# Case id -> (a last assistant turn with one call the extractor refuses, (violation type, index, tool_call_id, tool_name)). +_UNPARSABLE_CALLS = { + "invalid_json_arguments": ( + _with_second_call_arguments('{"city": "SECRET-VALUE"'), + ("malformed_arguments", 1, "call_2", "get_weather"), + ), + "invalid_json_before_a_valid_call": ( + assistant_tool_calls(wire_tool_call(arguments='{"city": "SECRET-VALUE"'), wire_tool_call(call_id="call_2")), + ("malformed_arguments", 0, "call_1", "get_weather"), + ), + "json_array_arguments": ( + _with_second_call_arguments('["SECRET-VALUE"]'), + ("malformed_arguments", 1, "call_2", "get_weather"), + ), + "non_dict_arguments": ( + _with_second_call_arguments([_SECRET_VALUE]), + ("malformed_arguments", 1, "call_2", "get_weather"), + ), + # A completed turn carries "{}" for a call without arguments; "" is not JSON, as in non-streaming generation. + "empty_string_arguments": (_with_second_call_arguments(""), ("malformed_arguments", 1, "call_2", "get_weather")), + "non_dict_entry": (assistant_tool_calls(wire_tool_call(), _SECRET_VALUE), ("malformed_tool_call", 1, None, None)), + "non_dict_function": ( + assistant_tool_calls({**wire_tool_call(), "function": _SECRET_VALUE}), + ("malformed_tool_call", 0, "call_1", None), + ), + "no_function_object": (assistant_tool_calls(_FLAT_CALL), ("malformed_tool_call", 0, "call_1", None)), + "non_string_function_name": ( + assistant_tool_calls({**wire_tool_call(), "function": {"name": {"x": 1}, "arguments": "{}"}}), + ("malformed_tool_call", 0, "call_1", None), + ), + "empty_function_name": ( + assistant_tool_calls({**wire_tool_call(), "function": {"name": "", "arguments": "{}"}}), + ("malformed_tool_call", 0, "call_1", None), + ), + "non_string_type": ( + assistant_tool_calls({**wire_tool_call(), "type": {"x": 1}}), + ("malformed_tool_call", 0, "call_1", "get_weather"), + ), + "non_string_id": ( + assistant_tool_calls({**wire_tool_call(), "id": {"x": 1}}), + ("malformed_tool_call", 0, None, "get_weather"), + ), +} + +# Case id -> (a last assistant turn carrying a legacy function_call, the tool name its error reports). +_LEGACY_FUNCTION_CALL_TURNS = { + "legacy_function_call": ({"role": "assistant", "function_call": _LEGACY_CALL}, "get_weather"), + "legacy_function_call_beside_tool_calls": ( + {**assistant_tool_calls(wire_tool_call()), "function_call": {**_LEGACY_CALL, "name": "delete_files"}}, + "delete_files", + ), +} + + +class TestExtractLatestToolCalls: + @pytest.fixture + def engine(self) -> ModelEngine: + """An OpenAI-engine ModelEngine; extraction needs no started client.""" + return ModelEngine(_make_model(engine="openai")) + + def test_parses_the_last_assistant_messages_tool_calls(self, engine): + """The last assistant message's tool calls are returned with their JSON arguments parsed into a dict.""" + calls = engine.extract_latest_tool_calls(make_tool_conversation()) + assert [(c.id, c.type, c.function.name, c.function.arguments) for c in calls] == [ + ("call_1", "function", "get_weather", {"city": "Paris"}) + ] + + def test_nim_uses_the_same_shape(self): + """NIM reads the OpenAI Chat Completions tool-call shape.""" + engine = ModelEngine(_make_model(engine="nim")) + assert [c.id for c in engine.extract_latest_tool_calls(make_tool_conversation())] == ["call_1"] + + def test_unknown_engine_falls_back_to_openai_extractor(self): + """An engine with no registered extractor reads the OpenAI Chat Completions shape.""" + engine = ModelEngine(_make_model(engine="vllm", parameters={"base_url": "http://localhost:8000"})) + assert [c.id for c in engine.extract_latest_tool_calls(make_tool_conversation())] == ["call_1"] + + def test_only_the_last_assistant_message_is_parsed(self, engine): + """Tool calls on an earlier assistant message are neither returned nor parsed, even when malformed.""" + calls = engine.extract_latest_tool_calls(malformed_prior_tool_call_messages()) + assert [(c.id, c.function.arguments) for c in calls] == [("call_1", {"city": "London"})] + + def test_context_message_is_ignored(self, engine): + """A server ``context`` message, whose role ChatMessage rejects, does not break extraction.""" + messages = [{"role": "context", "content": {"user_id": "u1"}}, *make_tool_conversation()] + assert [c.id for c in engine.extract_latest_tool_calls(messages)] == ["call_1"] + + def test_last_assistant_message_without_tool_calls_returns_empty(self, engine): + """A final assistant text turn has no calls to check, even when an earlier turn made some.""" + messages = [*make_tool_conversation(), {"role": "assistant", "content": "It's 18C in Paris."}] + assert engine.extract_latest_tool_calls(messages) == [] + + def test_no_assistant_message_returns_empty(self, engine): + """A conversation with no assistant message has no tool calls.""" + assert engine.extract_latest_tool_calls([]) == [] + assert engine.extract_latest_tool_calls([{"role": "user", "content": "hi"}]) == [] + + def test_missing_arguments_become_an_empty_object(self, engine): + """A call with no ``arguments`` key parses as ``{}``.""" + messages = [assistant_tool_calls({"id": "call_1", "type": "function", "function": {"name": "list_files"}})] + assert [c.function.arguments for c in engine.extract_latest_tool_calls(messages)] == [{}] + + def test_call_without_an_id_is_named_by_its_position(self, engine): + """A malformed call with no id is named in its error by its position in ``tool_calls``.""" + call = {"type": "function", "function": {"name": "get_weather", "arguments": "not json"}} + + [error] = engine.extract_latest_tool_calls([assistant_tool_calls(call)]) + + assert str(error) == "tool call at index 0 has malformed arguments" + + def test_long_call_id_is_capped_in_the_message_but_kept_whole(self, engine): + """A call id, which the model chose, is cut to 64 characters in the message but kept whole as identity.""" + call_id = "i" * 200 + + [error] = engine.extract_latest_tool_calls([assistant_tool_calls(wire_tool_call("f", "not json", call_id))]) + + assert (str(error), error.tool_call_id) == (f"tool call '{'i' * 64}...' has malformed arguments", call_id) + + @pytest.mark.parametrize(("message", "expected"), list(_UNPARSABLE_CALLS.values()), ids=list(_UNPARSABLE_CALLS)) + def test_unparsable_call_is_returned_as_its_error_in_place(self, engine, message, expected): + """An unparsable call's error takes its place, naming it without quoting its text; every other call parses.""" + violation_type, index, tool_call_id, tool_name = expected + + extracted = engine.extract_latest_tool_calls([message]) + + error = extracted[index] + assert isinstance(error, ToolCallExtractionError) + assert (error.violation_type, error.index, error.tool_call_id, error.tool_name) == ( + ToolViolationType(violation_type), + index, + tool_call_id, + tool_name, + ) + assert _SECRET_VALUE not in str(error) + assert all(isinstance(entry, ToolCall) for position, entry in enumerate(extracted) if position != index) + + @pytest.mark.parametrize( + ("message", "tool_name"), list(_LEGACY_FUNCTION_CALL_TURNS.values()), ids=list(_LEGACY_FUNCTION_CALL_TURNS) + ) + def test_legacy_function_call_raises(self, engine, message, tool_name): + """A legacy ``function_call`` makes the whole turn unvalidatable, so extraction raises.""" + with pytest.raises(ToolCallExtractionError) as excinfo: + engine.extract_latest_tool_calls([message]) + + error = excinfo.value + assert (error.violation_type, error.index, error.tool_name) == ( + ToolViolationType.LEGACY_FUNCTION_CALL, + None, + tool_name, + ) + + class TestModelEngineLLMModelProtocol: """ModelEngine implements LLMModel, the interface library rail actions call through.""" diff --git a/tests/guardrails/test_per_tool_regex_rails.py b/tests/guardrails/test_per_tool_regex_rails.py index f9b12f0dc5..a83c83f585 100644 --- a/tests/guardrails/test_per_tool_regex_rails.py +++ b/tests/guardrails/test_per_tool_regex_rails.py @@ -37,9 +37,16 @@ from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.manifests import RailDirection as SurfaceDirection from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.tracing.constants import GuardrailsAttributes from nemoguardrails.types import ToolCall, ToolCallFunction -from tests.guardrails.tool_helpers import assert_result_blocked +from tests.guardrails.tool_helpers import ( + assert_result_blocked, + call_violation, + result_violation, + tool_call_turn, + wire_tool_call, +) STACK_CONFIG = {"models": [{"type": "main", "engine": "nim", "model": "meta/llama-3.3-70b-instruct"}]} @@ -50,33 +57,33 @@ "tool_input": {"run_sql": {"patterns": [r"ssn:\s*\d{3}-\d{2}-\d{4}"]}}, } +# A run_sql schema with no "nope" argument, so a per-tool rail scoped to `$argument=nope` raises. +_CLOSED_QUERY_SCHEMA = {"type": "object", "properties": {"query": {"type": "string"}}, "additionalProperties": False} + def _sql_call(query: str, call_id: str = "call_1") -> ToolCall: return ToolCall(id=call_id, type="function", function=ToolCallFunction(name="run_sql", arguments={"query": query})) +def _call_with_arguments(name: str, arguments: dict, call_id: str) -> ToolCall: + return ToolCall(id=call_id, type="function", function=ToolCallFunction(name=name, arguments=arguments)) + + def _other_call(call_id: str = "call_2") -> ToolCall: return ToolCall( id=call_id, type="function", function=ToolCallFunction(name="list_tables", arguments={"query": "DROP TABLE x"}) ) -def _llm_params(*tool_names: str) -> dict: - """Declare *tool_names* with a schema permissive enough to accept any arguments. +def _llm_params(*tool_names: str, parameters: dict | None = None) -> dict: + """Declare *tool_names* with a schema permissive enough to accept any arguments, or with *parameters*. @tool_output_validation blocks a call whose tool isn't declared here, so per-tool tests that expect the regex check itself to run (not the schema gate) need their tool declared. """ - return { - "tools": [ - { - "type": "function", - "function": {"name": name, "parameters": {"type": "object", "additionalProperties": True}}, - } - for name in tool_names - ] - } + schema = parameters or {"type": "object", "additionalProperties": True} + return {"tools": [{"type": "function", "function": {"name": name, "parameters": schema}} for name in tool_names]} def _build_manager( @@ -140,6 +147,101 @@ async def test_matching_tool_and_pattern_blocks(self): assert result.records[0].rail_type == "tool_output" assert result.records[0].return_value["detections"] == [r"DROP\s+TABLE"] + @pytest.mark.asyncio + async def test_per_tool_block_reports_the_blocked_call(self): + """A per-tool block gives one ``per_tool_rail`` violation naming the rail and the blocked call's position.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output"]}, regex_detection=RUN_SQL_PATTERN_CONFIG + ) + calls = [_sql_call("SELECT 1", call_id="call_1"), _sql_call("DROP TABLE users", call_id="call_2")] + result = await manager.are_tool_calls_safe(calls, _llm_params("run_sql")) + assert result.tool_violations == ( + call_violation( + "per_tool_rail", + "regex check tool output", + tool_call_id="call_2", + tool_name="run_sql", + index=1, + rail="regex check tool output", + ), + ) + + @pytest.mark.asyncio + async def test_unparsable_call_does_not_hide_a_later_per_tool_block(self): + """A per-tool rail still checks the calls after an unparsable one, reporting each call at its own position.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output"]}, regex_detection=RUN_SQL_PATTERN_CONFIG + ) + messages = tool_call_turn( + wire_tool_call("run_sql", '{"query": "SELECT', "call_1"), + wire_tool_call("run_sql", '{"query": "DROP TABLE users"}', "call_2"), + ) + + result = await manager.are_latest_tool_calls_safe(messages, _llm_params("run_sql")) + + assert [(v.violation_type, v.tool_call_id, v.index, v.rail) for v in result.tool_violations] == [ + (ToolViolationType.MALFORMED_ARGUMENTS, "call_1", 0, None), + (ToolViolationType.PER_TOOL_RAIL, "call_2", 1, "regex check tool output"), + ] + + @pytest.mark.asyncio + async def test_non_string_tool_name_blocks_instead_of_raising(self): + """A call whose name is not a string is a ``malformed_tool_call`` block, not an error from the per-tool lookup.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output"]}, regex_detection=RUN_SQL_PATTERN_CONFIG + ) + messages = tool_call_turn({**wire_tool_call(), "function": {"name": {"x": 1}, "arguments": "{}"}}) + + result = await manager.are_latest_tool_calls_safe(messages, _llm_params("run_sql")) + + assert [(v.violation_type, v.index) for v in result.tool_violations] == [ + (ToolViolationType.MALFORMED_TOOL_CALL, 0) + ] + + @pytest.mark.asyncio + async def test_failed_rail_beside_an_unparsable_call_stays_a_failure(self): + """A rail that raises beside an unparsable call keeps the result a failure, with the failed rail's reason.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output $argument=nope"]}, + regex_detection=RUN_SQL_PATTERN_CONFIG, + ) + messages = tool_call_turn( + wire_tool_call("run_sql", '{"query": "SELECT', "call_1"), + wire_tool_call("run_sql", '{"query": "SELECT 1"}', "call_2"), + ) + + result = await manager.are_latest_tool_calls_safe( + messages, _llm_params("run_sql", parameters=_CLOSED_QUERY_SCHEMA) + ) + + assert (result.failed, result.reason) == (True, "regex check tool output error") + assert [(v.violation_type, v.index) for v in result.tool_violations] == [ + (ToolViolationType.MALFORMED_ARGUMENTS, 0), + (ToolViolationType.RAIL_FAILED, 1), + ] + + @pytest.mark.asyncio + async def test_failed_per_tool_rail_reports_rail_failed(self): + """A per-tool rail that raises reports ``rail_failed`` for the call it was checking, not a policy block.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output $argument=nope"]}, + regex_detection=RUN_SQL_PATTERN_CONFIG, + ) + result = await manager.are_tool_calls_safe( + [_sql_call("SELECT 1")], _llm_params("run_sql", parameters=_CLOSED_QUERY_SCHEMA) + ) + assert result.failed is True + assert result.tool_violations == ( + call_violation( + "rail_failed", + "regex check tool output error", + tool_call_id="call_1", + tool_name="run_sql", + index=0, + rail="regex check tool output", + ), + ) + @pytest.mark.asyncio async def test_matching_tool_non_matching_pattern_allows(self): manager = _build_manager( @@ -208,6 +310,52 @@ async def test_undeclared_tool_blocks_per_tool_regex_check(self): result = await manager.are_tool_calls_safe([_sql_call("SELECT 1")], {}) assert_result_blocked(result) + @pytest.mark.asyncio + @pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"]) + @pytest.mark.parametrize( + ("calls", "llm_params", "expected"), + [ + ( + [_sql_call("SELECT 1")], + {}, + call_violation( + "tool_not_allowed", + "tool call 'run_sql' is not an allowed tool", + tool_call_id="call_1", + tool_name="run_sql", + index=0, + rail="regex check tool output", + ), + ), + ( + [_sql_call("SELECT 1", "call_1"), _call_with_arguments("run_sql", {"query": 5}, "call_2")], + _llm_params("run_sql", parameters=_CLOSED_QUERY_SCHEMA), + call_violation( + "arguments_invalid", + "arguments for tool 'run_sql' do not match its schema: 'type' failed at '/query'", + tool_call_id="call_2", + tool_name="run_sql", + index=1, + argument_path="/query", + schema_keyword="type", + rail="regex check tool output", + ), + ), + ], + ids=["undeclared_tool", "schema_invalid_arguments"], + ) + async def test_per_tool_schema_gate_reports_the_underlying_violation(self, calls, llm_params, expected, parallel): + """A per-tool rail's own allowlist and schema gate reports the underlying violation, naming that rail.""" + manager = _build_manager( + per_tool_call_flows={"run_sql": ["regex check tool output"]}, + regex_detection=RUN_SQL_PATTERN_CONFIG, + tool_output_parallel=parallel, + ) + + result = await manager.are_tool_calls_safe(calls, llm_params) + + assert result.tool_violations == (expected,) + @pytest.mark.asyncio async def test_global_flow_blocks_before_per_tool_runs(self): manager = _build_manager( @@ -240,6 +388,7 @@ def _tool_result_messages(content: str, *, name: str | None = "run_sql") -> list class TestAreToolResultsSafe: @pytest.mark.asyncio async def test_matching_tool_and_pattern_blocks(self): + """A per-tool result block gives one ``per_tool_rail`` violation at the tool message's position.""" manager = _build_manager( per_tool_result_flows={"run_sql": ["regex check tool input"]}, regex_detection=RUN_SQL_RESULT_PATTERN_CONFIG, @@ -248,6 +397,16 @@ async def test_matching_tool_and_pattern_blocks(self): assert result.is_safe is False assert result.records[0].tool_name == "run_sql" assert result.records[0].rail_type == "tool_input" + assert result.tool_violations == ( + result_violation( + "per_tool_rail", + "regex check tool input", + tool_call_id="call_1", + tool_name="run_sql", + index=2, + rail="regex check tool input", + ), + ) @pytest.mark.asyncio async def test_non_matching_pattern_allows(self): @@ -359,6 +518,44 @@ async def test_unresolvable_call_id_blocks_when_a_per_tool_policy_is_enabled(sel messages[-1]["tool_call_id"] = "call_unknown" result = await manager.are_tool_results_safe(messages) assert result.is_safe is False + assert result.tool_violations == ( + result_violation( + "unlinkable_result", + "tool result cannot be linked to exactly one prior call; per-tool policy cannot be verified", + tool_call_id="call_unknown", + index=2, + ), + ) + + @pytest.mark.asyncio + async def test_non_string_call_id_still_gives_a_structured_block(self): + """A result linked by a non-string call id is blocked by its per-tool rail with a violation, not an exception.""" + manager = _build_manager( + per_tool_result_flows={"run_sql": ["regex check tool input"]}, + regex_detection=RUN_SQL_RESULT_PATTERN_CONFIG, + ) + messages = _tool_result_messages("ssn: 123-45-6789") + messages[1]["tool_calls"][0]["id"] = 5 + messages[2]["tool_call_id"] = 5 + + result = await manager.are_tool_results_safe(messages) + + assert [(v.violation_type, v.tool_call_id, v.index) for v in result.tool_violations] == [ + (ToolViolationType.PER_TOOL_RAIL, None, 2) + ] + + @pytest.mark.asyncio + async def test_unresolvable_call_id_block_names_no_rail(self): + """The manager refuses an unresolvable result before any per-tool rail runs, so the block names no rail.""" + manager = _build_manager( + per_tool_result_flows={"run_sql": ["regex check tool input"]}, + regex_detection=RUN_SQL_RESULT_PATTERN_CONFIG, + ) + messages = _tool_result_messages("no sensitive data", name=None) + messages[-1]["tool_call_id"] = "call_unknown" + result = await manager.are_tool_results_safe(messages) + assert result.is_safe is False + assert result.triggered_rail is None @pytest.mark.asyncio async def test_global_only_config_does_not_require_result_linkage(self): diff --git a/tests/guardrails/test_rails_manager.py b/tests/guardrails/test_rails_manager.py index 13850f7c0f..b880804ea9 100644 --- a/tests/guardrails/test_rails_manager.py +++ b/tests/guardrails/test_rails_manager.py @@ -19,7 +19,7 @@ import json import logging import warnings -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from opentelemetry.sdk.trace import TracerProvider @@ -27,6 +27,8 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.guardrails.actions.tool_call_action import ToolCallRailAction +from nemoguardrails.guardrails.actions.tool_result_action import ToolResultRailAction from nemoguardrails.guardrails.compiled_rail import RailCompilationError from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.guardrails_types import RailCallRecord, RailDirection, RailResult, serialize_prompt @@ -49,6 +51,7 @@ from nemoguardrails.logging.explain import LLMCallInfo from nemoguardrails.manifests import RailDirection as SurfaceDirection from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.tracing.constants import GuardrailsAttributes from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction from tests.guardrails.async_helpers import ( @@ -75,9 +78,13 @@ from tests.guardrails.tool_helpers import ( WEATHER_SCHEMA, assert_result_blocked, + call_violation, make_tool_conversation, malformed_prior_tool_call_messages, multi_turn_reused_call_id_messages, + result_violation, + tool_call_turn, + wire_tool_call, ) @@ -829,15 +836,27 @@ async def test_no_flows_returns_safe_without_parsing(self): @pytest.mark.asyncio async def test_allows_valid_call(self): + """A declared call with valid arguments is allowed and carries no tool violations.""" mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) result = await mgr.are_tool_calls_safe([_call("get_weather", {"city": "Paris"})], {"tools": [WEATHER_TOOL]}) assert result.is_safe is True + assert result.tool_violations == () @pytest.mark.asyncio async def test_blocks_undeclared_call(self): + """An undeclared call blocks, and the validator's violation is lifted onto the result.""" mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) result = await mgr.are_tool_calls_safe([_call("rm_rf", {})], {"tools": [WEATHER_TOOL]}) assert_result_blocked(result, "rm_rf") + assert result.tool_violations == ( + call_violation( + "tool_not_allowed", + "tool call 'rm_rf' is not an allowed tool", + tool_call_id="c1", + tool_name="rm_rf", + index=0, + ), + ) @pytest.mark.asyncio async def test_no_tool_calls_returns_safe(self): @@ -863,13 +882,25 @@ async def test_a_short_circuit_leaves_no_unawaited_coroutine(self): @pytest.mark.asyncio async def test_fails_closed_on_duplicate_tool_definitions(self): - # parse_tools raises ValueError on a duplicate tool name; the method must - # convert that into a block, not propagate. + """A toolset that declares a tool twice blocks with an ``invalid_toolset`` violation rather than raising.""" mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) result = await mgr.are_tool_calls_safe( [_call("get_weather", {"city": "Paris"})], {"tools": [WEATHER_TOOL, WEATHER_TOOL]} ) assert_result_blocked(result, "tool parsing failed") + assert result.tool_violations == ( + call_violation("invalid_toolset", "tool parsing failed: duplicate tool 'get_weather' in toolset"), + ) + + @pytest.mark.asyncio + async def test_unexpected_parsing_error_blocks_with_a_fixed_reason(self): + """A tool-parsing error other than a duplicate tool blocks with a fixed reason carrying none of its text.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + mgr.engine_registry.parse_tools = MagicMock(side_effect=RuntimeError(_SECRET_VALUE)) + result = await mgr.are_tool_calls_safe([_call("get_weather", {"city": "Paris"})], {"tools": [WEATHER_TOOL]}) + + assert result.tool_violations == (call_violation("invalid_toolset", "tool parsing failed"),) @pytest.mark.asyncio async def test_disabled_toggle_skips_validation(self): @@ -903,15 +934,26 @@ async def test_no_tool_results_returns_safe(self): @pytest.mark.asyncio async def test_allows_linked_result(self): + """A result linked to a prior call is allowed and carries no tool violations.""" mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) result = await mgr.are_tool_results_safe(make_tool_conversation(result_call_id="call_1")) assert result.is_safe is True + assert result.tool_violations == () @pytest.mark.asyncio async def test_blocks_unlinked_result(self): + """An unlinked result blocks, and its violation points at the tool message's position.""" mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) result = await mgr.are_tool_results_safe(make_tool_conversation(result_call_id="call_999")) assert_result_blocked(result, "call_999") + assert result.tool_violations == ( + result_violation( + "unknown_call_id", + "tool result for call_id 'call_999' does not correspond to a prior tool call", + tool_call_id="call_999", + index=2, + ), + ) @pytest.mark.asyncio async def test_recycled_call_ids_across_turns_are_safe(self): @@ -941,16 +983,238 @@ async def test_disabled_toggle_skips_validation(self): @pytest.mark.asyncio async def test_fails_closed_when_exchange_extraction_raises(self): - # If the engine adapter's exchange extraction itself blows up, the method must - # fail closed (block) rather than let the error escape. + """An exchange extraction error blocks with a fixed reason and an ``extraction_failed`` violation.""" mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) - def _boom(*args, **kwargs): - raise RuntimeError("extract boom") - - mgr.engine_registry.extract_tool_exchanges = _boom + mgr.engine_registry.extract_tool_exchanges = MagicMock(side_effect=RuntimeError(_SECRET_VALUE)) result = await mgr.are_tool_results_safe(make_tool_conversation()) assert_result_blocked(result, "tool exchange extraction failed") + assert result.tool_violations == (result_violation("extraction_failed", "tool exchange extraction failed"),) + + +_SECRET_VALUE = "SECRET-VALUE" + + +class TestRailsManagerLatestToolCalls: + """``are_latest_tool_calls_safe``: strictly extract the last assistant message's calls, then validate them.""" + + @pytest.mark.asyncio + async def test_delegates_the_extracted_calls_to_are_tool_calls_safe(self): + """The latest calls, the request's tool parameters and the toggle are handed to ``are_tool_calls_safe``.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + verdict = RailResult.allow() + mgr.are_tool_calls_safe = AsyncMock(return_value=verdict) + llm_params = {"tools": [WEATHER_TOOL]} + + result = await mgr.are_latest_tool_calls_safe( + make_tool_conversation(), llm_params, enabled=["tool call validation"] + ) + + expected_calls = [ + ToolCall(id="call_1", function=ToolCallFunction(name="get_weather", arguments={"city": "Paris"})) + ] + mgr.are_tool_calls_safe.assert_awaited_once_with( + expected_calls, llm_params, enabled=["tool call validation"], model_type="main" + ) + assert result is verdict + + @pytest.mark.asyncio + async def test_blocks_a_call_the_toolset_does_not_declare(self): + """A latest call missing from the declared tools is blocked by the global tool-call rail.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + result = await mgr.are_latest_tool_calls_safe(make_tool_conversation(), {"tools": []}) + + assert_result_blocked(result, "get_weather", "not an allowed tool") + assert result.triggered_rail == "tool call validation" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("messages", "violation"), + [ + ( + tool_call_turn(wire_tool_call(arguments='{"city": "SECRET-VALUE"')), + call_violation( + "malformed_arguments", + "tool call extraction failed: tool call 'call_1' has malformed arguments", + tool_call_id="call_1", + tool_name="get_weather", + index=0, + ), + ), + ( + [{"role": "assistant", "function_call": {"name": "get_weather", "arguments": "{}"}}], + call_violation( + "legacy_function_call", + "tool call extraction failed: legacy function_call is not validated by tool rails", + tool_name="get_weather", + ), + ), + ], + ids=["malformed_arguments", "legacy_function_call"], + ) + async def test_extraction_error_blocks_with_its_violation(self, messages, violation): + """A latest turn the extractor refuses blocks before any rail runs, with the extractor's content-free reason.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + result = await mgr.are_latest_tool_calls_safe(messages, {"tools": [WEATHER_TOOL]}) + + assert (result.is_safe, result.reason, result.triggered_rail) == (False, violation.reason, None) + assert result.tool_violations == (violation,) + + @pytest.mark.asyncio + async def test_unparsable_calls_do_not_hide_the_other_calls_violations(self): + """The calls that parse are still validated, and every violation is reported in call order.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + messages = tool_call_turn( + wire_tool_call("delete_files", "{}", "call_1"), + wire_tool_call(arguments='{"city": "SECRET-VALUE"', call_id="call_2"), + wire_tool_call(call_id="call_3"), + wire_tool_call("rm_rf", "{}", "call_4"), + {"id": "call_5", "type": "function", "name": "get_weather", "arguments": "{}"}, + ) + + result = await mgr.are_latest_tool_calls_safe(messages, {"tools": [WEATHER_TOOL]}) + + not_allowed = "tool call '{}' is not an allowed tool" + assert result.tool_violations == ( + call_violation( + "tool_not_allowed", + not_allowed.format("delete_files"), + tool_call_id="call_1", + tool_name="delete_files", + index=0, + ), + call_violation( + "malformed_arguments", + "tool call extraction failed: tool call 'call_2' has malformed arguments", + tool_call_id="call_2", + tool_name="get_weather", + index=1, + ), + call_violation( + "tool_not_allowed", not_allowed.format("rm_rf"), tool_call_id="call_4", tool_name="rm_rf", index=3 + ), + call_violation( + "malformed_tool_call", + "tool call extraction failed: tool call 'call_5' does not carry a function object", + tool_call_id="call_5", + index=4, + ), + ) + assert (result.is_safe, result.reason, result.triggered_rail) == ( + False, + not_allowed.format("delete_files"), + None, + ) + + @pytest.mark.asyncio + async def test_unparsable_calls_log_one_warning_per_check(self, caplog): + """However many calls cannot be parsed, a check logs one WARNING, so a large request cannot flood the log.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + messages = tool_call_turn(*(wire_tool_call(arguments="not json", call_id=f"call_{n}") for n in range(5))) + + with caplog.at_level(logging.WARNING, logger="nemoguardrails.guardrails.rails_manager"): + await mgr.are_latest_tool_calls_safe(messages, {"tools": [WEATHER_TOOL]}) + + warnings = [record.getMessage() for record in caplog.records if record.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "5 tool calls" in warnings[0] + + @pytest.mark.asyncio + async def test_whole_turn_violation_leads_an_unparsable_calls_violation(self): + """A toolset violation, which no call position owns, is listed and reported ahead of an unparsable call's.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + messages = tool_call_turn(wire_tool_call(arguments="not json"), wire_tool_call(call_id="call_2")) + + result = await mgr.are_latest_tool_calls_safe(messages, {"tools": [WEATHER_TOOL, WEATHER_TOOL]}) + + assert [(v.violation_type, v.index) for v in result.tool_violations] == [ + (ToolViolationType.INVALID_TOOLSET, None), + (ToolViolationType.MALFORMED_ARGUMENTS, 0), + ] + assert result.reason == "tool parsing failed: duplicate tool 'get_weather' in toolset" + + @pytest.mark.asyncio + async def test_unexpected_extraction_error_blocks_with_a_fixed_reason(self): + """Any other extraction error blocks with a fixed reason that carries none of the exception's text.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + mgr.engine_registry.extract_latest_tool_calls = MagicMock(side_effect=RuntimeError(_SECRET_VALUE)) + + result = await mgr.are_latest_tool_calls_safe(make_tool_conversation(), {"tools": [WEATHER_TOOL]}) + + assert result.is_safe is False + assert result.reason == "tool call extraction failed" + assert result.triggered_rail is None + assert result.tool_violations == (call_violation("extraction_failed", "tool call extraction failed"),) + + @pytest.mark.asyncio + async def test_no_tool_calls_on_the_last_assistant_message_is_safe(self): + """A final assistant text turn has no calls to validate, so even an empty toolset allows it.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + messages = [*make_tool_conversation(), {"role": "assistant", "content": "It's 18C in Paris."}] + + result = await mgr.are_latest_tool_calls_safe(messages, {"tools": []}) + + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_no_tool_call_rails_returns_safe_without_extracting(self): + """With no tool-call rail to run nothing is extracted, so a malformed call and a missing engine both pass.""" + mgr = _tool_rails_manager() + + result = await mgr.are_latest_tool_calls_safe( + tool_call_turn(wire_tool_call(arguments='{"city": "SECRET-VALUE"')), {"tools": [WEATHER_TOOL]} + ) + + assert result.is_safe is True + + +class TestRailsManagerToolViolations: + """Violations the manager attaches for a global tool validator that breaks, or that blocks in several turns.""" + + @pytest.mark.asyncio + async def test_a_failing_tool_call_validator_reports_rail_failed(self): + """A global tool-call validator that raises gives one ``rail_failed`` violation naming the rail.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + with patch.object(ToolCallRailAction, "_validate", side_effect=RuntimeError(_SECRET_VALUE)): + result = await mgr.are_tool_calls_safe([_call("get_weather", {"city": "Paris"})], {"tools": [WEATHER_TOOL]}) + + assert result.failed is True + assert result.tool_violations == ( + call_violation("rail_failed", "tool call validation error", rail="tool call validation"), + ) + + @pytest.mark.asyncio + async def test_a_failing_tool_result_validator_reports_rail_failed(self): + """A global tool-result validator that raises gives one ``rail_failed`` violation naming the rail.""" + mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) + + with patch.object(ToolResultRailAction, "_validate", side_effect=RuntimeError(_SECRET_VALUE)): + result = await mgr.are_tool_results_safe(make_tool_conversation()) + + assert result.failed is True + assert result.tool_violations == ( + result_violation("rail_failed", "tool result validation error", rail="tool result validation"), + ) + + @pytest.mark.asyncio + async def test_bad_results_in_every_turn_are_reported(self): + """The result rail checks every turn, not just up to the first bad one, and the reason is the first's.""" + mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) + messages = multi_turn_reused_call_id_messages() + messages[2] = {**messages[2], "tool_call_id": "call_x"} + messages[6] = {**messages[6], "tool_call_id": "call_y"} + + result = await mgr.are_tool_results_safe(messages) + + assert [(v.violation_type, v.tool_call_id, v.index) for v in result.tool_violations] == [ + (ToolViolationType.UNKNOWN_CALL_ID, "call_x", 2), + (ToolViolationType.UNKNOWN_CALL_ID, "call_y", 6), + ] + assert result.reason == "tool result for call_id 'call_x' does not correspond to a prior tool call" class TestRailsManagerToolToggleNormalization: @@ -1157,7 +1421,7 @@ def _mutual_wait_stub(started: dict): resolve if both calls are in flight at once; under a sequential regression the second coroutine is never scheduled, so the wait times out and raises.""" - async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None): + async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None, index=None): name = tool_call.function.name started[name].set() other = "list_tables" if name == "run_sql" else "run_sql" @@ -1235,7 +1499,7 @@ async def test_in_flight_checks_are_capped(self): slots_full = asyncio.Event() release = asyncio.Event() - async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None): + async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None, index=None): nonlocal in_flight, peak in_flight += 1 peak = max(peak, in_flight) @@ -1267,7 +1531,7 @@ async def test_checks_still_waiting_for_a_slot_are_closed_on_a_block(self, recwa slots_full = asyncio.Event() never_set = asyncio.Event() - async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None): + async def stub(direction, flow, tool_call, tool_result=None, tool_definition=None, index=None): # call_0 blocks only once every slot is taken, so the remaining checks are still # waiting for one; the others hold their slots until the block cancels them. nonlocal in_flight @@ -1364,7 +1628,7 @@ def _manager_returning(results: dict[str, RailResult], parallel: bool = False) - ) created: list[str] = [] - def stub(direction, flow, tool_call, tool_result=None, tool_definition=None): + def stub(direction, flow, tool_call, tool_result=None, tool_definition=None, index=None): created.append(tool_call.function.name) return _coroutine_returning(results[tool_call.function.name]) @@ -1398,13 +1662,14 @@ class TestTriggeredRail: @pytest.mark.asyncio async def test_input_block_sets_triggered_rail(self, content_safety_rails_manager): - """An input-rail block records the flow's base name in triggered_rail.""" + """An input-rail block records the flow's base name in triggered_rail, and carries no tool violations.""" mock_rail_model( content_safety_rails_manager.engine_registry, AsyncMock(return_value=LLMResponse(content=UNSAFE_INPUT_JSON)) ) result = await content_safety_rails_manager.is_input_safe(MESSAGES) assert not result.is_safe assert result.triggered_rail == "content safety check input" + assert result.tool_violations == () @pytest.mark.asyncio async def test_output_block_sets_triggered_rail(self, content_safety_rails_manager): @@ -1428,6 +1693,74 @@ async def test_safe_result_has_no_triggered_rail(self, content_safety_rails_mana assert result.triggered_rail is None +class TestToolRailTriggeredRail: + """A blocking global tool rail records its base flow name; a block the manager raises itself names none.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "configured_flow", + ["tool call validation", "tool call validation $model=main", "tool call validation(main)"], + ) + async def test_tool_call_block_sets_triggered_rail(self, configured_flow): + """A global tool-call block names its rail by the base flow name, with any suffix stripped.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=[configured_flow]) + + result = await mgr.are_tool_calls_safe([_call("rm_rf", {})], {"tools": [WEATHER_TOOL]}) + + assert result.is_safe is False + assert result.triggered_rail == "tool call validation" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "configured_flow", + ["tool result validation", "tool result validation $model=main", "tool result validation(main)"], + ) + async def test_tool_result_block_sets_triggered_rail(self, configured_flow): + """A global tool-result block names its rail by the base flow name, with any suffix stripped.""" + mgr = _tool_rails_manager_with_main(tool_result_flows=[configured_flow]) + + result = await mgr.are_tool_results_safe(make_tool_conversation(result_call_id="call_999")) + + assert result.is_safe is False + assert result.triggered_rail == "tool result validation" + + @pytest.mark.asyncio + async def test_tool_call_block_keeps_the_validator_reason(self): + """Naming the rail leaves the validator's own reason as the reason.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + result = await mgr.are_tool_calls_safe([_call("rm_rf", {})], {"tools": [WEATHER_TOOL]}) + + assert result.reason == "tool call 'rm_rf' is not an allowed tool" + + @pytest.mark.asyncio + async def test_tool_parsing_failure_names_no_rail(self): + """A toolset that fails to parse is refused before any rail runs, so the block names no rail.""" + mgr = _tool_rails_manager_with_main(tool_call_flows=["tool call validation"]) + + result = await mgr.are_tool_calls_safe( + [_call("get_weather", {"city": "Paris"})], {"tools": [WEATHER_TOOL, WEATHER_TOOL]} + ) + + assert result.is_safe is False + assert result.triggered_rail is None + + @pytest.mark.asyncio + async def test_exchange_extraction_failure_names_no_rail(self): + """A conversation that fails exchange extraction is refused before any rail runs, so the block names no rail.""" + mgr = _tool_rails_manager_with_main(tool_result_flows=["tool result validation"]) + + def _boom(*args, **kwargs): + raise RuntimeError("extract boom") + + mgr.engine_registry.extract_tool_exchanges = _boom + + result = await mgr.are_tool_results_safe(make_tool_conversation()) + + assert result.is_safe is False + assert result.triggered_rail is None + + class TestOutcomeToResult: """`_rail_result` maps a rail's verdict onto IORails' result type.""" diff --git a/tests/guardrails/test_tool_call_action.py b/tests/guardrails/test_tool_call_action.py index e59bf64ffe..ac5931fcf0 100644 --- a/tests/guardrails/test_tool_call_action.py +++ b/tests/guardrails/test_tool_call_action.py @@ -15,20 +15,23 @@ """Unit tests for ToolCallRailAction (allowlist + argument-schema validation).""" +from typing import Any + import pytest from nemoguardrails.guardrails.actions.tool_call_action import ToolCallRailAction from nemoguardrails.guardrails.tool_schema import Tool, Toolset +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.types import ToolCall, ToolCallFunction -from tests.guardrails.tool_helpers import WEATHER_SCHEMA, assert_outcome_blocked +from tests.guardrails.tool_helpers import WEATHER_SCHEMA, assert_outcome_blocked, call_violation, violations_in def _toolset() -> Toolset: return Toolset(tools=[Tool(name="get_weather", arguments_schema=WEATHER_SCHEMA), Tool(name="ping")]) -def _call(name: str, arguments: dict) -> ToolCall: - return ToolCall(id="c1", function=ToolCallFunction(name=name, arguments=arguments)) +def _call(name: str, arguments: dict, call_id: str = "c1") -> ToolCall: + return ToolCall(id=call_id, function=ToolCallFunction(name=name, arguments=arguments)) class TestToolCallRailAction: @@ -42,6 +45,15 @@ async def test_no_parameter_function_tool_with_arguments_is_blocked(self): """A function tool that declares no parameters accepts no arguments, so a call supplying any is blocked.""" result = await ToolCallRailAction().run(_toolset(), [_call("ping", {"anything": 1})]) assert_outcome_blocked(result, "ping", "no arguments") + assert violations_in(result) == [ + call_violation( + "unexpected_arguments", + "tool 'ping' accepts no arguments but the call supplied 1 argument", + tool_call_id="c1", + tool_name="ping", + index=0, + ) + ] @pytest.mark.asyncio async def test_no_parameter_function_tool_without_arguments_is_safe(self): @@ -51,13 +63,55 @@ async def test_no_parameter_function_tool_without_arguments_is_safe(self): @pytest.mark.asyncio async def test_undeclared_tool_is_blocked(self): + """A call to an undeclared tool blocks with a ``tool_not_allowed`` violation naming the call.""" result = await ToolCallRailAction().run(_toolset(), [_call("rm_rf", {})]) assert_outcome_blocked(result, "rm_rf", "not an allowed tool") + assert violations_in(result) == [ + call_violation( + "tool_not_allowed", + "tool call 'rm_rf' is not an allowed tool", + tool_call_id="c1", + tool_name="rm_rf", + index=0, + ) + ] + + @pytest.mark.asyncio + async def test_undeclared_tool_name_is_capped(self): + """An undeclared name, which the model made up, is cut to 64 characters in the reason and the violation.""" + capped = "x" * 64 + "..." + + result = await ToolCallRailAction().run(_toolset(), [_call("x" * 200, {})]) + + assert result.reason == f"tool call '{capped}' is not an allowed tool" + assert [violation.tool_name for violation in violations_in(result)] == [capped] @pytest.mark.asyncio async def test_invalid_arguments_are_blocked(self): + """Schema-invalid arguments block with an ``arguments_invalid`` violation pointing at the failing argument.""" result = await ToolCallRailAction().run(_toolset(), [_call("get_weather", {})]) assert_outcome_blocked(result, "get_weather") + assert violations_in(result) == [ + call_violation( + "arguments_invalid", + "arguments for tool 'get_weather' do not match its schema: 'required' failed at '/city'", + tool_call_id="c1", + tool_name="get_weather", + index=0, + argument_path="/city", + schema_keyword="required", + ) + ] + + @pytest.mark.asyncio + async def test_undeclared_call_with_a_non_string_id_is_still_reported(self): + """A call whose id is not a string blocks as ``tool_not_allowed`` with no id, rather than breaking the rail.""" + non_string_id: Any = 5 + call = ToolCall(id=non_string_id, function=ToolCallFunction(name="rm_rf", arguments={})) + result = await ToolCallRailAction().run(_toolset(), [call]) + assert violations_in(result) == [ + call_violation("tool_not_allowed", "tool call 'rm_rf' is not an allowed tool", tool_name="rm_rf", index=0) + ] @pytest.mark.asyncio async def test_empty_tool_calls_is_safe(self): @@ -65,10 +119,19 @@ async def test_empty_tool_calls_is_safe(self): assert result.is_blocked is False @pytest.mark.asyncio - async def test_one_bad_call_blocks_the_batch(self): - calls = [_call("get_weather", {"city": "Paris"}), _call("rm_rf", {})] + async def test_every_bad_call_is_reported_in_order(self): + """Each bad call in the list gets its own violation, in list order, and the reason is the first one's.""" + calls = [ + _call("rm_rf", {}, call_id="c1"), + _call("get_weather", {"city": "Paris"}, call_id="c2"), + _call("ping", {"anything": 1}, call_id="c3"), + ] result = await ToolCallRailAction().run(_toolset(), calls) - assert_outcome_blocked(result, "rm_rf") + assert [(v.violation_type, v.index, v.tool_call_id) for v in violations_in(result)] == [ + (ToolViolationType.TOOL_NOT_ALLOWED, 0, "c1"), + (ToolViolationType.UNEXPECTED_ARGUMENTS, 2, "c3"), + ] + assert result.reason == "tool call 'rm_rf' is not an allowed tool" @pytest.mark.asyncio async def test_hosted_tool_matched_by_type_when_name_is_empty(self): diff --git a/tests/guardrails/test_tool_rails_iorails.py b/tests/guardrails/test_tool_rails_iorails.py index 5a11f0897e..1633156819 100644 --- a/tests/guardrails/test_tool_rails_iorails.py +++ b/tests/guardrails/test_tool_rails_iorails.py @@ -494,6 +494,29 @@ async def test_undeclared_streamed_tool_call_blocked(self, iorails): # The tool-call chunk is suppressed: no chunk carries the tool call. assert not any(isinstance(c, str) and '"tool_calls"' in c for c in chunks) + @pytest.mark.asyncio + async def test_undeclared_streamed_tool_call_violation_carries_the_validator_reason(self, iorails): + """The violation message carries the validator's reason, not the name of the rail that blocked.""" + _inject_sse_stream(iorails, _tool_call_sse_lines("rm_rf", ["{}"])) + chunks = await _collect(iorails.stream_async(MESSAGES, options={"llm_params": LLM_PARAMS})) + + violations = _stream_violation_chunks(chunks) + assert violations[0]["error"]["message"] == ( + "Blocked by tool output rails: tool call 'rm_rf' is not an allowed tool" + ) + + @pytest.mark.asyncio + async def test_streamed_schema_mismatch_violation_omits_the_argument_value(self, iorails): + """A streamed schema-mismatch violation names the failing argument but never quotes its value.""" + _inject_sse_stream(iorails, _tool_call_sse_lines("get_weather", ['{"city": ["SECRET-VALUE"]}'])) + chunks = await _collect(iorails.stream_async(MESSAGES, options={"llm_params": LLM_PARAMS})) + + violations = _stream_violation_chunks(chunks) + assert violations[0]["error"]["message"] == ( + "Blocked by tool output rails: " + "arguments for tool 'get_weather' do not match its schema: 'type' failed at '/city'" + ) + @pytest.mark.asyncio async def test_truncated_streamed_tool_call_fails_closed(self, iorails): # Truncated tool-call args must fail closed (parity with non-streaming): no tool-call @@ -519,6 +542,17 @@ async def test_unlinked_tool_result_blocks_stream(self, iorails): assert REFUSAL_MESSAGE not in chunks forbidden_post.assert_not_called() + @pytest.mark.asyncio + async def test_unlinked_tool_result_violation_carries_the_validator_reason(self, iorails): + """The violation message carries the validator's reason, not the name of the rail that blocked.""" + _inject_forbidden_transport(iorails) + chunks = await _collect(iorails.stream_async(make_tool_conversation(result_call_id="call_999"))) + + violations = _stream_violation_chunks(chunks) + assert violations[0]["error"]["message"] == ( + "Blocked by tool input rails: tool result for call_id 'call_999' does not correspond to a prior tool call" + ) + @pytest.mark.asyncio async def test_nameless_linked_tool_result_streams_the_model_reply(self, iorails): """A tool result without a name, linked by its call id, passes the tool-result rail and the reply streams.""" diff --git a/tests/guardrails/test_tool_result_action.py b/tests/guardrails/test_tool_result_action.py index dd549bf445..c6fed85f32 100644 --- a/tests/guardrails/test_tool_result_action.py +++ b/tests/guardrails/test_tool_result_action.py @@ -15,12 +15,18 @@ """Unit tests for ToolResultRailAction (call_id linkage + structural validation).""" +from typing import Any + import pytest from nemoguardrails.guardrails.actions.tool_result_action import ToolResultRailAction from nemoguardrails.guardrails.tool_schema import ToolResult +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.types import ToolCall, ToolCallFunction -from tests.guardrails.tool_helpers import assert_outcome_blocked +from tests.guardrails.tool_helpers import result_violation, violations_in + +# Content of a type ToolResult does not declare, as a malformed client payload carries. +_MALFORMED_CONTENT: Any = {"unexpected": "shape"} def _prior_calls() -> list: @@ -30,198 +36,192 @@ def _prior_calls() -> list: ] -def _result(call_id, name=None, content: "str | list[dict] | None" = "18C") -> ToolResult: - return ToolResult(call_id=call_id, name=name, content=content) +def _prior_call(call_id: str, name: str) -> ToolCall: + return ToolCall(id=call_id, function=ToolCallFunction(name=name, arguments={})) + + +def _result( + call_id, name=None, content: "str | list[dict] | None" = "18C", message_index: "int | None" = None +) -> ToolResult: + return ToolResult(call_id=call_id, name=name, content=content, message_index=message_index) + + +# Content types ToolResult does not declare, as a malformed client payload carries them. +_NON_DICT_BLOCKS: Any = [1, 2, 3] + +# A call id longer than a reason quotes: the reason keeps 64 characters, the violation the whole id. +_LONG_ID = "i" * 200 +_CAPPED_ID = "i" * 64 + "..." + +# Case id -> (results, prior calls) the validator accepts. +_ACCEPTED = { + "linked_result_with_matching_name": ([_result("c1", name="get_weather")], _prior_calls()), + "result_without_name_linked_by_call_id": ([_result("c2")], _prior_calls()), + "prior_call_without_a_name": ([_result("c1", name="get_weather")], [_prior_call("c1", "")]), + "list_of_content_blocks": ( + [_result("c1", name="get_weather", content=[{"type": "text", "text": "18C"}])], + _prior_calls(), + ), + "no_content": ([_result("c1", name="get_weather", content=None)], _prior_calls()), + "empty_content_list": ([_result("c1", name="get_weather", content=[])], _prior_calls()), + "no_results": ([], _prior_calls()), + "prior_call_with_empty_id": ([], [_prior_call("", "get_weather")]), +} + +# Case id -> (results, prior calls, the one violation the validator reports). +_REJECTED = { + "missing_call_id": ( + [_result("", message_index=3)], + _prior_calls(), + result_violation("missing_call_id", "tool result is missing a call_id", index=3), + ), + "unknown_call_id": ( + [_result("c9", message_index=3)], + _prior_calls(), + result_violation( + "unknown_call_id", + "tool result for call_id 'c9' does not correspond to a prior tool call", + tool_call_id="c9", + index=3, + ), + ), + "name_mismatch": ( + [_result("c1", name="search", message_index=3)], + _prior_calls(), + result_violation( + "name_mismatch", + "tool result name 'search' does not match the called tool 'get_weather' for call_id 'c1'", + tool_call_id="c1", + tool_name="get_weather", + index=3, + ), + ), + "malformed_content": ( + [_result("c1", name="get_weather", content=_MALFORMED_CONTENT, message_index=3)], + _prior_calls(), + result_violation( + "malformed_content", + "tool result for call_id 'c1' has malformed content", + tool_call_id="c1", + tool_name="get_weather", + index=3, + ), + ), + "list_of_non_dicts": ( + [_result("c1", name="get_weather", content=_NON_DICT_BLOCKS)], + _prior_calls(), + result_violation( + "malformed_content", + "tool result for call_id 'c1' has malformed content", + tool_call_id="c1", + tool_name="get_weather", + ), + ), + "duplicate_prior_call_id": ( + [_result("c1")], + [_prior_call("c1", "get_weather"), _prior_call("c1", "search")], + result_violation( + "duplicate_prior_call_id", + "duplicate prior tool call id 'c1' makes tool-result linkage ambiguous", + tool_call_id="c1", + ), + ), + "duplicate_result": ( + [_result("c1", name="get_weather", message_index=2), _result("c1", name="get_weather", message_index=3)], + _prior_calls(), + result_violation( + "duplicate_result", + "duplicate tool result for call_id 'c1': each tool call must have exactly one result", + tool_call_id="c1", + tool_name="get_weather", + index=3, + ), + ), + "long_unknown_call_id_capped_in_the_reason": ( + [_result(_LONG_ID, message_index=3)], + _prior_calls(), + result_violation( + "unknown_call_id", + f"tool result for call_id '{_CAPPED_ID}' does not correspond to a prior tool call", + tool_call_id=_LONG_ID, + index=3, + ), + ), + "long_names_capped_in_the_reason": ( + [_result("c1", name="r" * 200, message_index=3)], + [_prior_call("c1", "n" * 200)], + result_violation( + "name_mismatch", + f"tool result name '{'r' * 64}...' does not match the called tool '{'n' * 64}...' for call_id 'c1'", + tool_call_id="c1", + tool_name="n" * 64 + "...", + index=3, + ), + ), +} class TestToolResultRailAction: @pytest.mark.asyncio - async def test_linked_result_with_matching_name_is_safe(self): - result = await ToolResultRailAction().run([_result("c1", name="get_weather")], _prior_calls()) - assert result.is_blocked is False - - @pytest.mark.asyncio - async def test_result_without_name_is_identified_by_its_call_id(self): - """A result without a name is safe when its call_id links it to a prior call.""" - result = await ToolResultRailAction().run([_result("c2")], _prior_calls()) - assert result.is_blocked is False - - @pytest.mark.asyncio - async def test_list_content_is_well_formed(self): - result = await ToolResultRailAction().run( - [_result("c1", name="get_weather", content=[{"type": "text", "text": "18C"}])], _prior_calls() - ) - assert result.is_blocked is False - - @pytest.mark.asyncio - async def test_empty_results_is_safe(self): - result = await ToolResultRailAction().run([], _prior_calls()) - assert result.is_blocked is False + @pytest.mark.parametrize(("results", "prior"), list(_ACCEPTED.values()), ids=list(_ACCEPTED)) + async def test_well_formed_linked_results_are_safe(self, results, prior): + """Results that link to one prior call, don't contradict its name and carry well-formed content pass.""" + outcome = await ToolResultRailAction().run(results, prior) + assert outcome.is_blocked is False @pytest.mark.asyncio - async def test_missing_call_id_is_blocked(self): - result = await ToolResultRailAction().run([_result("")], _prior_calls()) - assert_outcome_blocked(result, "missing a call_id") + @pytest.mark.parametrize(("results", "prior", "violation"), list(_REJECTED.values()), ids=list(_REJECTED)) + async def test_a_bad_result_is_blocked_with_its_violation(self, results, prior, violation): + """A result that breaks one rule blocks, reporting that rule's violation at the result's message index.""" + outcome = await ToolResultRailAction().run(results, prior) + assert outcome.is_blocked + assert outcome.reason == violation.reason + assert violations_in(outcome) == [violation] @pytest.mark.asyncio - async def test_unlinked_call_id_is_blocked(self): - """A result without a name whose call_id matches no prior call is blocked.""" - result = await ToolResultRailAction().run([_result("c9")], _prior_calls()) - assert_outcome_blocked(result, "c9", "does not correspond to a prior tool call") - - @pytest.mark.asyncio - async def test_name_mismatch_is_blocked(self): - result = await ToolResultRailAction().run([_result("c1", name="search")], _prior_calls()) - assert_outcome_blocked(result, "does not match the called tool", "get_weather") - - @pytest.mark.asyncio - async def test_malformed_content_is_blocked(self): - result = await ToolResultRailAction().run( - [_result("c1", name="get_weather", content={"unexpected": "shape"})], # type: ignore[arg-type] - _prior_calls(), - ) - assert_outcome_blocked(result, "malformed content") - - @pytest.mark.asyncio - async def test_list_of_non_dicts_is_blocked(self): - result = await ToolResultRailAction().run( - [_result("c1", name="get_weather", content=[1, 2, 3])], # type: ignore[arg-type] - _prior_calls(), - ) - assert_outcome_blocked(result, "malformed content") + async def test_every_bad_result_is_reported_in_order(self): + """Each bad result gets its own violation, in list order, and the reason is the first one's.""" + results = [ + _result("c9", message_index=2), + _result("c1", name="get_weather", message_index=3), + _result("", message_index=4), + ] + result = await ToolResultRailAction().run(results, _prior_calls()) + assert [(v.violation_type, v.index) for v in violations_in(result)] == [ + (ToolViolationType.UNKNOWN_CALL_ID, 2), + (ToolViolationType.MISSING_CALL_ID, 4), + ] + assert result.reason == "tool result for call_id 'c9' does not correspond to a prior tool call" @pytest.mark.asyncio - async def test_empty_content_list_is_well_formed(self): - result = await ToolResultRailAction().run([_result("c1", name="get_weather", content=[])], _prior_calls()) - assert result.is_blocked is False + async def test_a_result_reports_only_its_first_failing_check(self): + """A result that fails both the name and the content check reports only ``name_mismatch``.""" + bad = _result("c1", name="search", content=_MALFORMED_CONTENT, message_index=3) + result = await ToolResultRailAction().run([bad], _prior_calls()) + assert [v.violation_type for v in violations_in(result)] == [ToolViolationType.NAME_MISMATCH] @pytest.mark.asyncio - async def test_one_bad_result_blocks_the_batch(self): - results = [_result("c1", name="get_weather"), _result("c9")] + async def test_results_without_call_ids_are_each_missing_not_duplicates(self): + """Two results with no call_id each report ``missing_call_id``; neither counts as a duplicate.""" + results = [_result(None, message_index=2), _result(None, message_index=3)] result = await ToolResultRailAction().run(results, _prior_calls()) - assert_outcome_blocked(result, "c9") + assert [(v.violation_type, v.index) for v in violations_in(result)] == [ + (ToolViolationType.MISSING_CALL_ID, 2), + (ToolViolationType.MISSING_CALL_ID, 3), + ] @pytest.mark.asyncio - async def test_duplicate_prior_call_id_is_blocked(self): + async def test_duplicate_prior_call_ids_are_each_reported_and_skip_result_checks(self): + """Every duplicated prior id is reported once, and no result is checked against the ambiguous calls.""" prior = [ - ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})), - ToolCall(id="c1", function=ToolCallFunction(name="search", arguments={})), + _prior_call("c1", "get_weather"), + _prior_call("c1", "search"), + _prior_call("c2", "get_weather"), + _prior_call("c2", "search"), + _prior_call("c2", "list_files"), ] - result = await ToolResultRailAction().run([_result("c1")], prior) - assert_outcome_blocked(result, "duplicate prior tool call id", "c1") - - @pytest.mark.asyncio - async def test_prior_call_with_empty_id_is_skipped(self): - prior = [ToolCall(id="", function=ToolCallFunction(name="get_weather", arguments={}))] - result = await ToolResultRailAction().run([], prior) - assert result.is_blocked is False - - @pytest.mark.asyncio - async def test_duplicate_result_ids_are_blocked(self): - results = [_result("c1", name="get_weather"), _result("c1", name="get_weather")] - result = await ToolResultRailAction().run(results, _prior_calls()) - assert_outcome_blocked(result, "duplicate tool result", "c1") - - -class TestValidateToolResultIds: - def setup_method(self): - self.action = ToolResultRailAction() - - def test_unique_ids_returns_none(self): - assert self.action._validate_tool_result_ids([_result("c1"), _result("c2")]) is None - - def test_empty_list_returns_none(self): - assert self.action._validate_tool_result_ids([]) is None - - def test_result_without_call_id_is_skipped(self): - assert self.action._validate_tool_result_ids([ToolResult(call_id=None)]) is None - - def test_duplicate_call_id_is_blocked(self): - results = [ - ToolResult(call_id="call_123", name="search", content="real result"), - ToolResult(call_id="call_123", name="search", content="duplicate/injected result"), + result = await ToolResultRailAction().run([_result("c9", message_index=6)], prior) + assert [(v.violation_type, v.tool_call_id) for v in violations_in(result)] == [ + (ToolViolationType.DUPLICATE_PRIOR_CALL_ID, "c1"), + (ToolViolationType.DUPLICATE_PRIOR_CALL_ID, "c2"), ] - assert_outcome_blocked(self.action._validate_tool_result_ids(results), "duplicate tool result", "call_123") - - -class TestValidateResultCallId: - def setup_method(self): - self.action = ToolResultRailAction() - self.calls_by_id = { - "c1": ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})), - } - - def test_linked_result_returns_none(self): - assert self.action._validate_result_call_id(_result("c1"), self.calls_by_id) is None - - def test_missing_call_id_is_blocked(self): - assert_outcome_blocked( - self.action._validate_result_call_id(_result(""), self.calls_by_id), - "missing a call_id", - ) - - def test_orphaned_call_id_is_blocked(self): - assert_outcome_blocked( - self.action._validate_result_call_id(_result("c9"), self.calls_by_id), - "c9", - "does not correspond to a prior tool call", - ) - - -class TestValidateResultName: - def setup_method(self): - self.action = ToolResultRailAction() - self.prior = ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})) - - def test_matching_names_returns_none(self): - assert self.action._validate_result_name(_result("c1", name="get_weather"), self.prior) is None - - def test_result_without_name_returns_none(self): - """A result without a name defers to its linked call's name, so there is nothing to contradict.""" - assert self.action._validate_result_name(_result("c1"), self.prior) is None - - def test_prior_without_function_name_returns_none(self): - """If a prior call is missing a name, a matching result with a name is allowed through""" - prior = ToolCall(id="c1", function=ToolCallFunction(name="", arguments={})) - assert self.action._validate_result_name(_result("c1", name="get_weather"), prior) is None - - def test_both_names_absent_returns_none(self): - prior = ToolCall(id="c1", function=ToolCallFunction(name="", arguments={})) - assert self.action._validate_result_name(_result("c1"), prior) is None - - def test_name_mismatch_is_blocked(self): - assert_outcome_blocked( - self.action._validate_result_name(_result("c1", name="search"), self.prior), - "search", - "get_weather", - "does not match", - ) - - -class TestValidateResultContent: - def setup_method(self): - self.action = ToolResultRailAction() - - def test_string_content_returns_none(self): - assert self.action._validate_result_content(_result("c1", content="18C")) is None - - def test_list_of_dicts_returns_none(self): - assert self.action._validate_result_content(_result("c1", content=[{"type": "text", "text": "18C"}])) is None - - def test_none_content_returns_none(self): - assert self.action._validate_result_content(_result("c1", content=None)) is None - - def test_empty_list_returns_none(self): - assert self.action._validate_result_content(_result("c1", content=[])) is None - - def test_dict_content_is_blocked(self): - assert_outcome_blocked( - self.action._validate_result_content(_result("c1", content={"key": "val"})), # type: ignore[arg-type] - "malformed content", - ) - - def test_list_of_non_dicts_is_blocked(self): - assert_outcome_blocked( - self.action._validate_result_content(_result("c1", content=[1, 2, 3])), # type: ignore[arg-type] - "malformed content", - ) diff --git a/tests/guardrails/test_tool_schema.py b/tests/guardrails/test_tool_schema.py index 5ebf83c571..d276dc65fd 100644 --- a/tests/guardrails/test_tool_schema.py +++ b/tests/guardrails/test_tool_schema.py @@ -15,9 +15,13 @@ """Unit tests for the tool_schema module (canonical tool types + arg validation).""" +import uuid from dataclasses import FrozenInstanceError +from typing import Any +from unittest.mock import patch import pytest +from jsonschema import Draft202012Validator from nemoguardrails.actions.rail_outcome import RailOutcome from nemoguardrails.guardrails.tool_schema import ( @@ -28,8 +32,11 @@ tool_output_validation, validate_arguments, ) +from nemoguardrails.rails.llm.options import ToolViolationType from nemoguardrails.types import ToolCall, ToolCallFunction +_SECRET_VALUE = "SECRET-VALUE" + _WEATHER_SCHEMA = { "type": "object", "properties": { @@ -139,14 +146,26 @@ def test_valid_with_only_required(self): assert validate_arguments(_weather_tool(), {"city": "Paris"}) is None def test_type_mismatch_is_rejected(self): - reason = validate_arguments(_weather_tool(), {"city": 123}) - assert reason is not None - assert "get_weather" in reason + """A type mismatch names the failing argument and keyword, never the value that failed.""" + violation = validate_arguments(_weather_tool(), {"city": "Paris", "days": _SECRET_VALUE}) + assert (violation.violation_type, violation.argument_path, violation.schema_keyword) == ( + ToolViolationType.ARGUMENTS_INVALID, + "/days", + "type", + ) + assert violation.reason == "arguments for tool 'get_weather' do not match its schema: 'type' failed at '/days'" def test_missing_required_is_rejected(self): - reason = validate_arguments(_weather_tool(), {}) - assert reason is not None - assert "get_weather" in reason + """A missing required argument is named in the path, although jsonschema reports it at the parent object.""" + violation = validate_arguments(_weather_tool(), {}) + assert (violation.violation_type, violation.argument_path, violation.schema_keyword) == ( + ToolViolationType.ARGUMENTS_INVALID, + "/city", + "required", + ) + assert violation.reason == ( + "arguments for tool 'get_weather' do not match its schema: 'required' failed at '/city'" + ) def test_hosted_tool_without_schema_skips_validation(self): """A hosted/server tool (name is None) declares no schema and the provider owns the call shape, so it stays allowlist-only: any arguments are accepted here.""" @@ -161,34 +180,31 @@ def test_function_tool_without_parameters_allows_empty_arguments(self): def test_function_tool_without_parameters_rejects_arguments(self): """A function tool that declares no parameters accepts no arguments, so supplying any blocks instead of skipping.""" tool = Tool(name="get_time", arguments_schema=None) - reason = validate_arguments(tool, {"path": "/etc/shadow"}) - assert reason is not None - assert "get_time" in reason - assert "no arguments" in reason + violation = validate_arguments(tool, {"path": "/etc/shadow"}) + assert (violation.violation_type, violation.argument_path, violation.schema_keyword) == ( + ToolViolationType.UNEXPECTED_ARGUMENTS, + None, + None, + ) + assert violation.reason == "tool 'get_time' accepts no arguments but the call supplied 1 argument" def test_function_tool_empty_dict_schema_rejects_arguments(self): """An explicit empty schema ({}) declares no properties; jsonschema would accept anything, so it is treated as no-args.""" tool = Tool(name="ping", arguments_schema={}) - reason = validate_arguments(tool, {"x": 1}) - assert reason is not None - assert "ping" in reason - assert "no arguments" in reason + violation = validate_arguments(tool, {"x": 1}) + assert violation.violation_type is ToolViolationType.UNEXPECTED_ARGUMENTS def test_function_tool_empty_properties_rejects_arguments(self): """An object schema with empty properties and no additionalProperties:false defaults to permissive, so it is treated as no-args.""" tool = Tool(name="ping", arguments_schema={"type": "object", "properties": {}}) - reason = validate_arguments(tool, {"x": 1}) - assert reason is not None - assert "ping" in reason - assert "no arguments" in reason + violation = validate_arguments(tool, {"x": 1}) + assert violation.violation_type is ToolViolationType.UNEXPECTED_ARGUMENTS def test_function_tool_object_schema_without_properties_rejects_arguments(self): """An object schema with no properties key at all declares no inputs, so it is treated as no-args.""" tool = Tool(name="ping", arguments_schema={"type": "object"}) - reason = validate_arguments(tool, {"x": 1}) - assert reason is not None - assert "ping" in reason - assert "no arguments" in reason + violation = validate_arguments(tool, {"x": 1}) + assert violation.violation_type is ToolViolationType.UNEXPECTED_ARGUMENTS def test_function_tool_additional_properties_schema_accepts_arguments(self): """A schema that opens additionalProperties declares an input channel, so it is not treated as no-args and free-form arguments pass.""" @@ -206,13 +222,114 @@ def test_function_tool_composition_schema_accepts_arguments(self): assert validate_arguments(tool, {"x": 1}) is None def test_malformed_schema_is_reported_not_raised(self): - # A declared schema that is not valid JSON Schema degrades to a reason - # rather than raising (e.g. a non-JSON-Schema dialect reaching the - # validator before its engine adapter normalizes it). + """A declared schema that is not valid JSON Schema gives ``invalid_tool_schema`` rather than raising.""" bad = Tool(name="bad", arguments_schema={"type": "not-a-real-type"}) - reason = validate_arguments(bad, {}) - assert reason is not None - assert "bad" in reason + violation = validate_arguments(bad, {}) + assert violation.violation_type is ToolViolationType.INVALID_TOOL_SCHEMA + assert violation.reason == "declared schema for tool 'bad' is not valid JSON Schema" + + def test_schema_that_is_not_json_is_reported_as_invalid(self): + """A schema that cannot be written as JSON, here one holding a set, gives ``invalid_tool_schema``.""" + odd = Tool(name="odd", arguments_schema={"type": "object", "default": {"a", "b"}}) + violation = validate_arguments(odd, {}) + assert violation.violation_type is ToolViolationType.INVALID_TOOL_SCHEMA + assert violation.reason == "declared schema for tool 'odd' is not valid JSON Schema" + + def test_a_schema_is_checked_once_however_many_calls_use_it(self): + """Validating several calls against one schema checks the schema against its metaschema only once.""" + # A schema no other test uses, so no earlier validation has cached it. + tool = Tool(name="get_weather", arguments_schema={**_WEATHER_SCHEMA, "$comment": str(uuid.uuid4())}) + check_schema = Draft202012Validator.check_schema + + with patch.object(Draft202012Validator, "check_schema", wraps=check_schema) as schema_checks: + violations = [validate_arguments(tool, {"city": city}) for city in ("Paris", "Rome", "Oslo")] + + assert (violations, schema_checks.call_count) == ([None, None, None], 1) + + +_CLOSED_WEATHER_SCHEMA = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "additionalProperties": False, +} + + +def _object(properties: dict, **keywords: Any) -> dict: + """An object schema declaring *properties*, plus any other keywords.""" + return {"type": "object", "properties": properties, **keywords} + + +_DRAFT3 = {"$schema": "http://json-schema.org/draft-03/schema#"} + +# Case id -> (arguments schema, arguments, (argument_path, schema_keyword)). +_ARGUMENT_PATHS = { + "unexpected_key_not_named": ( + _CLOSED_WEATHER_SCHEMA, + {"city": "Paris", _SECRET_VALUE: 1}, + ("", "additionalProperties"), + ), + "map_key_masked": ( + _object({"contacts": {"type": "object", "additionalProperties": {"type": "integer"}}}), + {"contacts": {_SECRET_VALUE: "x"}}, + ("/contacts/*", "type"), + ), + "array_index_kept": ( + _object({"cities": {"type": "array", "items": {"type": "string"}}}), + {"cities": ["Paris", 3]}, + ("/cities/1", "type"), + ), + "nested_declared_names": ( + _object({"opts": _object({"units": {"type": "string"}})}), + {"opts": {"units": 5}}, + ("/opts/units", "type"), + ), + "missing_dependency_named": ( + _object({"a": {}, "b": {}}, dependentRequired={"a": ["b"]}), + {"a": 1}, + ("/b", "dependentRequired"), + ), + "segments_escaped": (_object({"a/b~c": {"type": "string"}}), {"a/b~c": 1}, ("/a~1b~0c", "type")), + "property_name_failure_at_parent": ( + _object({"a": {}}, propertyNames={"maxLength": 3}), + {_SECRET_VALUE: 1}, + ("", "maxLength"), + ), + "unevaluated_property_at_parent": ( + _object({"city": {"type": "string"}}, unevaluatedProperties=False), + {"city": "Paris", "debug": 1}, + ("", "unevaluatedProperties"), + ), + "draft3_required_at_property": ( + _object({"a": {"type": "string", "required": True}}, **_DRAFT3), + {}, + ("/a", "required"), + ), +} + + +class TestArgumentPath: + """Where ``validate_arguments`` points ``argument_path`` for each kind of schema failure.""" + + @pytest.mark.parametrize( + ("schema", "arguments", "expected"), list(_ARGUMENT_PATHS.values()), ids=list(_ARGUMENT_PATHS) + ) + def test_argument_path_names_only_what_the_schema_declares(self, schema, arguments, expected): + """The path names only declared properties and indices, and the reason quotes no argument value or key.""" + violation = validate_arguments(Tool(name="tool", arguments_schema=schema), arguments) + assert (violation.violation_type, violation.argument_path, violation.schema_keyword) == ( + ToolViolationType.ARGUMENTS_INVALID, + *expected, + ) + assert _SECRET_VALUE not in violation.reason + + def test_top_level_failure_reads_as_the_top_level(self): + """A failure at the root keeps an empty ``argument_path``, and its reason says "the top level".""" + schema = _object({"a": {}, "b": {}}, minProperties=2) + violation = validate_arguments(Tool(name="pair", arguments_schema=schema), {"a": 1}) + assert violation.argument_path == "" + assert violation.reason == ( + "arguments for tool 'pair' do not match its schema: 'minProperties' failed at the top level" + ) def _weather_call(arguments: dict) -> ToolCall: @@ -243,6 +360,15 @@ class TestToolOutputValidation: async def _action(tool_call, tool_definition, argument_name=None, **kwargs): return RailOutcome.allow() + @pytest.mark.asyncio + async def test_undeclared_tool_name_is_capped_in_the_reason(self): + """An undeclared name, which the model made up, is cut to 64 characters in the block reason.""" + call = ToolCall(id="call_1", type="function", function=ToolCallFunction(name="x" * 200, arguments={})) + + outcome = await self._action(tool_call=call, tool_definition=None) + + assert outcome.reason == f"tool call '{'x' * 64}...' is not an allowed tool" + @pytest.mark.asyncio async def test_no_argument_name_skips_the_check(self): outcome = await self._action(tool_call=_weather_call({"city": "Paris"}), tool_definition=_weather_tool()) diff --git a/tests/guardrails/tool_helpers.py b/tests/guardrails/tool_helpers.py index f59b05b410..efb2671aad 100644 --- a/tests/guardrails/tool_helpers.py +++ b/tests/guardrails/tool_helpers.py @@ -21,7 +21,9 @@ carry explicit messages since this module is not assertion-rewritten by pytest. """ -from typing import Optional +from typing import Any, Optional + +from nemoguardrails.rails.llm.options import ToolViolation, ToolViolationType WEATHER_SCHEMA = { "type": "object", @@ -54,6 +56,43 @@ def assert_result_blocked(result, *substrings: str) -> None: _assert_reason_contains(result.reason, substrings, result) +def call_violation(violation_type: str, reason: str, **identity: Any) -> ToolViolation: + """A tool-call ``ToolViolation``; *violation_type* is its wire value, such as ``"tool_not_allowed"``.""" + return ToolViolation(kind="tool_call", violation_type=ToolViolationType(violation_type), reason=reason, **identity) + + +def result_violation(violation_type: str, reason: str, **identity: Any) -> ToolViolation: + """A tool-result ``ToolViolation``; *violation_type* is its wire value, such as ``"unknown_call_id"``.""" + return ToolViolation( + kind="tool_result", violation_type=ToolViolationType(violation_type), reason=reason, **identity + ) + + +TOOL_CALL_QUESTION = "What's the weather in Paris?" + +UNREAD_TOOLS_MESSAGE = "tools is read only by a tool_call check; include tool_call in rail_types or leave tools out" + + +def wire_tool_call(name: str = "get_weather", arguments: Any = '{"city": "Paris"}', call_id: str = "call_1") -> dict: + """One OpenAI Chat Completions tool call as it appears on an assistant message.""" + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + + +def assistant_tool_calls(*calls: Any, content: Optional[str] = None) -> dict: + """An assistant message carrying *calls*, with *content* as its text.""" + return {"role": "assistant", "content": content, "tool_calls": list(calls)} + + +def tool_call_turn(*calls: Any, content: Optional[str] = None) -> list: + """A user question, then an assistant turn carrying *calls*.""" + return [{"role": "user", "content": TOOL_CALL_QUESTION}, assistant_tool_calls(*calls, content=content)] + + +def violations_in(outcome) -> list[ToolViolation]: + """The ``ToolViolation``s a tool validator attached to its ``RailOutcome`` metadata.""" + return [ToolViolation.model_validate(violation) for violation in outcome.metadata.get("tool_violations", [])] + + def make_tool_conversation(result_call_id: str = "call_1", result_name: Optional[str] = "get_weather") -> list: """A user turn, an assistant ``get_weather`` tool call (id ``call_1``), then a tool result. diff --git a/tests/server/test_guardrail_checks.py b/tests/server/test_guardrail_checks.py index c5fa995f1a..8f5f8f736a 100644 --- a/tests/server/test_guardrail_checks.py +++ b/tests/server/test_guardrail_checks.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -20,12 +21,17 @@ pytest.importorskip("openai", reason="openai is required for server tests") from fastapi.testclient import TestClient -from nemoguardrails.exceptions import RailTypeNotConfiguredError +from nemoguardrails.exceptions import ( + InvalidCheckRequestError, + RailTypeNotConfiguredError, + RailTypeNotSupportedError, +) from nemoguardrails.rails import LLMRails from nemoguardrails.rails.llm.config import RailsConfig -from nemoguardrails.rails.llm.options import RailsResult, RailStatus, RailType +from nemoguardrails.rails.llm.options import RailsResult, RailStatus, RailType, ToolViolation from nemoguardrails.server import api from nemoguardrails.testing.fake_model import FakeLLMModel +from tests.guardrails.tool_helpers import call_violation, result_violation client = TestClient(api.app) @@ -52,6 +58,16 @@ def _post(body: dict, **kwargs): return client.post(ENDPOINT, json=body, **kwargs) +def _check_body(guardrails: Optional[dict] = None, **fields: Any) -> dict: + """A check body for one user turn against config ``test``, with extra ``guardrails`` options and top-level fields.""" + return { + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": {"config_id": "test", **(guardrails or {})}, + **fields, + } + + def _checked(result, config_id="test", colang_version="1.0"): with patch.object(api, "_get_rails", new_callable=AsyncMock, return_value=_mock_rails(result, colang_version)): resp = _post( @@ -136,9 +152,57 @@ def test_reason_returned_on_blocked(): ids=["passed", "blocked-without-reason"], ) def test_reason_absent_when_none(result): - """A result with no reason, as on a pass or an LLMRails block, omits reason from the response.""" + """A result with no reason or tool violations, as on a pass or an LLMRails block, omits both from the response.""" data = _checked(result) assert "reason" not in data + assert "tool_violations" not in data + + +_UNLINKED_RESULT = result_violation( + "unknown_call_id", + "tool result for call_id 'call_9' does not correspond to a prior tool call", + tool_call_id="call_9", + index=3, +) +_UNDECLARED_CALL = call_violation( + "tool_not_allowed", "tool call 'rm_rf' is not an allowed tool", tool_call_id="call_1", tool_name="rm_rf", index=0 +) + + +def _tool_block(*violations: ToolViolation) -> RailsResult: + """A blocked tool check carrying *violations*.""" + return RailsResult( + status=RailStatus.BLOCKED, + content="I'm sorry, I can't respond to that.", + rail="tool result validation", + reason=violations[0].reason, + tool_violations=list(violations), + ) + + +def test_tool_violations_returned_on_blocked(): + """A blocked tool check's violations reach the response, without their unset fields.""" + data = _checked(_tool_block(_UNDECLARED_CALL)) + assert data["tool_violations"] == [ + { + "kind": "tool_call", + "violation_type": "tool_not_allowed", + "reason": "tool call 'rm_rf' is not an allowed tool", + "tool_call_id": "call_1", + "tool_name": "rm_rf", + "index": 0, + } + ] + + +def test_result_violation_index_skips_the_prepended_context(): + """With context prepended, a result's index still counts the request's messages; a call's index is unchanged.""" + mock = _mock_rails(_tool_block(_UNLINKED_RESULT, _UNDECLARED_CALL)) + + with patch.object(api, "_get_rails", new_callable=AsyncMock, return_value=mock): + resp = _post(_check_body({"context": {"topic": "weather"}})) + + assert [violation["index"] for violation in resp.json()["tool_violations"]] == [2, 0] # --- Config resolution --- @@ -370,6 +434,8 @@ def test_context_prepended_to_messages(): (["input"], [RailType.INPUT]), (["output"], [RailType.OUTPUT]), (["input", "output"], [RailType.INPUT, RailType.OUTPUT]), + (["tool_call"], [RailType.TOOL_CALL]), + (["tool_result"], [RailType.TOOL_RESULT]), (None, None), ], ) @@ -409,21 +475,46 @@ def test_rail_types_invalid_value_returns_422(): # --- Unsatisfiable rail_types --- -def test_unsatisfiable_rail_types_returns_422(): - result = RailsResult(status=RailStatus.PASSED, content="hi") - mock = _mock_rails(result) - mock.check_async = AsyncMock( - side_effect=RailTypeNotConfiguredError("Requested rail type 'output' has no configured rails.") - ) +@pytest.mark.parametrize( + ("error", "rail_types"), + [ + (RailTypeNotConfiguredError("Requested rail type 'output' has no configured rails."), ["output"]), + (RailTypeNotSupportedError("LLMRails supports input and output rails only, not tool_call"), ["tool_call"]), + (InvalidCheckRequestError("tools is read only by a tool_call check"), ["input"]), + ], + ids=["not_configured", "not_supported", "invalid_check_request"], +) +def test_check_request_errors_return_422(error, rail_types): + """An unconfigured or unsupported rail type, or contradictory check arguments, give a 422 carrying the message.""" + mock = _mock_rails(RailsResult(status=RailStatus.PASSED, content="hi")) + mock.check_async = AsyncMock(side_effect=error) with patch.object(api, "_get_rails", new_callable=AsyncMock, return_value=mock): - resp = _post( - { - "model": "test", - "messages": [{"role": "user", "content": "hi"}], - "guardrails": {"config_id": "test", "rail_types": ["output"]}, - } - ) + resp = _post(_check_body({"rail_types": rail_types})) + + assert resp.status_code == 422 + assert resp.json()["error"]["message"] == str(error) + + +_WEATHER_TOOL = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} + + +@pytest.mark.parametrize("tools", [None, [_WEATHER_TOOL]], ids=["absent", "present"]) +def test_tools_passed_through(tools): + """The request's tools reach the engine's check_async unchanged; the engine decides whether it can use them.""" + mock = _mock_rails(RailsResult(status=RailStatus.PASSED, content="hi")) + body = _check_body() + if tools is not None: + body["tools"] = tools + + with patch.object(api, "_get_rails", new_callable=AsyncMock, return_value=mock): + resp = _post(body) + + assert resp.status_code == 200 + assert mock.check_async.call_args.kwargs["tools"] == tools + +def test_custom_tools_are_rejected(): + """The request's tools are still validated as chat-completion tools, so a custom tool is a 422.""" + resp = _post(_check_body({"rail_types": ["tool_call"]}, tools=[{"type": "custom", "custom": {"name": "grep"}}])) assert resp.status_code == 422 - assert "output" in resp.json()["error"]["message"] diff --git a/tests/server/test_iorails_engine_compat.py b/tests/server/test_iorails_engine_compat.py index eb5c9d4b98..4d05b6c293 100644 --- a/tests/server/test_iorails_engine_compat.py +++ b/tests/server/test_iorails_engine_compat.py @@ -34,12 +34,19 @@ from nemoguardrails.guardrails import guardrails as guardrails_module from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails from nemoguardrails.guardrails.model_engine import ModelEngine +from nemoguardrails.rails.llm.llmrails import LLMRails from nemoguardrails.rails.llm.options import GenerationResponse from nemoguardrails.server import api +from nemoguardrails.testing.fake_model import FakeLLMModel from nemoguardrails.types import LLMResponse from tests.guardrails.test_data import CONTENT_SAFETY_CONFIG -from tests.guardrails.test_tool_rails_iorails import TOOL_CONFIG -from tests.guardrails.tool_helpers import make_tool_conversation +from tests.guardrails.test_tool_rails_iorails import TOOL_CONFIG, WEATHER_TOOL +from tests.guardrails.tool_helpers import ( + UNREAD_TOOLS_MESSAGE, + make_tool_conversation, + tool_call_turn, + wire_tool_call, +) REASONING_TRACE = "The user asked for a capital city." LLM_ANSWER = "Paris." @@ -376,3 +383,150 @@ def test_chat_completion_refuses_an_unlinked_tool_result_under_iorails_alias(too assert response.status_code == 200 assert response.json()["choices"][0]["message"]["content"] == REFUSAL_MESSAGE stubbed_main_model.assert_not_awaited() + + +def _post_check( + messages: list, rail_types: "list | None", *, tools: "list | None" = None, context: "dict | None" = None +): + """POST *messages* to /v1/checks against the tool-rails config.""" + guardrails = {"config_id": "tools", "rail_types": rail_types} + if context is not None: + guardrails["context"] = context + body = {"model": "test-model", "messages": messages, "guardrails": guardrails} + if tools is not None: + body["tools"] = tools + # Each request of an un-entered TestClient runs on a fresh event loop, which an IORails + # engine cached by an earlier request cannot serve, so each request builds its own. + api.llm_rails_instances.clear() + client = TestClient(api.app, raise_server_exceptions=False) + return client.post("/v1/checks", json=body) + + +def test_check_blocks_a_disallowed_tool_call_under_iorails_alias(tool_rails_alias, stubbed_main_model): + """/v1/checks refuses an undeclared call with the rail, the reason and a structured violation, without the model.""" + response = _post_check(tool_call_turn(wire_tool_call("delete_files", "{}")), ["tool_call"], tools=[WEATHER_TOOL]) + + assert response.status_code == 200 + assert response.json() == { + "status": "blocked", + "content": REFUSAL_MESSAGE, + "rail": "tool call validation", + "reason": "tool call 'delete_files' is not an allowed tool", + "tool_violations": [ + { + "kind": "tool_call", + "violation_type": "tool_not_allowed", + "reason": "tool call 'delete_files' is not an allowed tool", + "tool_call_id": "call_1", + "tool_name": "delete_files", + "index": 0, + } + ], + } + stubbed_main_model.assert_not_awaited() + + +def test_check_honours_the_request_tools_under_iorails_alias(tool_rails_alias, stubbed_main_model): + """The request's top-level tools are the allowlist: the same call is refused without them and passes with them.""" + without_tools = _post_check(tool_call_turn(wire_tool_call()), ["tool_call"]) + with_tools = _post_check(tool_call_turn(wire_tool_call()), ["tool_call"], tools=[WEATHER_TOOL]) + + assert (without_tools.json()["status"], with_tools.json()["status"]) == ("blocked", "passed") + stubbed_main_model.assert_not_awaited() + + +def test_check_round_trip_under_iorails_alias(tool_rails_alias, stubbed_main_model): + """A harness checks its model's tool call, then the spec-shaped tool message it appends, and both pass.""" + messages = tool_call_turn(wire_tool_call()) + + call_check = _post_check(messages, ["tool_call"], tools=[WEATHER_TOOL]) + messages.append({"role": "tool", "tool_call_id": "call_1", "content": "18C"}) + result_check = _post_check(messages, ["tool_result"]) + + assert (call_check.json()["status"], result_check.json()["status"]) == ("passed", "passed") + assert "tool_violations" not in result_check.json() + stubbed_main_model.assert_not_awaited() + + +@pytest.mark.parametrize( + ("rail_type", "tools"), [("tool_call", [WEATHER_TOOL]), ("tool_result", None)], ids=["tool_call", "tool_result"] +) +def test_tool_checks_accept_guardrails_context_under_iorails_alias( + tool_rails_alias, stubbed_main_model, rail_type, tools +): + """The context message the server prepends does not break either tool check.""" + messages = [*tool_call_turn(wire_tool_call()), {"role": "tool", "tool_call_id": "call_1", "content": "18C"}] + + response = _post_check(messages, [rail_type], tools=tools, context={"user_id": "u1"}) + + assert response.status_code == 200 + assert response.json()["status"] == "passed" + + +def test_result_violation_index_counts_request_messages_under_iorails_alias(tool_rails_alias, stubbed_main_model): + """With context prepended, an unlinked result's index is still its position in the request's messages.""" + messages = [*tool_call_turn(wire_tool_call()), {"role": "tool", "tool_call_id": "call_999", "content": "18C"}] + + response = _post_check(messages, ["tool_result"], context={"user_id": "u1"}) + + assert [violation["index"] for violation in response.json()["tool_violations"]] == [2] + + +@pytest.mark.parametrize("rail_types", [None, ["tool_result"]], ids=["auto_detection", "tool_result"]) +def test_check_with_unread_tools_returns_422_under_iorails_alias(tool_rails_alias, stubbed_main_model, rail_types): + """IORails refuses tools on a check that runs no tool_call rails with a 422, rather than ignoring them.""" + response = _post_check(make_tool_conversation(result_name=None), rail_types, tools=[WEATHER_TOOL]) + + assert response.status_code == 422 + assert response.json()["error"]["message"] == UNREAD_TOOLS_MESSAGE + stubbed_main_model.assert_not_awaited() + + +# A Colang input rail, which LLMRails runs through its own runtime: the config the server holds when +# NEMO_GUARDRAILS_IORAILS_ENGINE is unset. +INPUT_RAIL_COLANG = """ +define flow input rail + if $user_message == "block" + bot refuse to respond + stop +""" +INPUT_RAIL_YAML = """ +rails: + input: + flows: + - input rail +""" + + +@pytest.fixture +def llmrails_server(monkeypatch): + """Serve an input-rail config through a real LLMRails inside the server, with a fake main model.""" + input_rail_config = RailsConfig.from_content(INPUT_RAIL_COLANG, INPUT_RAIL_YAML) + + def build_llmrails(config, verbose=False): + return LLMRails(config, llm=FakeLLMModel(responses=[]), verbose=verbose) + + monkeypatch.setattr(api, "LLMRails", build_llmrails) + monkeypatch.setattr(api.RailsConfig, "from_path", staticmethod(lambda full_path: input_rail_config)) + yield + + +def test_check_tool_rail_types_on_llmrails_return_422(llmrails_server): + """LLMRails serving /v1/checks refuses tool rail types with a 422 naming them, rather than passing or a 500.""" + response = _post_check(tool_call_turn(wire_tool_call()), ["input", "tool_result", "tool_call"]) + + assert response.status_code == 422 + assert response.json()["error"]["message"] == ( + "LLMRails supports input and output rails only, not tool_call, tool_result" + ) + + +def test_check_with_tools_on_llmrails_returns_422(llmrails_server): + """LLMRails serving /v1/checks refuses a request carrying tools, even for an input check it can run.""" + response = _post_check([{"role": "user", "content": "hi"}], ["input"], tools=[WEATHER_TOOL]) + + assert response.status_code == 422 + assert response.json()["error"]["message"] == ( + "LLMRails check() does not run tool rails, so it does not take tools; " + "tool_call checks run on the IORails engine only." + ) diff --git a/tests/test_generation_options.py b/tests/test_generation_options.py index 4a3d021fd7..59b1ff4403 100644 --- a/tests/test_generation_options.py +++ b/tests/test_generation_options.py @@ -14,6 +14,7 @@ # limitations under the License. import json import os.path +from typing import Any import pytest @@ -26,6 +27,8 @@ GenerationLog, GenerationResponse, GenerationStats, + ToolViolation, + ToolViolationType, ) from tests.utils import TestChat @@ -669,3 +672,32 @@ async def test_rails_options_combinations(input_opt, output_opt, dialog_opt, exp assert input_rails_ran == expect_input, f"Input rails: expected {expect_input}, got {rail_names}" assert output_rails_ran == expect_output, f"Output rails: expected {expect_output}, got {rail_names}" + + +def test_tool_violation_stores_non_string_identity_as_none(): + """A non-string call id or tool name is stored as None rather than failing to build the violation.""" + non_string_id: Any = 5 + non_string_name: Any = ["run_sql"] + + violation = ToolViolation( + kind="tool_call", + violation_type=ToolViolationType.TOOL_NOT_ALLOWED, + reason="tool call is not an allowed tool", + tool_call_id=non_string_id, + tool_name=non_string_name, + ) + + assert (violation.tool_call_id, violation.tool_name) == (None, None) + + +def test_tool_violation_caps_the_tool_name_but_keeps_the_call_id_whole(): + """A tool name, which the model can make up, is cut to 64 characters; the call id stays whole as identity.""" + violation = ToolViolation( + kind="tool_call", + violation_type=ToolViolationType.TOOL_NOT_ALLOWED, + reason="tool call is not an allowed tool", + tool_call_id="i" * 200, + tool_name="n" * 200, + ) + + assert (violation.tool_call_id, violation.tool_name) == ("i" * 200, "n" * 64 + "...") diff --git a/tests/test_llmrails_check_async.py b/tests/test_llmrails_check_async.py index 2047b1cb1f..244d635e89 100644 --- a/tests/test_llmrails_check_async.py +++ b/tests/test_llmrails_check_async.py @@ -18,7 +18,7 @@ import pytest from nemoguardrails import LLMRails, RailsConfig -from nemoguardrails.exceptions import RailTypeNotConfiguredError +from nemoguardrails.exceptions import RailTypeNotConfiguredError, RailTypeNotSupportedError from nemoguardrails.rails.llm.llmrails import ( _determine_rails_from_messages, _get_blocking_rail, @@ -27,6 +27,7 @@ _normalize_messages_for_rails, ) from nemoguardrails.rails.llm.options import ActivatedRail, GenerationLog, GenerationResponse, RailStatus, RailType +from tests.utils import check_entry_points, run_check class TestDetermineRailsFromMessages: @@ -587,6 +588,37 @@ async def test_auto_detect_skips_validation(self, input_only_rails): assert result.status == RailStatus.PASSED +_TOOL_RAIL_TYPE_REQUESTS = pytest.mark.parametrize( + ("rail_types", "unsupported"), + [ + ([RailType.TOOL_CALL], "tool_call"), + ([RailType.TOOL_RESULT], "tool_result"), + ([RailType.INPUT, RailType.TOOL_CALL], "tool_call"), + ([RailType.TOOL_RESULT, RailType.TOOL_CALL, RailType.TOOL_CALL], "tool_call, tool_result"), + ], + ids=["tool_call", "tool_result", "input_and_tool_call", "both_tool_types_repeated"], +) + + +class TestToolRailTypesAreNotSupported: + """LLMRails check() runs input and output rails only, so a tool rail type fails loud instead of passing.""" + + @check_entry_points + @_TOOL_RAIL_TYPE_REQUESTS + def test_tool_rail_types_raise(self, mock_rails, entry_point, rail_types, unsupported): + """``check`` and ``check_async`` name each unsupported tool rail type once, in sorted order.""" + with pytest.raises(RailTypeNotSupportedError) as excinfo: + run_check(mock_rails, entry_point, [{"role": "user", "content": "hello"}], rail_types=rail_types) + assert str(excinfo.value) == f"LLMRails supports input and output rails only, not {unsupported}" + + @check_entry_points + @pytest.mark.parametrize("rail_types", [None, [RailType.INPUT]], ids=["auto_detected", "input"]) + def test_tools_raise(self, mock_rails, entry_point, rail_types): + """Both entry points raise ``RailTypeNotSupportedError`` for ``tools``, which only a tool-call check reads.""" + with pytest.raises(RailTypeNotSupportedError, match="tools"): + run_check(mock_rails, entry_point, [{"role": "user", "content": "hello"}], rail_types=rail_types, tools=[]) + + @pytest.fixture def no_main_llm_config(): return RailsConfig.from_content( diff --git a/tests/utils.py b/tests/utils.py index 68a24d6cbc..7b9617ce9a 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -14,10 +14,14 @@ # limitations under the License. from __future__ import annotations +import asyncio +import inspect import json import sys from typing import Any, Dict, Iterable, List, Optional +import pytest + from nemoguardrails import RailsConfig from nemoguardrails.colang import parse_colang_file from nemoguardrails.colang.v2_x.runtime.flows import State @@ -123,3 +127,15 @@ def _init_state(colang_content, yaml_content: Optional[str] = None) -> State: json.dump(state.flow_configs, sys.stdout, indent=4, cls=EnhancedJsonEncoder) return state + + +# Runs a test once through each check entry point; call them with ``run_check``. +check_entry_points = pytest.mark.parametrize("entry_point", ["check", "check_async"]) + + +def run_check(rails: Any, entry_point: str, *args: Any, **kwargs: Any) -> Any: + """Call ``check`` or ``check_async`` from sync test code, so one test body covers both entry points.""" + outcome = getattr(rails, entry_point)(*args, **kwargs) + if inspect.iscoroutine(outcome): + return asyncio.run(outcome) + return outcome