From a88925e97e4cf5095eac4cf59a49aa72a7e48294 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Mon, 24 Aug 2026 14:33:10 -0400 Subject: [PATCH 1/3] refactor(guardrail): migrate guardrail consumers to typed GuardrailClient Migrate guardrail API call sites from sdk.guardrail.* (Stainless SDK) to client_from_platform(sdk, GuardrailClient).* (typed HTTP client), following the pattern established in #1277. AIRCORE-827 Signed-off-by: Max Dubrinsky --- e2e/guardrails/test_checks.py | 19 ++++-- .../tests/test_outputs.py | 51 +++++++++++----- .../tests/test_outputs.py | 51 +++++++++++----- .../tests/test_outputs.py | 59 +++++++++++++------ .../tests/test_outputs.py | 59 +++++++++++++------ 5 files changed, 164 insertions(+), 75 deletions(-) diff --git a/e2e/guardrails/test_checks.py b/e2e/guardrails/test_checks.py index 018c200e5f..eb57cafe4b 100644 --- a/e2e/guardrails/test_checks.py +++ b/e2e/guardrails/test_checks.py @@ -19,6 +19,9 @@ GuardrailCheckResponse, GuardrailsDataParam, ) +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.guardrail.client import GuardrailClient +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest from e2e.guardrails.utils import ( BACKEND_RESPONSE, @@ -51,11 +54,17 @@ def _post_check( if extra_guardrails: guardrails.update(extra_guardrails) - return test_case.sdk.guardrail.check( - workspace=test_case.workspace, - model=test_case.backend_model_ref, - messages=_check_messages(test_case), - guardrails=cast(GuardrailsDataParam, guardrails), + return ( + client_from_platform(test_case.sdk, GuardrailClient) + .check_guardrail( + workspace=test_case.workspace, + body=GuardrailCheckRequest( + model=test_case.backend_model_ref, + messages=_check_messages(test_case), + guardrails=cast(GuardrailsDataParam, guardrails), + ), + ) + .data() ) diff --git a/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py b/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py index a7afa4cba3..d31b4f57da 100644 --- a/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py @@ -16,6 +16,9 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.guardrail.client import GuardrailClient +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, ListGuardrailConfigsQueryParams WORKSPACE = "default" MODEL = "default/mock-llm" @@ -35,14 +38,18 @@ def guardrail_config_id(client: NeMoPlatform) -> str: Pre-existing configs (default, abc, self-check) use meta/llama3-70b-instruct, so only the agent's config will reference mock-llm. """ - response = client.guardrail.configs.list(page=1, page_size=50) + response = ( + client_from_platform(client, GuardrailClient) + .list_guardrail_configs(query_params=ListGuardrailConfigsQueryParams(page=1, page_size=50)) + .data() + ) configs = response.data assert configs, "No guardrail configurations found. The agent should have created one." # Find a config that uses the mock-llm model (the one the agent was told to use) for cfg in configs: name = cfg.name - full_config = client.guardrail.configs.retrieve(name=name) + full_config = client_from_platform(client, GuardrailClient).get_guardrail_config(name=name).data() models = (full_config.data.models or []) if full_config.data else [] for model in models: model_name = getattr(model, "model", "") @@ -70,13 +77,19 @@ def test_content_blocked_by_guardrails(client: NeMoPlatform, guardrail_config_id The mock inference backend always returns 'Yes' to the self-check prompt, causing ALL content to be blocked regardless of whether it's actually harmful. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, + ) + ) + .data() ) assert response.status == "blocked", f"Content should be blocked, but got status: {response.status}" print(f"Content correctly blocked with status: {response.status}") @@ -88,13 +101,19 @@ def test_safe_content_also_blocked(client: NeMoPlatform, guardrail_config_id: st Since the mock always returns 'Yes' (block), even safe content gets blocked. This verifies the guardrails pipeline is working end-to-end. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, + ) + ) + .data() ) assert response.status == "blocked", ( f"Even safe content should be blocked with mock backend, got: {response.status}" diff --git a/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py b/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py index a7afa4cba3..d31b4f57da 100644 --- a/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py @@ -16,6 +16,9 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.guardrail.client import GuardrailClient +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, ListGuardrailConfigsQueryParams WORKSPACE = "default" MODEL = "default/mock-llm" @@ -35,14 +38,18 @@ def guardrail_config_id(client: NeMoPlatform) -> str: Pre-existing configs (default, abc, self-check) use meta/llama3-70b-instruct, so only the agent's config will reference mock-llm. """ - response = client.guardrail.configs.list(page=1, page_size=50) + response = ( + client_from_platform(client, GuardrailClient) + .list_guardrail_configs(query_params=ListGuardrailConfigsQueryParams(page=1, page_size=50)) + .data() + ) configs = response.data assert configs, "No guardrail configurations found. The agent should have created one." # Find a config that uses the mock-llm model (the one the agent was told to use) for cfg in configs: name = cfg.name - full_config = client.guardrail.configs.retrieve(name=name) + full_config = client_from_platform(client, GuardrailClient).get_guardrail_config(name=name).data() models = (full_config.data.models or []) if full_config.data else [] for model in models: model_name = getattr(model, "model", "") @@ -70,13 +77,19 @@ def test_content_blocked_by_guardrails(client: NeMoPlatform, guardrail_config_id The mock inference backend always returns 'Yes' to the self-check prompt, causing ALL content to be blocked regardless of whether it's actually harmful. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, + ) + ) + .data() ) assert response.status == "blocked", f"Content should be blocked, but got status: {response.status}" print(f"Content correctly blocked with status: {response.status}") @@ -88,13 +101,19 @@ def test_safe_content_also_blocked(client: NeMoPlatform, guardrail_config_id: st Since the mock always returns 'Yes' (block), even safe content gets blocked. This verifies the guardrails pipeline is working end-to-end. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, + ) + ) + .data() ) assert response.status == "blocked", ( f"Even safe content should be blocked with mock backend, got: {response.status}" diff --git a/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py b/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py index 8ff781e6b7..3c3126d8cb 100644 --- a/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py @@ -22,6 +22,9 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.guardrail.client import GuardrailClient +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest from trace_reader import get_session WORKSPACE = "default" @@ -52,7 +55,7 @@ def client() -> NeMoPlatform: @pytest.fixture def config(client: NeMoPlatform): """Retrieve the agent-created guardrail config.""" - return client.guardrail.configs.retrieve(name=CONFIG_NAME) + return client_from_platform(client, GuardrailClient).get_guardrail_config(name=CONFIG_NAME).data() # --- Config structure checks --- @@ -147,12 +150,18 @@ def test_input_rail_blocks_fruit_mention(client: NeMoPlatform) -> None: A message about apples should trigger a 'Yes' response from the self-check, causing guardrails to mark the request blocked. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "blocked", f"Message mentioning fruit should be blocked, got: {response.status}" print(f"Input rail correctly blocked fruit mention: {response.status}") @@ -164,12 +173,18 @@ def test_normal_message_passes_through(client: NeMoPlatform) -> None: A geography question doesn't mention fruit (passes input rail) and the response won't be about bread baking (passes output rail). """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "success", f"Normal message should NOT be blocked, got: {response.status}" print(f"Normal message passed through: {response.status}") @@ -182,12 +197,18 @@ def test_output_rail_blocks_bread_content(client: NeMoPlatform) -> None: baking will elicit a response about baking bread, which the output self-check should mark as blocked. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "blocked", f"Response about baking bread should be blocked, got: {response.status}" print(f"Output rail correctly blocked bread content: {response.status}") diff --git a/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py b/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py index 35aa67825f..45b937e4ed 100644 --- a/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py @@ -22,6 +22,9 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.guardrail.client import GuardrailClient +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest from trace_reader import get_session WORKSPACE = "default" @@ -52,7 +55,7 @@ def client() -> NeMoPlatform: @pytest.fixture def config(client: NeMoPlatform): """Retrieve the agent-created guardrail config.""" - return client.guardrail.configs.retrieve(name=CONFIG_NAME) + return client_from_platform(client, GuardrailClient).get_guardrail_config(name=CONFIG_NAME).data() # --- Config structure checks --- @@ -147,12 +150,18 @@ def test_input_rail_blocks_fruit_mention(client: NeMoPlatform) -> None: A message about apples should trigger a 'Yes' response from the self-check, causing guardrails to mark the request blocked. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "blocked", f"Message mentioning fruit should be blocked, got: {response.status}" print(f"Input rail correctly blocked fruit mention: {response.status}") @@ -164,12 +173,18 @@ def test_normal_message_passes_through(client: NeMoPlatform) -> None: A geography question doesn't mention fruit (passes input rail) and the response won't be about bread baking (passes output rail). """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "success", f"Normal message should NOT be blocked, got: {response.status}" print(f"Normal message passed through: {response.status}") @@ -182,12 +197,18 @@ def test_output_rail_blocks_bread_content(client: NeMoPlatform) -> None: baking will elicit a response about baking bread, which the output self-check should mark as blocked. """ - response = client.guardrail.check( - model=MODEL, - messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, + response = ( + client_from_platform(client, GuardrailClient) + .check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, + ) + ) + .data() ) assert response.status == "blocked", f"Response about baking bread should be blocked, got: {response.status}" print(f"Output rail correctly blocked bread content: {response.status}") From 196df8e8e7df7fdc826313f5f6e30f69e8cbe8ad Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Wed, 26 Aug 2026 17:22:35 -0400 Subject: [PATCH 2/3] fix(guardrail): declare check request/response fields on typed client The migrated GuardrailCheckRequest and GuardrailCheckResponse were empty extra="allow" models, so ty rejected every keyword argument at the call sites and the lint-python-types gate failed. Declare the actual HTTP contract instead. Also fix consumers that the ty exclusion for tests/agentic-use hid: list_guardrail_configs returns a paginated response with no data(), and GuardrailConfig.data is a plain dict rather than a typed model, so attribute access silently returned empty strings. Signed-off-by: Max Dubrinsky --- e2e/guardrails/test_checks.py | 22 +-- .../nemo_platform_plugin/guardrail/types.py | 59 +++++++- .../tests/test_outputs.py | 71 ++++------ .../tests/test_outputs.py | 71 ++++------ .../tests/test_outputs.py | 132 +++++++----------- .../tests/test_outputs.py | 132 +++++++----------- 6 files changed, 225 insertions(+), 262 deletions(-) diff --git a/e2e/guardrails/test_checks.py b/e2e/guardrails/test_checks.py index eb57cafe4b..4746962b28 100644 --- a/e2e/guardrails/test_checks.py +++ b/e2e/guardrails/test_checks.py @@ -9,19 +9,13 @@ """ from collections.abc import Callable -from typing import Any, TypeAlias, cast +from typing import Any import nemo_platform import pytest -from nemo_platform.types.guardrail import ( - ChatCompletionAssistantMessageParam, - ChatCompletionUserMessageParam, - GuardrailCheckResponse, - GuardrailsDataParam, -) from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.guardrail.client import GuardrailClient -from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, GuardrailCheckResponse from e2e.guardrails.utils import ( BACKEND_RESPONSE, @@ -30,12 +24,6 @@ GuardrailsChatTestCase, ) -# `guardrails` is assembled dynamically from test fixture data (`content_safety_config()`, -# ad-hoc `extra_guardrails` overrides), so we build it as a plain dict and cast it once at -# the SDK call boundary rather than threading `GuardrailsDataParam`'s nested TypedDicts -# through the test fixtures. -_CheckMessage: TypeAlias = ChatCompletionUserMessageParam | ChatCompletionAssistantMessageParam - def _post_check( test_case: GuardrailsChatTestCase, @@ -61,15 +49,15 @@ def _post_check( body=GuardrailCheckRequest( model=test_case.backend_model_ref, messages=_check_messages(test_case), - guardrails=cast(GuardrailsDataParam, guardrails), + guardrails=guardrails, ), ) .data() ) -def _check_messages(test_case: GuardrailsChatTestCase) -> list[_CheckMessage]: - messages: list[_CheckMessage] = [{"role": "user", "content": test_case.user_input}] +def _check_messages(test_case: GuardrailsChatTestCase) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [{"role": "user", "content": test_case.user_input}] if "output" in test_case.rail_types: messages.append({"role": "assistant", "content": BACKEND_RESPONSE}) return messages diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/guardrail/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/guardrail/types.py index b3cd43e036..6922334bb0 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/guardrail/types.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/guardrail/types.py @@ -40,11 +40,54 @@ class GuardrailConfig(BaseModel): GuardrailConfigPage = Page[GuardrailConfig] +class RailStatus(BaseModel): + """Status of an individual rail.""" + + model_config = ConfigDict(extra="allow") + + status: str = Field(description="Status of the individual rail: success, blocked, or unknown.") + + +class ActivatedRail(BaseModel): + """A rail that ran during a check, as reported in the generation log.""" + + model_config = ConfigDict(extra="allow") + + name: str = "" + type: str = "" + stop: bool = False + + +class GenerationLog(BaseModel): + """Logging information about a guardrails generation.""" + + model_config = ConfigDict(extra="allow") + + activated_rails: list[ActivatedRail] = Field(default_factory=list) + + +class GuardrailsDataOutput(BaseModel): + """Guardrails-specific output attached to a check or chat response.""" + + model_config = ConfigDict(extra="allow") + + llm_output: dict[str, Any] | None = None + config_ids: list[str] | None = Field(default=None, description="Configuration ids that were used.") + output_data: dict[str, Any] | None = None + log: GenerationLog | None = Field(default=None, description="Populated when guardrails log options are requested.") + + class GuardrailCheckResponse(BaseModel): """Response from a guardrail check request.""" model_config = ConfigDict(extra="allow") + status: str = Field(description="Overall status: success if all rails passed, blocked if any failed.") + rails_status: dict[str, RailStatus] = Field( + default_factory=dict, description="Status of each rail, keyed by rail name." + ) + guardrails_data: GuardrailsDataOutput | None = None + # --------------------------------------------------------------------------- # Request types @@ -67,10 +110,24 @@ class UpdateGuardrailConfigRequest(BaseModel): class GuardrailCheckRequest(BaseModel): - """Guardrail check request body.""" + """Guardrail check request body. + + Shaped like an OpenAI chat-completions request. ``extra="allow"`` passes + through any additional sampling parameters the backend accepts. + """ model_config = ConfigDict(extra="allow") + model: str = Field(description="The model the checked conversation targets.") + messages: list[dict[str, Any]] = Field(description="The conversation to check, in OpenAI chat format.") + guardrails: dict[str, Any] = Field( + default_factory=dict, + description="Guardrails options for the request, e.g. config_id, config, or options.", + ) + max_tokens: int | None = None + temperature: float | None = None + top_p: float | None = None + # --------------------------------------------------------------------------- # Query parameter types diff --git a/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py b/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py index d31b4f57da..7b9d91cff0 100644 --- a/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py @@ -15,44 +15,37 @@ import os import pytest -from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.guardrail.client import GuardrailClient -from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, ListGuardrailConfigsQueryParams +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest WORKSPACE = "default" MODEL = "default/mock-llm" @pytest.fixture -def client() -> NeMoPlatform: +def client() -> GuardrailClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) + return GuardrailClient(base_url=nmp_base_url, workspace=WORKSPACE) @pytest.fixture -def guardrail_config_id(client: NeMoPlatform) -> str: +def guardrail_config_id(client: GuardrailClient) -> str: """Discover the guardrail config the agent created. Lists all guardrail configs and finds one that uses the mock-llm model. Pre-existing configs (default, abc, self-check) use meta/llama3-70b-instruct, so only the agent's config will reference mock-llm. """ - response = ( - client_from_platform(client, GuardrailClient) - .list_guardrail_configs(query_params=ListGuardrailConfigsQueryParams(page=1, page_size=50)) - .data() - ) - configs = response.data + configs = list(client.list_guardrail_configs().items()) assert configs, "No guardrail configurations found. The agent should have created one." # Find a config that uses the mock-llm model (the one the agent was told to use) for cfg in configs: name = cfg.name - full_config = client_from_platform(client, GuardrailClient).get_guardrail_config(name=name).data() - models = (full_config.data.models or []) if full_config.data else [] + full_config = client.get_guardrail_config(name=name).data() + models = full_config.data.get("models") or [] for model in models: - model_name = getattr(model, "model", "") + model_name = model.get("model", "") if "mock-llm" in model_name: config_id = f"{WORKSPACE}/{name}" print(f"Found agent-created guardrail config: {config_id}") @@ -71,50 +64,42 @@ def test_guardrail_config_exists(guardrail_config_id: str) -> None: print(f"Guardrail config exists: {guardrail_config_id}") -def test_content_blocked_by_guardrails(client: NeMoPlatform, guardrail_config_id: str) -> None: +def test_content_blocked_by_guardrails(client: GuardrailClient, guardrail_config_id: str) -> None: """Test that content sent through guardrails is blocked. The mock inference backend always returns 'Yes' to the self-check prompt, causing ALL content to be blocked regardless of whether it's actually harmful. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, ) - .data() - ) + ).data() assert response.status == "blocked", f"Content should be blocked, but got status: {response.status}" print(f"Content correctly blocked with status: {response.status}") -def test_safe_content_also_blocked(client: NeMoPlatform, guardrail_config_id: str) -> None: +def test_safe_content_also_blocked(client: GuardrailClient, guardrail_config_id: str) -> None: """Test that even safe content is blocked (expected with mock backend). Since the mock always returns 'Yes' (block), even safe content gets blocked. This verifies the guardrails pipeline is working end-to-end. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, ) - .data() - ) + ).data() assert response.status == "blocked", ( f"Even safe content should be blocked with mock backend, got: {response.status}" ) diff --git a/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py b/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py index d31b4f57da..7b9d91cff0 100644 --- a/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py @@ -15,44 +15,37 @@ import os import pytest -from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.guardrail.client import GuardrailClient -from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, ListGuardrailConfigsQueryParams +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest WORKSPACE = "default" MODEL = "default/mock-llm" @pytest.fixture -def client() -> NeMoPlatform: +def client() -> GuardrailClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) + return GuardrailClient(base_url=nmp_base_url, workspace=WORKSPACE) @pytest.fixture -def guardrail_config_id(client: NeMoPlatform) -> str: +def guardrail_config_id(client: GuardrailClient) -> str: """Discover the guardrail config the agent created. Lists all guardrail configs and finds one that uses the mock-llm model. Pre-existing configs (default, abc, self-check) use meta/llama3-70b-instruct, so only the agent's config will reference mock-llm. """ - response = ( - client_from_platform(client, GuardrailClient) - .list_guardrail_configs(query_params=ListGuardrailConfigsQueryParams(page=1, page_size=50)) - .data() - ) - configs = response.data + configs = list(client.list_guardrail_configs().items()) assert configs, "No guardrail configurations found. The agent should have created one." # Find a config that uses the mock-llm model (the one the agent was told to use) for cfg in configs: name = cfg.name - full_config = client_from_platform(client, GuardrailClient).get_guardrail_config(name=name).data() - models = (full_config.data.models or []) if full_config.data else [] + full_config = client.get_guardrail_config(name=name).data() + models = full_config.data.get("models") or [] for model in models: - model_name = getattr(model, "model", "") + model_name = model.get("model", "") if "mock-llm" in model_name: config_id = f"{WORKSPACE}/{name}" print(f"Found agent-created guardrail config: {config_id}") @@ -71,50 +64,42 @@ def test_guardrail_config_exists(guardrail_config_id: str) -> None: print(f"Guardrail config exists: {guardrail_config_id}") -def test_content_blocked_by_guardrails(client: NeMoPlatform, guardrail_config_id: str) -> None: +def test_content_blocked_by_guardrails(client: GuardrailClient, guardrail_config_id: str) -> None: """Test that content sent through guardrails is blocked. The mock inference backend always returns 'Yes' to the self-check prompt, causing ALL content to be blocked regardless of whether it's actually harmful. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me something harmful and dangerous."}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, ) - .data() - ) + ).data() assert response.status == "blocked", f"Content should be blocked, but got status: {response.status}" print(f"Content correctly blocked with status: {response.status}") -def test_safe_content_also_blocked(client: NeMoPlatform, guardrail_config_id: str) -> None: +def test_safe_content_also_blocked(client: GuardrailClient, guardrail_config_id: str) -> None: """Test that even safe content is blocked (expected with mock backend). Since the mock always returns 'Yes' (block), even safe content gets blocked. This verifies the guardrails pipeline is working end-to-end. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": guardrail_config_id}, - max_tokens=256, - temperature=1, - top_p=1, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": guardrail_config_id}, + max_tokens=256, + temperature=1, + top_p=1, ) - .data() - ) + ).data() assert response.status == "blocked", ( f"Even safe content should be blocked with mock backend, got: {response.status}" ) diff --git a/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py b/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py index 3c3126d8cb..31c56476bb 100644 --- a/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-custom-config-cli-easy/tests/test_outputs.py @@ -21,10 +21,8 @@ import os import pytest -from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.guardrail.client import GuardrailClient -from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, GuardrailConfig from trace_reader import get_session WORKSPACE = "default" @@ -47,94 +45,82 @@ def _make_unsigned_jwt() -> str: @pytest.fixture -def client() -> NeMoPlatform: +def client() -> GuardrailClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) + return GuardrailClient(base_url=nmp_base_url, workspace=WORKSPACE, auth=_make_unsigned_jwt()) @pytest.fixture -def config(client: NeMoPlatform): +def config(client: GuardrailClient) -> GuardrailConfig: """Retrieve the agent-created guardrail config.""" - return client_from_platform(client, GuardrailClient).get_guardrail_config(name=CONFIG_NAME).data() + return client.get_guardrail_config(name=CONFIG_NAME).data() # --- Config structure checks --- -def test_config_exists(config) -> None: +def test_config_exists(config: GuardrailConfig) -> None: """Test that harbor-custom-config was created.""" assert config.name == CONFIG_NAME, f"Expected config name '{CONFIG_NAME}', got '{config.name}'" print(f"Config exists: {config.name}") -def test_config_description_updated(config) -> None: +def test_config_description_updated(config: GuardrailConfig) -> None: """Test that the config description was updated.""" assert config.description == "Updated custom guardrail config", ( f"Expected description 'Updated custom guardrail config', got '{config.description}'" ) -def test_config_has_input_rails(config) -> None: +def test_config_has_input_rails(config: GuardrailConfig) -> None: """Test that the config has input rails configured.""" - data = config.data - assert data is not None, "Config data should not be None" - rails = data.rails - input_rails = rails.input if rails else None - input_flows = input_rails.flows if input_rails else [] + rails = config.data.get("rails") or {} + input_flows = (rails.get("input") or {}).get("flows") or [] assert any("self check input" in f for f in input_flows), ( f"Expected 'self check input' in input rail flows, got {input_flows}" ) print(f"Input rails configured: {input_flows}") -def test_config_has_output_rails(config) -> None: +def test_config_has_output_rails(config: GuardrailConfig) -> None: """Test that the config has output rails configured.""" - data = config.data - assert data is not None, "Config data should not be None" - rails = data.rails - output_rails = rails.output if rails else None - output_flows = output_rails.flows if output_rails else [] + rails = config.data.get("rails") or {} + output_flows = (rails.get("output") or {}).get("flows") or [] assert any("self check output" in f for f in output_flows), ( f"Expected 'self check output' in output rail flows, got {output_flows}" ) print(f"Output rails configured: {output_flows}") -def test_config_uses_guardrails_model(config) -> None: +def test_config_uses_guardrails_model(config: GuardrailConfig) -> None: """Test that the config uses the guardrails-llm model.""" - data = config.data - assert data is not None, "Config data should not be None" - models = data.models or [] - model_names = [getattr(m, "model", "") for m in models] + models = config.data.get("models") or [] + model_names = [m.get("model", "") for m in models] assert any("guardrails-llm" in name for name in model_names), ( f"Expected a model containing 'guardrails-llm', got {model_names}" ) print(f"Models configured: {model_names}") -def test_config_has_input_prompt_about_fruit(config) -> None: +def test_config_has_input_prompt_about_fruit(config: GuardrailConfig) -> None: """Test that the self_check_input prompt checks for fruit mentions.""" - data = config.data - assert data is not None, "Config data should not be None" - prompts = data.prompts or [] - input_prompts = [p for p in prompts if "self_check_input" in getattr(p, "task", "")] + prompts = config.data.get("prompts") or [] + input_prompts = [p for p in prompts if "self_check_input" in p.get("task", "")] assert len(input_prompts) > 0, ( - f"Expected a prompt with task 'self_check_input', got tasks: {[getattr(p, 'task', None) for p in prompts]}" + f"Expected a prompt with task 'self_check_input', got tasks: {[p.get('task') for p in prompts]}" ) - content = (getattr(input_prompts[0], "content", "") or "").lower() + content = (input_prompts[0].get("content") or "").lower() assert "fruit" in content, f"Expected self_check_input prompt to mention 'fruit', got: {content[:200]}" -def test_config_has_output_prompt_about_bread(config) -> None: +def test_config_has_output_prompt_about_bread(config: GuardrailConfig) -> None: """Test that the self_check_output prompt checks for bread baking content.""" - data = config.data - assert data is not None, "Config data should not be None" - prompts = data.prompts or [] - output_prompts = [p for p in prompts if "self_check_output" in getattr(p, "task", "")] + prompts = config.data.get("prompts") or [] + output_prompts = [p for p in prompts if "self_check_output" in p.get("task", "")] assert len(output_prompts) > 0, ( - f"Expected a prompt with task 'self_check_output', got tasks: {[getattr(p, 'task', None) for p in prompts]}" + f"Expected a prompt with task 'self_check_output', got tasks: {[p.get('task') for p in prompts]}" ) - content = (getattr(output_prompts[0], "content", "") or "").lower() + content = (output_prompts[0].get("content") or "").lower() assert "bread" in content or "baking" in content, ( f"Expected self_check_output prompt to mention 'bread' or 'baking', got: {content[:200]}" ) @@ -143,73 +129,61 @@ def test_config_has_output_prompt_about_bread(config) -> None: # --- Functional inference checks --- -def test_input_rail_blocks_fruit_mention(client: NeMoPlatform) -> None: +def test_input_rail_blocks_fruit_mention(client: GuardrailClient) -> None: """Test that a message mentioning fruit is blocked by the input rail. The self_check_input prompt tells the LLM to block messages mentioning fruit. A message about apples should trigger a 'Yes' response from the self-check, causing guardrails to mark the request blocked. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "blocked", f"Message mentioning fruit should be blocked, got: {response.status}" print(f"Input rail correctly blocked fruit mention: {response.status}") -def test_normal_message_passes_through(client: NeMoPlatform) -> None: +def test_normal_message_passes_through(client: GuardrailClient) -> None: """Test that a normal message (no fruit, no bread) passes through both rails. A geography question doesn't mention fruit (passes input rail) and the response won't be about bread baking (passes output rail). """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "success", f"Normal message should NOT be blocked, got: {response.status}" print(f"Normal message passed through: {response.status}") -def test_output_rail_blocks_bread_content(client: NeMoPlatform) -> None: +def test_output_rail_blocks_bread_content(client: GuardrailClient) -> None: """Test that a response about baking bread is blocked by the output rail. The message doesn't mention fruit (passes input rail), but asking about bread baking will elicit a response about baking bread, which the output self-check should mark as blocked. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "blocked", f"Response about baking bread should be blocked, got: {response.status}" print(f"Output rail correctly blocked bread content: {response.status}") diff --git a/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py b/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py index 45b937e4ed..b18f2d8c42 100644 --- a/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py +++ b/tests/agentic-use/guardrails-custom-config-cli/tests/test_outputs.py @@ -21,10 +21,8 @@ import os import pytest -from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.guardrail.client import GuardrailClient -from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest +from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, GuardrailConfig from trace_reader import get_session WORKSPACE = "default" @@ -47,94 +45,82 @@ def _make_unsigned_jwt() -> str: @pytest.fixture -def client() -> NeMoPlatform: +def client() -> GuardrailClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE, access_token=_make_unsigned_jwt()) + return GuardrailClient(base_url=nmp_base_url, workspace=WORKSPACE, auth=_make_unsigned_jwt()) @pytest.fixture -def config(client: NeMoPlatform): +def config(client: GuardrailClient) -> GuardrailConfig: """Retrieve the agent-created guardrail config.""" - return client_from_platform(client, GuardrailClient).get_guardrail_config(name=CONFIG_NAME).data() + return client.get_guardrail_config(name=CONFIG_NAME).data() # --- Config structure checks --- -def test_config_exists(config) -> None: +def test_config_exists(config: GuardrailConfig) -> None: """Test that harbor-custom-config was created.""" assert config.name == CONFIG_NAME, f"Expected config name '{CONFIG_NAME}', got '{config.name}'" print(f"Config exists: {config.name}") -def test_config_description_updated(config) -> None: +def test_config_description_updated(config: GuardrailConfig) -> None: """Test that the config description was updated.""" assert config.description == "Updated custom guardrail config", ( f"Expected description 'Updated custom guardrail config', got '{config.description}'" ) -def test_config_has_input_rails(config) -> None: +def test_config_has_input_rails(config: GuardrailConfig) -> None: """Test that the config has input rails configured.""" - data = config.data - assert data is not None, "Config data should not be None" - rails = data.rails - input_rails = rails.input if rails else None - input_flows = input_rails.flows if input_rails else [] + rails = config.data.get("rails") or {} + input_flows = (rails.get("input") or {}).get("flows") or [] assert any("self check input" in f for f in input_flows), ( f"Expected 'self check input' in input rail flows, got {input_flows}" ) print(f"Input rails configured: {input_flows}") -def test_config_has_output_rails(config) -> None: +def test_config_has_output_rails(config: GuardrailConfig) -> None: """Test that the config has output rails configured.""" - data = config.data - assert data is not None, "Config data should not be None" - rails = data.rails - output_rails = rails.output if rails else None - output_flows = output_rails.flows if output_rails else [] + rails = config.data.get("rails") or {} + output_flows = (rails.get("output") or {}).get("flows") or [] assert any("self check output" in f for f in output_flows), ( f"Expected 'self check output' in output rail flows, got {output_flows}" ) print(f"Output rails configured: {output_flows}") -def test_config_uses_guardrails_model(config) -> None: +def test_config_uses_guardrails_model(config: GuardrailConfig) -> None: """Test that the config uses the guardrails-llm model.""" - data = config.data - assert data is not None, "Config data should not be None" - models = data.models or [] - model_names = [getattr(m, "model", "") for m in models] + models = config.data.get("models") or [] + model_names = [m.get("model", "") for m in models] assert any("guardrails-llm" in name for name in model_names), ( f"Expected a model containing 'guardrails-llm', got {model_names}" ) print(f"Models configured: {model_names}") -def test_config_has_input_prompt_about_fruit(config) -> None: +def test_config_has_input_prompt_about_fruit(config: GuardrailConfig) -> None: """Test that the self_check_input prompt checks for fruit mentions.""" - data = config.data - assert data is not None, "Config data should not be None" - prompts = data.prompts or [] - input_prompts = [p for p in prompts if "self_check_input" in getattr(p, "task", "")] + prompts = config.data.get("prompts") or [] + input_prompts = [p for p in prompts if "self_check_input" in p.get("task", "")] assert len(input_prompts) > 0, ( - f"Expected a prompt with task 'self_check_input', got tasks: {[getattr(p, 'task', None) for p in prompts]}" + f"Expected a prompt with task 'self_check_input', got tasks: {[p.get('task') for p in prompts]}" ) - content = (getattr(input_prompts[0], "content", "") or "").lower() + content = (input_prompts[0].get("content") or "").lower() assert "fruit" in content, f"Expected self_check_input prompt to mention 'fruit', got: {content[:200]}" -def test_config_has_output_prompt_about_bread(config) -> None: +def test_config_has_output_prompt_about_bread(config: GuardrailConfig) -> None: """Test that the self_check_output prompt checks for bread baking content.""" - data = config.data - assert data is not None, "Config data should not be None" - prompts = data.prompts or [] - output_prompts = [p for p in prompts if "self_check_output" in getattr(p, "task", "")] + prompts = config.data.get("prompts") or [] + output_prompts = [p for p in prompts if "self_check_output" in p.get("task", "")] assert len(output_prompts) > 0, ( - f"Expected a prompt with task 'self_check_output', got tasks: {[getattr(p, 'task', None) for p in prompts]}" + f"Expected a prompt with task 'self_check_output', got tasks: {[p.get('task') for p in prompts]}" ) - content = (getattr(output_prompts[0], "content", "") or "").lower() + content = (output_prompts[0].get("content") or "").lower() assert "bread" in content or "baking" in content, ( f"Expected self_check_output prompt to mention 'bread' or 'baking', got: {content[:200]}" ) @@ -143,73 +129,61 @@ def test_config_has_output_prompt_about_bread(config) -> None: # --- Functional inference checks --- -def test_input_rail_blocks_fruit_mention(client: NeMoPlatform) -> None: +def test_input_rail_blocks_fruit_mention(client: GuardrailClient) -> None: """Test that a message mentioning fruit is blocked by the input rail. The self_check_input prompt tells the LLM to block messages mentioning fruit. A message about apples should trigger a 'Yes' response from the self-check, causing guardrails to mark the request blocked. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Tell me about the health benefits of apples"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "blocked", f"Message mentioning fruit should be blocked, got: {response.status}" print(f"Input rail correctly blocked fruit mention: {response.status}") -def test_normal_message_passes_through(client: NeMoPlatform) -> None: +def test_normal_message_passes_through(client: GuardrailClient) -> None: """Test that a normal message (no fruit, no bread) passes through both rails. A geography question doesn't mention fruit (passes input rail) and the response won't be about bread baking (passes output rail). """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "What is the capital of France?"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "What is the capital of France?"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "success", f"Normal message should NOT be blocked, got: {response.status}" print(f"Normal message passed through: {response.status}") -def test_output_rail_blocks_bread_content(client: NeMoPlatform) -> None: +def test_output_rail_blocks_bread_content(client: GuardrailClient) -> None: """Test that a response about baking bread is blocked by the output rail. The message doesn't mention fruit (passes input rail), but asking about bread baking will elicit a response about baking bread, which the output self-check should mark as blocked. """ - response = ( - client_from_platform(client, GuardrailClient) - .check_guardrail( - body=GuardrailCheckRequest( - model=MODEL, - messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], - guardrails={"config_id": CONFIG_ID}, - max_tokens=256, - temperature=0, - ) + response = client.check_guardrail( + body=GuardrailCheckRequest( + model=MODEL, + messages=[{"role": "user", "content": "Give me a step-by-step guide for baking sourdough bread"}], + guardrails={"config_id": CONFIG_ID}, + max_tokens=256, + temperature=0, ) - .data() - ) + ).data() assert response.status == "blocked", f"Response about baking bread should be blocked, got: {response.status}" print(f"Output rail correctly blocked bread content: {response.status}") From 55f1122de01f5704da21482ba3a72fbeb63cdc87 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Wed, 26 Aug 2026 18:21:16 -0400 Subject: [PATCH 3/3] fix(guardrail): expect typed BadRequestError in e2e config-id rejection test test_checks_rejects_unknown_config_id still expected the pre-migration Stainless nemo_platform.APIStatusError, but _post_check now goes through the typed GuardrailClient, which raises nemo_platform_plugin's BadRequestError on a 400. pytest.raises therefore never matched and the test failed despite correct server behavior. Catch BadRequestError (the narrow 400 class) and keep the status_code == 400 assertion. Signed-off-by: Max Dubrinsky --- e2e/guardrails/test_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/e2e/guardrails/test_checks.py b/e2e/guardrails/test_checks.py index 4746962b28..8b53c6ae0f 100644 --- a/e2e/guardrails/test_checks.py +++ b/e2e/guardrails/test_checks.py @@ -11,9 +11,9 @@ from collections.abc import Callable from typing import Any -import nemo_platform import pytest from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import BadRequestError from nemo_platform_plugin.guardrail.client import GuardrailClient from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, GuardrailCheckResponse @@ -149,7 +149,7 @@ def test_checks_rejects_unknown_config_id( config_mode="referenced", outcome="safe", rail_types=("input",) ) - with pytest.raises(nemo_platform.APIStatusError) as exc_info: + with pytest.raises(BadRequestError) as exc_info: _post_check( test_case, extra_guardrails={"config_id": f"{test_case.workspace}/missing-guardrails-config"},