From 4adca330b34db31f5ca3bec1231600f35ca09e94 Mon Sep 17 00:00:00 2001 From: Adam Lanicek Date: Wed, 19 Aug 2026 19:51:39 +0200 Subject: [PATCH 1/4] Add output-side shield moderation for LLM response classification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses pentest finding OFFSEC-310 (LCORE-2750, CVSS 8.5 Important): the model can be manipulated into generating creative content (emails, speeches, roleplay) outside its intended scope as a RHEL assistant. Adds output_shields configuration and run_output_shield_moderation() that reuses the existing QuestionValidity capability to classify LLM responses before returning them to the user. Changes: - constants.py: DEFAULT_OUTPUT_MODEL_PROMPT and DEFAULT_OUTPUT_REJECTION_MESSAGE for output classification - config.py: output_shields field on Configuration (same format as input shields), validator ensures unique names across both lists - shields.py: run_output_shield_moderation() — iterates output shields, gracefully handles errors (logs warning, doesn't block response) - rlsapi_v1.py: output moderation after response text extraction, replaces response with rejection message if blocked - responses.py: output moderation for non-streaming (blocks) and streaming (log-only — cannot retroactively block streamed content) 7 unit tests covering: passed, blocked, first-block short-circuit, error resilience (AgentRunError, RuntimeError), text passthrough. RSPEED-3399 --- src/app/endpoints/responses.py | 25 ++- src/app/endpoints/rlsapi_v1.py | 14 +- src/constants.py | 54 +++++ src/models/config.py | 15 +- src/utils/shields.py | 53 +++++ .../utils/test_output_shield_moderation.py | 188 ++++++++++++++++++ 6 files changed, 346 insertions(+), 3 deletions(-) create mode 100644 tests/unit/utils/test_output_shield_moderation.py diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 47035eb9b..feaf4809e 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -105,7 +105,7 @@ select_model_for_responses, ) from utils.rh_identity import get_rh_identity_context -from utils.shields import run_shield_moderation_v2 +from utils.shields import run_output_shield_moderation, run_shield_moderation_v2 from utils.suid import ( normalize_conversation_id, ) @@ -974,6 +974,19 @@ async def response_generator( ) chunk_dict["response"]["output_text"] = turn_summary.llm_response + # Output shield check on completed stream (OFFSEC-310). + # Cannot retroactively block already-streamed content; + # log a warning for monitoring/alerting. + output_moderation = await run_output_shield_moderation( + turn_summary.llm_response or "", + configuration.configuration.output_shields, + ) + if output_moderation.decision == "blocked": + logger.warning( + "Output shield triggered on streamed response " + "(cannot retroactively block)" + ) + yield f"event: {chunk.type or 'error'}\ndata: {json.dumps(chunk_dict)}\n\n" except Exception: if not inference_metric_recorded: @@ -1109,6 +1122,16 @@ async def handle_non_streaming_response( token_usage=token_usage, ) output_text = extract_text_from_response_items(api_response.output) + + # Run output shields on LLM response (OFFSEC-310 / LCORE-2750). + output_moderation = await run_output_shield_moderation( + output_text or "", + configuration.configuration.output_shields, + ) + if output_moderation.decision == "blocked": + logger.info("Output shield blocked response") + output_text = output_moderation.message + # Explicitly append the turn to conversation if context passed by previous response await _append_previous_response_turn( api_params, diff --git a/src/app/endpoints/rlsapi_v1.py b/src/app/endpoints/rlsapi_v1.py index 8e9a3d446..be0179530 100644 --- a/src/app/endpoints/rlsapi_v1.py +++ b/src/app/endpoints/rlsapi_v1.py @@ -63,7 +63,7 @@ get_mcp_tools, ) from utils.rh_identity import AUTH_DISABLED, get_rh_identity_context -from utils.shields import run_shield_moderation_v2 +from utils.shields import run_output_shield_moderation, run_shield_moderation_v2 from utils.suid import get_suid logger = get_logger(__name__) @@ -795,6 +795,18 @@ async def infer_endpoint( # pylint: disable=R0914,R0915 logger.warning("Empty response from LLM for request %s", request_id) response_text = constants.UNABLE_TO_PROCESS_RESPONSE + # Run output shields on LLM response (OFFSEC-310 / LCORE-2750). + # Detects non-technical content (creative writing, persona roleplay) + # that bypassed input filters. + output_moderation = await run_output_shield_moderation( + response_text, configuration.configuration.output_shields + ) + if output_moderation.decision == "blocked": + logger.info( + "Output shield blocked response for request %s", request_id + ) + response_text = output_moderation.message + # Consume quota tokens after successful inference. if quota_id is not None: logger.info( diff --git a/src/constants.py b/src/constants.py index 2d71f6a48..813b41062 100644 --- a/src/constants.py +++ b/src/constants.py @@ -390,3 +390,57 @@ ) SAVED_PROMPTS_DEFAULT_MAX_CONTENT_LENGTH: Final[int] = 10_000 SAVED_PROMPTS_MAX_CONTENT_LENGTH_UPPER_BOUND: Final[int] = 30_000 + +# Output-side classification (OFFSEC-310 / LCORE-2750) +DEFAULT_OUTPUT_MODEL_PROMPT: Final[str] = """ +Instructions: +- You are a response classifying tool +- You are an expert in Red Hat Enterprise Linux (RHEL) and Linux system administration +- Your job is to determine whether an AI assistant's response is related to RHEL, \ +Linux system administration, or related technologies and to provide a one-word response. +- If the response provides technical information about RHEL, Linux, or related \ +technologies (package management, systemd, SELinux, networking, storage, containers, \ +security, troubleshooting, shell scripting, kernel, etc.), answer with the word \ +${allowed}, otherwise answer with the word ${rejected}. +- Do not explain your answer, just provide the one-word response. \ +Do not give any other response. + +Example Response: +To configure SELinux policies on RHEL 9, edit /etc/selinux/config and set SELINUX=enforcing. +Classification: +${allowed} + +Example Response: +Dear Mr. Smith, I am writing to inform you about our marketing strategy for Q3... +Classification: +${rejected} + +Example Response: +Once upon a time in a land far away, there lived a brave knight... +Classification: +${rejected} + +Example Response: +You can use systemctl to manage services. Run: sudo systemctl restart httpd +Classification: +${allowed} + +Example Response: +Sure! As your personal travel agent, I recommend visiting Prague in the spring... +Classification: +${rejected} + +Example Response: +The dnf package manager replaced yum in RHEL 8. Use sudo dnf install to install. +Classification: +${allowed} + +Response: +${message} +Classification: +""" + +DEFAULT_OUTPUT_REJECTION_MESSAGE: Final[str] = ( + "This response was filtered because it contains content outside " + "the scope of RHEL technical assistance." +) diff --git a/src/models/config.py b/src/models/config.py index a70f57e09..df1278a23 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -3320,17 +3320,30 @@ class Configuration(ConfigurationBase): "and a type-specific 'config'.", ) + output_shields: list[ShieldConfiguration] = Field( + default_factory=list, + title="Output shields configuration", + description="Shields that run on LLM output before returning to the " + "user. Same format as input shields but applied post-inference. " + "Typically uses question_validity with an output-classification " + "prompt to detect non-technical or off-topic responses.", + ) + @model_validator(mode="after") def validate_shield_names_unique(self) -> Self: """Reject shields lists containing duplicate names. + Checks both input shields and output shields, and ensures no + name collision across the two lists. + Returns: Self: The model instance after validation. Raises: ValueError: If two or more shields share the same name. """ - names = [shield.name for shield in self.shields] + all_shields = list(self.shields) + list(self.output_shields) + names = [shield.name for shield in all_shields] duplicates = {name for name in names if names.count(name) > 1} if duplicates: raise ValueError( diff --git a/src/utils/shields.py b/src/utils/shields.py index a9a16cc0c..3f6d36afe 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -110,6 +110,59 @@ async def run_shield_moderation_v2( return ShieldModerationPassed() +async def run_output_shield_moderation( + response_text: str, + output_shield_configs: list[ShieldConfiguration], +) -> ShieldModerationResult: + """Run shield moderation on LLM output text. + + Iterates through configured output shields and checks the LLM + response for non-technical or off-topic content before it is + returned to the user. + + Addresses pentest finding OFFSEC-310 (LCORE-2750): the model can + be manipulated into generating creative content outside its intended + scope as a RHEL technical assistant. + + Parameters: + response_text: The LLM response text to classify. + output_shield_configs: List of output shield configurations. + + Returns: + Result indicating if the output was blocked or passed. + """ + if not output_shield_configs: + return ShieldModerationPassed() + + for shield_config in output_shield_configs: + shield = build_shield(shield_config) + + try: + shield_result = await shield.run(response_text) + except (AgentRunError, RuntimeError) as exc: + model_id = getattr( + shield_config.config, "model_id", "unknown-shield-model" + ) + logger.warning( + "Output shield %s failed (model=%s): %s", + shield_config.name, + model_id, + exc, + ) + # Don't block the response if the output shield itself fails — + # return the original response rather than an error. + continue + + if shield_result.decision == "blocked": + logger.info( + "Output shield %s blocked response", + shield_config.name, + ) + return shield_result + + return ShieldModerationPassed() + + def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability: """Build a safety capability instance from a shield configuration. diff --git a/tests/unit/utils/test_output_shield_moderation.py b/tests/unit/utils/test_output_shield_moderation.py new file mode 100644 index 000000000..b3bd30083 --- /dev/null +++ b/tests/unit/utils/test_output_shield_moderation.py @@ -0,0 +1,188 @@ +"""Unit tests for run_output_shield_moderation() in utils/shields.py. + +Tests the output-side shield moderation function that checks LLM +responses for non-technical content before returning to the user. + +RSPEED-3399 / OFFSEC-310 / LCORE-2750 +""" + +import pytest +from pydantic_ai.exceptions import AgentRunError +from pytest_mock import MockerFixture + +from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.config import ( + QuestionValidityConfig, + QuestionValidityShieldConfiguration, + ShieldConfiguration, +) +from utils.shields import run_output_shield_moderation + + +def _output_shield_config(name: str = "output-guard") -> ShieldConfiguration: + """Create a QuestionValidityShieldConfiguration for output testing. + + Parameters: + name: The shield name. + + Returns: + A ShieldConfiguration instance. + """ + return QuestionValidityShieldConfiguration( + name=name, + provider_id="question_validity", + config=QuestionValidityConfig( + model_id="test-model", + model_prompt="Classify: ${message} -> ${allowed} or ${rejected}", + invalid_question_response="Non-technical content detected.", + ), + ) + + +class TestRunOutputShieldModeration: + """Tests for run_output_shield_moderation function.""" + + @pytest.mark.asyncio + async def test_returns_passed_when_no_shields(self) -> None: + """Empty output_shields list should return passed immediately.""" + result = await run_output_shield_moderation("any text", []) + assert isinstance(result, ShieldModerationPassed) + assert result.decision == "passed" + + @pytest.mark.asyncio + async def test_returns_passed_when_shield_allows( + self, mocker: MockerFixture + ) -> None: + """Shield returning passed should result in passed.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + return_value=ShieldModerationPassed() + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + result = await run_output_shield_moderation( + "To configure SELinux, edit /etc/selinux/config...", + [_output_shield_config()], + ) + + assert isinstance(result, ShieldModerationPassed) + mock_shield.run.assert_called_once_with( + "To configure SELinux, edit /etc/selinux/config..." + ) + + @pytest.mark.asyncio + async def test_returns_blocked_when_shield_rejects( + self, mocker: MockerFixture + ) -> None: + """Shield returning blocked should result in blocked.""" + blocked = ShieldModerationBlocked( + decision="blocked", + message="Non-technical content detected.", + moderation_id="test-id", + ) + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=blocked) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + result = await run_output_shield_moderation( + "Dear Mr. Smith, here is the marketing email...", + [_output_shield_config()], + ) + + assert isinstance(result, ShieldModerationBlocked) + assert result.decision == "blocked" + assert result.message == "Non-technical content detected." + + @pytest.mark.asyncio + async def test_stops_on_first_block( + self, mocker: MockerFixture + ) -> None: + """Should return blocked on first shield that blocks.""" + blocked = ShieldModerationBlocked( + decision="blocked", + message="Blocked.", + moderation_id="test-id", + ) + shield1 = mocker.Mock() + shield1.run = mocker.AsyncMock(return_value=blocked) + shield2 = mocker.Mock() + shield2.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + + mocker.patch( + "utils.shields.build_shield", side_effect=[shield1, shield2] + ) + + result = await run_output_shield_moderation( + "some text", + [_output_shield_config("s1"), _output_shield_config("s2")], + ) + + assert result.decision == "blocked" + shield1.run.assert_called_once() + shield2.run.assert_not_called() + + @pytest.mark.asyncio + async def test_continues_on_shield_error( + self, mocker: MockerFixture + ) -> None: + """Shield errors should be logged but not block the response.""" + failing_shield = mocker.Mock() + failing_shield.run = mocker.AsyncMock( + side_effect=AgentRunError("model error") + ) + passing_shield = mocker.Mock() + passing_shield.run = mocker.AsyncMock( + return_value=ShieldModerationPassed() + ) + + mocker.patch( + "utils.shields.build_shield", + side_effect=[failing_shield, passing_shield], + ) + + result = await run_output_shield_moderation( + "some text", + [_output_shield_config("fail"), _output_shield_config("pass")], + ) + + assert isinstance(result, ShieldModerationPassed) + failing_shield.run.assert_called_once() + passing_shield.run.assert_called_once() + + @pytest.mark.asyncio + async def test_continues_on_runtime_error( + self, mocker: MockerFixture + ) -> None: + """RuntimeError from shield should be logged but not block.""" + failing_shield = mocker.Mock() + failing_shield.run = mocker.AsyncMock( + side_effect=RuntimeError("unexpected") + ) + + mocker.patch( + "utils.shields.build_shield", return_value=failing_shield + ) + + result = await run_output_shield_moderation( + "some text", [_output_shield_config()] + ) + + assert isinstance(result, ShieldModerationPassed) + + @pytest.mark.asyncio + async def test_passes_response_text_to_shield( + self, mocker: MockerFixture + ) -> None: + """The response text should be passed to shield.run().""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + return_value=ShieldModerationPassed() + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + response_text = "Use sudo dnf install httpd to install Apache." + await run_output_shield_moderation( + response_text, [_output_shield_config()] + ) + + mock_shield.run.assert_called_once_with(response_text) From 63bfde1878c216817533ee6377cadfff8bc6596c Mon Sep 17 00:00:00 2001 From: Adam Lanicek Date: Wed, 19 Aug 2026 20:29:20 +0200 Subject: [PATCH 2/4] Address CodeRabbit review feedback - Replace api_response.output with refusal_response when output shield blocks, preventing blocked content from leaking via the structured response.output field (security fix) - Replace O(n^2) duplicate-name scan with single-pass set lookup in validate_shield_names_unique (performance) Not addressed (intentionally): - Output-specific defaults for QuestionValidityConfig: the prompt is always explicitly configured in lightspeed-stack.yaml, so the default is never used. Adding a separate config type would be unnecessary boilerplate. --- src/app/endpoints/responses.py | 3 +++ src/models/config.py | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index feaf4809e..14b94399b 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -1131,6 +1131,9 @@ async def handle_non_streaming_response( if output_moderation.decision == "blocked": logger.info("Output shield blocked response") output_text = output_moderation.message + # Replace the structured output too, so the blocked + # content is not leaked via response.output field. + api_response.output = [output_moderation.refusal_response] # Explicitly append the turn to conversation if context passed by previous response await _append_previous_response_turn( diff --git a/src/models/config.py b/src/models/config.py index df1278a23..ce038b10d 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -3343,8 +3343,12 @@ def validate_shield_names_unique(self) -> Self: ValueError: If two or more shields share the same name. """ all_shields = list(self.shields) + list(self.output_shields) - names = [shield.name for shield in all_shields] - duplicates = {name for name in names if names.count(name) > 1} + seen: set[str] = set() + duplicates: set[str] = set() + for shield in all_shields: + if shield.name in seen: + duplicates.add(shield.name) + seen.add(shield.name) if duplicates: raise ValueError( f"Shield names must be unique, found duplicates: {sorted(duplicates)}" From 09de390480b0ad811951438e9796d5cbcbc540bf Mon Sep 17 00:00:00 2001 From: Adam Lanicek Date: Wed, 19 Aug 2026 20:38:19 +0200 Subject: [PATCH 3/4] Reject output shields using input-side default prompt/rejection Add validate_output_shield_prompts_explicit validator that rejects QuestionValidityConfig output shields relying on DEFAULT_MODEL_PROMPT or DEFAULT_INVALID_QUESTION_RESPONSE. These input-side defaults are inappropriate for output classification. Addresses CodeRabbit follow-up: a valid output_shields entry could omit model_prompt and silently use the input-side classification prompt, classifying responses with questions-oriented logic. 3 new tests verify detection of default values. --- src/models/config.py | 37 ++++++++++++++++++ .../utils/test_output_shield_moderation.py | 39 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/src/models/config.py b/src/models/config.py index ce038b10d..d2581d899 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -3355,6 +3355,43 @@ def validate_shield_names_unique(self) -> Self: ) return self + @model_validator(mode="after") + def validate_output_shield_prompts_explicit(self) -> Self: + """Require explicit model_prompt and invalid_question_response for output shields. + + Output shields use the same QuestionValidityConfig as input shields, + but the input-side defaults (DEFAULT_MODEL_PROMPT and + DEFAULT_INVALID_QUESTION_RESPONSE) are inappropriate for output + classification. This validator ensures that output shields explicitly + set both fields. + + Returns: + Self: The model instance after validation. + + Raises: + ValueError: If an output shield uses input-side default prompt + or rejection message. + """ + for shield in self.output_shields: + if not isinstance(shield.config, QuestionValidityConfig): + continue + if shield.config.model_prompt == constants.DEFAULT_MODEL_PROMPT: + raise ValueError( + f"Output shield '{shield.name}' must explicitly set " + f"'model_prompt' — the input-side default prompt is not " + f"suitable for output classification." + ) + if ( + shield.config.invalid_question_response + == constants.DEFAULT_INVALID_QUESTION_RESPONSE + ): + raise ValueError( + f"Output shield '{shield.name}' must explicitly set " + f"'invalid_question_response' — the input-side default " + f"rejection message is not suitable for output classification." + ) + return self + @model_validator(mode="after") def validate_mcp_auth_headers(self) -> Self: """ diff --git a/tests/unit/utils/test_output_shield_moderation.py b/tests/unit/utils/test_output_shield_moderation.py index b3bd30083..24987333f 100644 --- a/tests/unit/utils/test_output_shield_moderation.py +++ b/tests/unit/utils/test_output_shield_moderation.py @@ -10,6 +10,7 @@ from pydantic_ai.exceptions import AgentRunError from pytest_mock import MockerFixture +import constants from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed from models.config import ( QuestionValidityConfig, @@ -186,3 +187,41 @@ async def test_passes_response_text_to_shield( ) mock_shield.run.assert_called_once_with(response_text) + + +class TestOutputShieldConfigValidation: + """Tests that output shields require explicit prompt and rejection message.""" + + def test_default_prompt_detected(self) -> None: + """Output shield using input-side default prompt should be caught.""" + config = QuestionValidityConfig( + model_id="test-model", + # model_prompt not set — uses DEFAULT_MODEL_PROMPT + invalid_question_response="Custom rejection.", + ) + assert config.model_prompt == constants.DEFAULT_MODEL_PROMPT + + def test_default_rejection_detected(self) -> None: + """Output shield using input-side default rejection should be caught.""" + config = QuestionValidityConfig( + model_id="test-model", + model_prompt="Custom prompt: ${message} ${allowed} ${rejected}", + # invalid_question_response not set — uses default + ) + assert ( + config.invalid_question_response + == constants.DEFAULT_INVALID_QUESTION_RESPONSE + ) + + def test_explicit_fields_differ_from_defaults(self) -> None: + """Output shield with explicit fields should differ from defaults.""" + config = QuestionValidityConfig( + model_id="test-model", + model_prompt="Custom output prompt: ${message} ${allowed} ${rejected}", + invalid_question_response="Custom rejection message.", + ) + assert config.model_prompt != constants.DEFAULT_MODEL_PROMPT + assert ( + config.invalid_question_response + != constants.DEFAULT_INVALID_QUESTION_RESPONSE + ) From d98d6ac5bae0b62d015bad8a8aeeaea399f120cc Mon Sep 17 00:00:00 2001 From: Adam Lanicek Date: Thu, 20 Aug 2026 01:08:31 +0200 Subject: [PATCH 4/4] Test Configuration validator for output shield defaults Replace field-level default assertions with Configuration-level validation tests per CodeRabbit feedback. Tests now construct a full Configuration with output_shields and verify that: - Default model_prompt raises ValueError - Default invalid_question_response raises ValueError - Explicit values pass validation This ensures the validate_output_shield_prompts_explicit validator is actually exercised. --- .../utils/test_output_shield_moderation.py | 90 ++++++++++++------- 1 file changed, 60 insertions(+), 30 deletions(-) diff --git a/tests/unit/utils/test_output_shield_moderation.py b/tests/unit/utils/test_output_shield_moderation.py index 94ac5bdb1..34562db04 100644 --- a/tests/unit/utils/test_output_shield_moderation.py +++ b/tests/unit/utils/test_output_shield_moderation.py @@ -10,7 +10,6 @@ from pydantic_ai.exceptions import AgentRunError from pytest_mock import MockerFixture -import constants from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed from models.config import ( QuestionValidityConfig, @@ -166,38 +165,69 @@ async def test_passes_response_text_to_shield(self, mocker: MockerFixture) -> No class TestOutputShieldConfigValidation: - """Tests that output shields require explicit prompt and rejection message.""" - - def test_default_prompt_detected(self) -> None: - """Output shield using input-side default prompt should be caught.""" - config = QuestionValidityConfig( - model_id="test-model", - # model_prompt not set — uses DEFAULT_MODEL_PROMPT - invalid_question_response="Custom rejection.", + """Tests that Configuration rejects output shields with input-side defaults.""" + + @staticmethod + def _minimal_config(**kwargs): + """Build a minimal valid Configuration with the given overrides.""" + from models.config import ( + Configuration, + LlamaStackConfiguration, + ServiceConfiguration, + UserDataCollection, ) - assert config.model_prompt == constants.DEFAULT_MODEL_PROMPT - def test_default_rejection_detected(self) -> None: - """Output shield using input-side default rejection should be caught.""" - config = QuestionValidityConfig( - model_id="test-model", - model_prompt="Custom prompt: ${message} ${allowed} ${rejected}", - # invalid_question_response not set — uses default - ) - assert ( - config.invalid_question_response - == constants.DEFAULT_INVALID_QUESTION_RESPONSE + return Configuration( + name="test", + service=ServiceConfiguration(), + llama_stack=LlamaStackConfiguration( + use_as_library_client=True, + library_client_config_path="/tmp/run.yaml", + ), + user_data_collection=UserDataCollection(), + **kwargs, ) - def test_explicit_fields_differ_from_defaults(self) -> None: - """Output shield with explicit fields should differ from defaults.""" - config = QuestionValidityConfig( - model_id="test-model", - model_prompt="Custom output prompt: ${message} ${allowed} ${rejected}", - invalid_question_response="Custom rejection message.", + def test_rejects_output_shield_with_default_prompt(self) -> None: + """Configuration should reject output shield using input-side default prompt.""" + shield = QuestionValidityShieldConfiguration( + name="bad-output", + provider_id="question_validity", + config=QuestionValidityConfig( + model_id="test-model", + # model_prompt omitted — defaults to DEFAULT_MODEL_PROMPT + invalid_question_response="Custom rejection.", + ), + ) + with pytest.raises(ValueError, match="must explicitly set 'model_prompt'"): + self._minimal_config(output_shields=[shield]) + + def test_rejects_output_shield_with_default_rejection(self) -> None: + """Configuration should reject output shield using input-side default rejection.""" + shield = QuestionValidityShieldConfiguration( + name="bad-output", + provider_id="question_validity", + config=QuestionValidityConfig( + model_id="test-model", + model_prompt="Custom: ${message} ${allowed} ${rejected}", + # invalid_question_response omitted — defaults to DEFAULT_INVALID_QUESTION_RESPONSE + ), ) - assert config.model_prompt != constants.DEFAULT_MODEL_PROMPT - assert ( - config.invalid_question_response - != constants.DEFAULT_INVALID_QUESTION_RESPONSE + with pytest.raises( + ValueError, match="must explicitly set 'invalid_question_response'" + ): + self._minimal_config(output_shields=[shield]) + + def test_accepts_output_shield_with_explicit_fields(self) -> None: + """Configuration should accept output shield with explicit prompt and rejection.""" + shield = QuestionValidityShieldConfiguration( + name="good-output", + provider_id="question_validity", + config=QuestionValidityConfig( + model_id="test-model", + model_prompt="Custom output: ${message} ${allowed} ${rejected}", + invalid_question_response="Custom rejection message.", + ), ) + config = self._minimal_config(output_shields=[shield]) + assert len(config.output_shields) == 1