refactor(guardrail): migrate guardrail consumers to typed GuardrailClient

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}"
     )