diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 47035eb9b..14b94399b 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,19 @@ 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 + # 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( api_params, diff --git a/src/app/endpoints/rlsapi_v1.py b/src/app/endpoints/rlsapi_v1.py index 8e9a3d446..28d3c54c9 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,16 @@ 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 4104db4b6..96457855a 100644 --- a/src/constants.py +++ b/src/constants.py @@ -396,3 +396,57 @@ "Your input contains characters or encoding patterns that cannot be " "processed. Please rephrase your question in plain text." ) + +# 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..d2581d899 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -3320,24 +3320,78 @@ 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] - duplicates = {name for name in names if names.count(name) > 1} + all_shields = list(self.shields) + list(self.output_shields) + 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)}" ) 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/src/utils/shields.py b/src/utils/shields.py index 570a281a8..84d76c4a5 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -127,6 +127,57 @@ 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..34562db04 --- /dev/null +++ b/tests/unit/utils/test_output_shield_moderation.py @@ -0,0 +1,233 @@ +"""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) + + +class TestOutputShieldConfigValidation: + """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, + ) + + 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_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 + ), + ) + 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