refactor(guardrail): migrate guardrail consumers to typed GuardrailClient (#1488)

* 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 <mdubrinsky@nvidia.com>

* 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 <mdubrinsky@nvidia.com>

* 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 <mdubrinsky@nvidia.com>

---------

Signed-off-by: Max Dubrinsky <mdubrinsky@nvidia.com>

Author: maxdubrinsky

e2e/guardrails/test_checks.py

@@ -9,16 +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.client.errors import BadRequestError
+from nemo_platform_plugin.guardrail.client import GuardrailClient
+from nemo_platform_plugin.guardrail.types import GuardrailCheckRequest, GuardrailCheckResponse
 
 from e2e.guardrails.utils import (
     BACKEND_RESPONSE,
@@ -27,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,
@@ -51,16 +42,22 @@ 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=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
@@ -152,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"},

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

tests/agentic-use/guardrails-content-safety-cli-easy/tests/test_outputs.py

@@ -15,37 +15,37 @@
 import os
 
 import pytest
-from nemo_platform import NeMoPlatform
+from nemo_platform_plugin.guardrail.client import GuardrailClient
+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.guardrail.configs.list(page=1, page_size=50)
-    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.guardrail.configs.retrieve(name=name)
-        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}")
@@ -64,38 +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.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.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}")
 
 
-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.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.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}"
     )

tests/agentic-use/guardrails-content-safety-cli/tests/test_outputs.py

@@ -15,37 +15,37 @@
 import os
 
 import pytest
-from nemo_platform import NeMoPlatform
+from nemo_platform_plugin.guardrail.client import GuardrailClient
+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.guardrail.configs.list(page=1, page_size=50)
-    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.guardrail.configs.retrieve(name=name)
-        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}")
@@ -64,38 +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.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.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}")
 
 
-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.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.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}"
     )