MAINT: make PromptShield parser dispatch exhaustive

pyrit/prompt_target/prompt_shield_target.py

@@ -4,7 +4,7 @@
 import json
 import logging
 from collections.abc import Callable
-from typing import Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, assert_never
 
 from pyrit.auth import (
     get_azure_token_provider,
@@ -291,24 +291,29 @@ def _input_parser(self, input_str: str) -> dict[str, Any]:
                 return {"userPrompt": input_str, "documents": []}
             case "documents":
                 return {"userPrompt": "", "documents": [input_str]}
-            case _:
-                # This can also be accomplished using regex, but Python is flexible.
-                user_prompt = ""
-                documents: list[str] = []
-                split_input = input_str.split("<document>")
-
-                for element in split_input:
-                    contents = element.split("</document>")
-
-                    if len(contents) == 1:
-                        user_prompt += contents[0]
-                    else:
-                        # Text before the closing tag is the document; any text after it
-                        # belongs to the user prompt and must not be silently dropped.
-                        documents.append(contents[0])
-                        user_prompt += contents[1]
-
-                return {"userPrompt": user_prompt, "documents": documents if documents else []}
+            case None:
+                pass
+            case unknown_field:
+                if TYPE_CHECKING:
+                    assert_never(unknown_field)
+
+        # This can also be accomplished using regex, but Python is flexible.
+        user_prompt = ""
+        documents: list[str] = []
+        split_input = input_str.split("<document>")
+
+        for element in split_input:
+            contents = element.split("</document>")
+
+            if len(contents) == 1:
+                user_prompt += contents[0]
+            else:
+                # Text before the closing tag is the document; any text after it
+                # belongs to the user prompt and must not be silently dropped.
+                documents.append(contents[0])
+                user_prompt += contents[1]
+
+        return {"userPrompt": user_prompt, "documents": documents if documents else []}
 
     def _add_auth_param_to_headers(self, headers: dict[str, str]) -> None:
         """

tests/unit/prompt_target/target/test_prompt_shield_target.py

@@ -106,6 +106,14 @@ async def test_prompt_shield_document_parsing_keeps_text_between_documents(promp
     assert result["documents"] == [" d1 ", " d2 "]
 
 
+def test_prompt_shield_unknown_entry_field_preserves_runtime_fallback(promptshield_target: PromptShieldTarget):
+    object.__setattr__(promptshield_target, "_force_entry_field", "future-field")
+
+    result = promptshield_target._input_parser("prompt <document> document </document> tail")
+
+    assert result == {"userPrompt": "prompt  tail", "documents": [" document "]}
+
+
 async def test_prompt_shield_response_validation(promptshield_target: PromptShieldTarget):
     # This tests handling both an empty request and an empty response
     promptshield_target._validate_response(request_body={}, response_body={})