From 86766ed009cbc5ec781660e16c11435a42c34ba1 Mon Sep 17 00:00:00 2001 From: asimurka Date: Mon, 26 Oct 2026 16:08:59 +0100 Subject: [PATCH] Implemented X-LCS-Shield-After-Last-Refusal logic --- docs/devel_doc/responses.md | 34 +++++ docs/user_doc/shields_guide.md | 6 + src/app/endpoints/responses.py | 11 +- src/models/common/moderation.py | 3 +- src/utils/responses.py | 60 ++++++++- tests/unit/app/endpoints/test_responses.py | 127 ++++++++++++++++++ .../question_validity/test_capability.py | 1 + tests/unit/utils/test_responses.py | 88 ++++++++++++ tests/unit/utils/test_shields.py | 1 + 9 files changed, 326 insertions(+), 5 deletions(-) diff --git a/docs/devel_doc/responses.md b/docs/devel_doc/responses.md index 4320a4b55..afed57998 100644 --- a/docs/devel_doc/responses.md +++ b/docs/devel_doc/responses.md @@ -595,12 +595,46 @@ When server-tool merging is active, the response stream and output include items In streaming mode, server-deployed MCP events (e.g. `mcp_call`, `mcp_list_tools`) are filtered from the SSE stream so clients only see standard output types (`message`, `function_call`, etc.) for their own tools. +### Shield Moderation After Last Refusal + +By default, shields moderate the **entire** extracted `input` (plus attachments). Clients that own their own context (`store: false`, no conversation id) resend full history on every call. Once a blocked turn is in that history, every later request is refused — even when the new user turn is clean. + +Set the `X-LCS-Shield-After-Last-Refusal: true` request header to window moderation: + +```bash +curl -X POST http://localhost:8090/v1/responses \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $TOKEN" \ + -H "X-LCS-Shield-After-Last-Refusal: true" \ + -d '{ + "input": [ + {"type": "message", "role": "user", "content": ""}, + {"type": "message", "role": "assistant", "id": "modr-abc", "content": "I cannot help with that."}, + {"type": "message", "role": "user", "content": "What is Kubernetes?"} + ], + "store": false + }' +``` + +When this header is set to `true`: + +1. LCORE finds the **last** input item whose `id` starts with `modr` (a prior shield refusal). +2. Shields run only on text and attachments extracted from items **after** that refusal. +3. If no `modr-` item is present, behavior matches today (moderate the full input). + +When the header is absent or `false`, full-input moderation applies (no breaking change for conversation / `store: true` clients). + +**Client guidance:** Keep refusal items (with their `modr-…` ids) in resent history so the window can advance. Shield refusals use response and output-item ids of the form `modr-…` (not `resp_…`). Deployments that require blocked content never enter the model should omit this header. + +See also the [Safety Shields Guide](../user_doc/shields_guide.md). + ### LCORE-Specific Extensions The API introduces extensions that are not part of the OpenResponses specification: - `generate_topic_summary` (request) — When set to `true` and a new conversation is created, a topic summary is automatically generated and stored in conversation metadata. - `shield_ids` (request) — Optional list of LCORE shield `name` values to apply. If omitted, all shields from `lightspeed-stack.yaml` are used. See the [Safety Shields Guide](../user_doc/shields_guide.md). +- `X-LCS-Shield-After-Last-Refusal` (request header) — When `true`, shields moderate only input items after the last prior refusal (`modr-…`). See [Shield Moderation After Last Refusal](#shield-moderation-after-last-refusal). - `solr` (request) — Object with optional `mode` (`semantic`, `hybrid`, or `lexical`) and `filters` (Solr vector_io provider payload). Legacy filter-only objects (no `mode`/`filters` wrapper) still work. - `available_quotas` (response) — Provides real-time quota information from all configured quota limiters. diff --git a/docs/user_doc/shields_guide.md b/docs/user_doc/shields_guide.md index f68ce50be..71bd84194 100644 --- a/docs/user_doc/shields_guide.md +++ b/docs/user_doc/shields_guide.md @@ -179,6 +179,12 @@ functionality directly** through a custom API (`run_shield_moderation`) before each request. When moderation blocks the input, the endpoint returns a refusal (and may persist the blocked turn) without calling the model. +On `/v1/responses`, clients that resend full history (`store: false`) can opt +into windowed moderation with the +`X-LCS-Shield-After-Last-Refusal: true` header. Shields then judge only input +items after the last prior refusal (`modr-…`). See +[Shield Moderation After Last Refusal](../devel_doc/responses.md#shield-moderation-after-last-refusal). + ## Per-endpoint behavior | Endpoint | How shields run | `shield_ids` | diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index c8a1b115e..715f3100b 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -104,7 +104,7 @@ build_turn_summary, check_model_configured, deduplicate_referenced_documents, - extract_attachments_text, + extract_shield_moderation_text, extract_text_from_response_items, extract_token_usage, extract_vector_store_ids_from_tools, @@ -651,13 +651,18 @@ async def handle_responses_with_tracing( # pylint: disable=too-many-locals ): client = await AsyncOgxClientHolder().update_azure_token() - attachments_text = extract_attachments_text(original_request.input) add_span_event(root_span, SpanEvents.VALIDATION_COMPLETED) endpoint_path = ENDPOINT_PATH_RESPONSES + shield_after_last_refusal = ( + request.headers.get("X-LCS-Shield-After-Last-Refusal", "").lower() == "true" + ) moderation_result = await run_shield_moderation_v2( - input_text + "\n\n" + attachments_text, + extract_shield_moderation_text( + original_request.input, + shield_after_last_refusal, + ), configuration.configuration.shields, responses_request.shield_ids, ) diff --git a/src/models/common/moderation.py b/src/models/common/moderation.py index 575d6e535..00e906760 100644 --- a/src/models/common/moderation.py +++ b/src/models/common/moderation.py @@ -23,10 +23,11 @@ class ShieldModerationBlocked(BaseModel): @property def refusal_response(self) -> ResponseMessage: - """Build a ResponseMessage carrying the shield's refusal text.""" + """Build a refusal message; ``id`` is the moderation id (``modr-…``).""" return ResponseMessage( role="assistant", content=self.message, + id=self.moderation_id, ) diff --git a/src/utils/responses.py b/src/utils/responses.py index 9cacadf91..9124715a7 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -135,7 +135,7 @@ normalize_vertex_ai_model_id, prepare_input, ) -from utils.suid import to_ogx_conversation_id +from utils.suid import is_moderation_id, to_ogx_conversation_id from utils.token_counter import TokenCounter logger = get_logger(__name__) @@ -1780,6 +1780,64 @@ def extract_attachments_text(response_input: ResponseInput) -> str: return "\n\n".join(file_data_parts) +def get_shield_moderation_input( + response_input: ResponseInput, + after_last_refusal: bool = False, +) -> ResponseInput: + """Return input for shield moderation, optionally after the last refusal. + + Args: + response_input: Responses API ``input`` (string or item list). + after_last_refusal: When True, window after the last ``modr-…`` item. + + Returns: + The full input, or the suffix list after the last moderated refusal. + + """ + if not after_last_refusal or isinstance(response_input, str): + return response_input + + item_count = len(response_input) + for offset, item in enumerate(reversed(response_input)): + if ( + isinstance(item, ResponseMessage) + and item.role == "assistant" + and item.id is not None + and is_moderation_id(item.id) + ): + return list(response_input[item_count - offset :]) + + return response_input + + +def extract_shield_moderation_text( + response_input: ResponseInput, + after_last_refusal: bool = False, +) -> str: + """Extract text and attachments for shield moderation. + + Optionally windows to items after the last ``modr-…`` refusal first. + + Args: + response_input: Responses API ``input`` (string or item list). + after_last_refusal: When True, window after the last ``modr-…`` item. + + Returns: + Combined message text and attachment ``file_data`` for shields. + + """ + shield_input = get_shield_moderation_input( + response_input, + after_last_refusal=after_last_refusal, + ) + text = ( + shield_input + if isinstance(shield_input, str) + else extract_text_from_response_items(shield_input) + ) + return text + "\n\n" + extract_attachments_text(shield_input) + + def _merge_tools( client_tools: list[InputTool], server_tools: list[InputTool], diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index 7f86a7d90..934be9f5e 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -671,6 +671,133 @@ async def test_responses_blocked_with_conversation_appends_refusal( assert "model" in payload, "Handler must set model on the response payload" ResponsesResponse.model_validate(payload) + @pytest.mark.asyncio + async def test_shield_after_last_refusal_header_windows_moderation( + self, + minimal_config: AppConfig, + mocker: MockerFixture, + ) -> None: + """X-LCS-Shield-After-Last-Refusal moderates only items after last modr-.""" + structured_input: list[Any] = [ + OpenAIResponseMessage(role="user", content="blocked earlier"), + OpenAIResponseMessage( + role="assistant", + content="I cannot help with that.", + id="modr-abc", + ), + OpenAIResponseMessage(role="user", content="What is Kubernetes?"), + ] + responses_request = ResponsesRequest( + input=cast("Any", structured_input), + store=False, + ) + req = Request( + scope={ + "type": "http", + "headers": [(b"x-lcs-shield-after-last-refusal", b"true")], + } + ) + req.state.authorized_actions = { + Action.RESPONSES, + Action.READ_OTHERS_CONVERSATIONS, + } + + _patch_base(mocker, minimal_config) + _patch_client(mocker) + _patch_resolve_response_context(mocker) + mocker.patch( + f"{MODULE}.select_model_for_responses", + new=mocker.AsyncMock(return_value="provider/model1"), + ) + mocker.patch( + f"{MODULE}.check_model_configured", + new=mocker.AsyncMock(return_value=True), + ) + _patch_rag(mocker) + mock_moderation = mocker.patch( + f"{MODULE}.run_shield_moderation_v2", + new=mocker.AsyncMock(return_value=ShieldModerationPassed()), + ) + mocker.patch( + f"{MODULE}.handle_non_streaming_response", + new=mocker.AsyncMock( + return_value=_make_responses_response( + output_text="K8s is Kubernetes.", + conversation="conv_new", + ) + ), + ) + + await responses_endpoint_handler( + request=req, + responses_request=responses_request, + auth=MOCK_AUTH, + mcp_headers={}, + ) + + moderated_text = mock_moderation.await_args.args[0] + assert moderated_text.startswith("What is Kubernetes?") + assert "blocked earlier" not in moderated_text + assert "I cannot help with that." not in moderated_text + + @pytest.mark.asyncio + async def test_without_shield_after_last_refusal_header_moderates_full_input( + self, + dummy_request: Request, + minimal_config: AppConfig, + mocker: MockerFixture, + ) -> None: + """Without the header, shields still see the full resent history.""" + structured_input: list[Any] = [ + OpenAIResponseMessage(role="user", content="blocked earlier"), + OpenAIResponseMessage( + role="assistant", + content="I cannot help with that.", + id="modr-abc", + ), + OpenAIResponseMessage(role="user", content="What is Kubernetes?"), + ] + responses_request = ResponsesRequest( + input=cast("Any", structured_input), + store=False, + ) + _patch_base(mocker, minimal_config) + _patch_client(mocker) + _patch_resolve_response_context(mocker) + mocker.patch( + f"{MODULE}.select_model_for_responses", + new=mocker.AsyncMock(return_value="provider/model1"), + ) + mocker.patch( + f"{MODULE}.check_model_configured", + new=mocker.AsyncMock(return_value=True), + ) + _patch_rag(mocker) + mock_moderation = mocker.patch( + f"{MODULE}.run_shield_moderation_v2", + new=mocker.AsyncMock(return_value=ShieldModerationPassed()), + ) + mocker.patch( + f"{MODULE}.handle_non_streaming_response", + new=mocker.AsyncMock( + return_value=_make_responses_response( + output_text="K8s is Kubernetes.", + conversation="conv_new", + ) + ), + ) + + await responses_endpoint_handler( + request=dummy_request, + responses_request=responses_request, + auth=MOCK_AUTH, + mcp_headers={}, + ) + + moderated_text = mock_moderation.await_args.args[0] + assert "blocked earlier" in moderated_text + assert "What is Kubernetes?" in moderated_text + @pytest.mark.asyncio async def test_tool_choice_none_without_tools_does_not_load_server_tools( self, diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py index f1116197f..27cff3a76 100644 --- a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py @@ -740,6 +740,7 @@ async def test_rejected_returns_blocked(self, mocker: MockerFixture) -> None: assert result.moderation_id.startswith("modr-") assert result.refusal_response.role == "assistant" assert result.refusal_response.content == DEFAULT_INVALID_QUESTION_RESPONSE + assert result.refusal_response.id == result.moderation_id @pytest.mark.asyncio async def test_unexpected_response_returns_blocked( diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index 758000efb..bbbc6e198 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -12,6 +12,7 @@ from ogx_api.openai_responses import ( AllowedToolsFilter, OpenAIResponseInputToolChoiceAllowedTools, + OpenAIResponseMessage, OpenAIResponseReasoning, ) from ogx_api.openai_responses import ApprovalFilter as OgxApprovalFilter @@ -85,6 +86,7 @@ build_tool_call_summary, build_tool_result_from_mcp_output_item_done, extract_rag_chunks_from_file_search_item, + extract_shield_moderation_text, extract_text_from_response_item, extract_text_from_response_items, extract_token_usage, @@ -92,6 +94,7 @@ filter_tools_by_allowed_entries, get_mcp_tools, get_rag_tools, + get_shield_moderation_input, get_topic_summary, is_server_deployed_output, maybe_get_topic_summary, @@ -385,6 +388,91 @@ def test_extract_text_from_response_items_with_list_content(self) -> None: assert result == "Part 1 Part 2" +class TestGetShieldModerationInput: + """Test cases for get_shield_moderation_input windowing.""" + + @staticmethod + def _msg( + role: str, content: str, *, item_id: Optional[str] = None + ) -> OpenAIResponseMessage: + """Build a real OpenAIResponseMessage for windowing tests.""" + return OpenAIResponseMessage( + role=cast("Any", role), content=content, id=item_id + ) + + def test_flag_false_returns_full_input(self) -> None: + """When after_last_refusal is False, return the full input unchanged.""" + items = [ + self._msg("user", "blocked"), + self._msg("assistant", "refusal", item_id="modr-1"), + self._msg("user", "clean"), + ] + assert get_shield_moderation_input(items, after_last_refusal=False) is items + + def test_string_input_unchanged(self) -> None: + """String input is returned unchanged even when the flag is set.""" + assert get_shield_moderation_input("hello", after_last_refusal=True) == "hello" + + def test_no_modr_item_returns_full_list(self) -> None: + """Without a modr- item, return the full list when the flag is set.""" + items = [ + self._msg("user", "one"), + self._msg("assistant", "two", item_id="resp_1"), + self._msg("user", "three"), + ] + assert get_shield_moderation_input(items, after_last_refusal=True) is items + + def test_windows_after_last_modr_item(self) -> None: + """Return only items after the last modr- refusal when the flag is set.""" + items = [ + self._msg("user", "blocked"), + self._msg("assistant", "refusal", item_id="modr-1"), + self._msg("user", "clean follow-up"), + ] + window = get_shield_moderation_input(items, after_last_refusal=True) + assert isinstance(window, list) + assert len(window) == 1 + assert extract_text_from_response_items(window) == "clean follow-up" + + def test_uses_last_modr_when_multiple(self) -> None: + """When multiple modr- items exist, window after the last one.""" + items = [ + self._msg("user", "bad-1"), + self._msg("assistant", "refusal-1", item_id="modr-1"), + self._msg("user", "bad-2"), + self._msg("assistant", "refusal-2", item_id="modr-2"), + self._msg("user", "clean"), + ] + window = get_shield_moderation_input(items, after_last_refusal=True) + assert isinstance(window, list) + assert len(window) == 1 + assert extract_text_from_response_items(window) == "clean" + + def test_empty_suffix_when_modr_is_last(self) -> None: + """When the last item is a modr- refusal, return an empty list.""" + items = [ + self._msg("user", "blocked"), + self._msg("assistant", "refusal", item_id="modr-1"), + ] + window = get_shield_moderation_input(items, after_last_refusal=True) + assert not window + + def test_extract_shield_moderation_text_windows_suffix(self) -> None: + """extract_shield_moderation_text returns text only from the windowed suffix.""" + items = [ + self._msg("user", "blocked earlier"), + self._msg("assistant", "refusal", item_id="modr-1"), + self._msg("user", "What is Kubernetes?"), + ] + text = extract_shield_moderation_text(items, after_last_refusal=True) + assert text.startswith("What is Kubernetes?") + assert "blocked earlier" not in text + + def test_extract_shield_moderation_text_string_input(self) -> None: + """String input is moderated as-is (plus empty attachments suffix).""" + assert extract_shield_moderation_text("hello") == "hello\n\n" + + class TestGetRAGTools: """Test cases for get_rag_tools utility function.""" diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py index 93d25eccb..451bdd776 100644 --- a/tests/unit/utils/test_shields.py +++ b/tests/unit/utils/test_shields.py @@ -157,6 +157,7 @@ async def test_returns_blocked_on_first_block(self, mocker: MockerFixture) -> No assert isinstance(result, ShieldModerationBlocked) assert result.message == "rejected" + assert result.refusal_response.id == "modr-123" mock_shield.run.assert_called_once() @pytest.mark.asyncio