# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. import abc import logging from collections.abc import Mapping from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, Identifiable, JsonResponseConfig, Message, MessagePiece, TargetIdentifier, ) from pyrit.prompt_target.common.target_capabilities import ( CapabilityName, TargetCapabilities, get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_history import filter_non_replayable_messages from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) # Authentication modes a target can expose to the create-target catalog / API. # ``api_key`` passes a key (from params or the target's env var); ``identity`` # omits the key so the target authenticates itself via an ambient Azure identity # (e.g. minting a Microsoft Entra ID token for its own endpoint, or falling back # to ``DefaultAzureCredential``). AuthMode = Literal["api_key", "identity"] class PromptTarget(Identifiable): """ Abstract base class for prompt targets. A prompt target is a destination where prompts can be sent to interact with various services, models, or APIs. This class defines the interface that all prompt targets must implement. """ _memory: MemoryInterface # A list of Converters that are supported by the prompt target. # An empty list implies that the prompt target supports all converters. supported_converters: list[Any] _identifier: ComponentIdentifier | None = None # Class-level default configuration for this target type. # # Subclasses **should** override this when their capabilities differ from the base # defaults (e.g., to declare multi-turn support or non-text modalities). # Overriding is *optional* — if a subclass does not define ``_DEFAULT_CONFIGURATION``, # it inherits the base-class default (text-only, single-turn, no JSON response). # # Per-instance overrides are also possible via the ``custom_configuration`` # constructor parameter, which takes precedence over the class-level value. _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration(capabilities=TargetCapabilities()) # Declarative auth facts consumed by the create-target service and catalog. # Kept off ``TargetCapabilities`` (auth is a construction/credential axis, not # a message-handling capability) and out of the identity hash. It is surfaced # on ``TargetIdentifier`` only as a ``Param.ClassAttr`` (``Evaluate.Exclude``) # so the registry can read it into ``TargetMetadata`` without building an # instance — never as an identity input or a constructor argument. # # ``supported_auth_modes`` lists the auth modes the create-target API accepts # for this type. Base default is api-key only; targets that can authenticate # via an ambient Azure identity when given no key (e.g. OpenAI, Azure ML, # Azure Blob Storage, Prompt Shield) override this to add ``"identity"``. supported_auth_modes: ClassVar[tuple[AuthMode, ...]] = ("api_key",) def __init_subclass__(cls, **kwargs: object) -> None: """ Validate that subclasses follow the keyword-only ``__init__`` contract. Args: **kwargs: Additional keyword arguments passed to the superclass. Raises: TypeError: If the subclass ``__init__`` accepts positional parameters after ``self``. """ super().__init_subclass__(**kwargs) # Local import to avoid a circular dependency at package init time. from pyrit.common.brick_contract import enforce_keyword_only_init enforce_keyword_only_init(cls, base_name="PromptTarget") def __init__( self, *, verbose: bool = False, max_requests_per_minute: int | None = None, endpoint: str = "", model_name: str = "", underlying_model: str | None = None, custom_configuration: TargetConfiguration | None = None, ) -> None: """ Initialize the PromptTarget. Args: verbose (bool): Enable verbose logging. Defaults to False. max_requests_per_minute (int | None): Maximum number of requests per minute. endpoint (str): The endpoint URL. Defaults to empty string. model_name (str): The model name. Defaults to empty string. underlying_model (str | None): The underlying model name (e.g., "gpt-4o") for identification purposes. This is useful when the deployment name in Azure differs from the actual model. If not provided, ``model_name`` will be used for the identifier. Defaults to None. custom_configuration (TargetConfiguration | None): Override the default configuration for this target instance. Useful for targets whose capabilities depend on deployment configuration (e.g., Playwright, HTTP). If None, uses the class-level ``_DEFAULT_CONFIGURATION``. Defaults to None. """ self._memory = CentralMemory.get_memory_instance() self._verbose = verbose self._max_requests_per_minute = max_requests_per_minute self._endpoint = endpoint self._model_name = model_name self._underlying_model = underlying_model self._configuration = ( custom_configuration if custom_configuration is not None else type(self).get_default_configuration(self._underlying_model) ) if self._verbose: logging.basicConfig(level=logging.INFO) @final async def send_prompt_async( self, *, message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. This is the public entry point called by the prompt normalizer. It: 1. Validates the message. 2. Loads memory history and runs the target's normalization pipeline. 3. Validates the normalized conversation against the target's capabilities. 4. Delegates to ``_send_prompt_to_target_async`` with the normalized conversation. Subclasses MUST NOT override this method. Override ``_send_prompt_to_target_async`` instead. Args: message (Message): The message to send. normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-owned history selection and send lifecycle state. Returns: list[Message]: Response messages from the target. Raises: ValueError: If the message or normalized conversation are empty. """ message.validate() conversation_id = message.get_piece().conversation_id or "" if send_context and send_context.conversation_id != conversation_id: raise ValueError("Target send context conversation_id does not match the current request conversation_id.") if send_context: send_context.begin_send() send_succeeded = False try: normalized_conversation = await self._get_normalized_conversation_async( message=message, normalizer_overrides=normalizer_overrides, send_context=send_context, ) if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) if send_context: send_context.mark_target_invoked() response = await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) send_succeeded = True return response finally: if send_context: send_context.finish_send(succeeded=send_succeeded) @abc.abstractmethod async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ Target-specific send logic. Called by ``send_prompt_async`` after validation and normalization. Args: normalized_conversation (list[Message]): The full conversation (history + current message) after running the normalization pipeline. The current message is the last element. Returns: list[Message]: Response messages from the target. """ def _validate_request(self, *, normalized_conversation: list[Message]) -> None: """ Validate the normalized conversation before sending to the target. Called after the normalization pipeline has run. Validates the last message (the current request) for piece count, data types, and checks whether the full conversation violates multi-turn constraints. Args: normalized_conversation: The normalized conversation to validate. The last element is the current request message. Raises: ValueError: if the target does not support the provided message pieces or if the conversation violates any constraints based on the target's capabilities. """ message = normalized_conversation[-1] n_pieces = len(message.message_pieces) custom_configuration_message = ( "If your target does support this, set the custom_configuration parameter accordingly." ) if not self.configuration.includes(capability=CapabilityName.MULTI_MESSAGE_PIECES) and n_pieces != 1: raise ValueError( f"This target only supports a single message piece. Received: {n_pieces} pieces. " f"{custom_configuration_message}" ) for piece in message.message_pieces: piece_type = piece.converted_value_data_type supported_types_flat = {t for combo in self.capabilities.input_modalities for t in combo} if piece_type not in supported_types_flat: supported_types = ", ".join(sorted(supported_types_flat)) raise ValueError( f"This target supports only the following data types: {supported_types}. Received: {piece_type}. " f"{custom_configuration_message}" ) if not self.configuration.includes(capability=CapabilityName.MULTI_TURN) and len(normalized_conversation) > 1: raise ValueError(f"This target only supports a single turn conversation. {custom_configuration_message}") async def _get_normalized_conversation_async( self, *, message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, send_context: TargetSendContext | None = None, ) -> list[Message]: """ Build the target-facing conversation and run the normalization pipeline. Memory history is loaded and the current message is appended before the target normalization pipeline runs. The original conversation in memory is never mutated. The returned list is an ephemeral copy intended only for building the API request body. After normalization, every output piece is stamped with the current conversation ID. Normalizers own all other output metadata; removed ``prompt_metadata`` keys are not restored. Args: message (Message): The current message to append. normalizer_overrides: Optional per-send target normalizer overrides. send_context: Optional internal coordination contract for caller-approved persisted history. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, history squashed, etc.). """ conversation_id = message.message_pieces[0].conversation_id persisted_messages = ( list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] ) persisted_messages = filter_non_replayable_messages(messages=persisted_messages) conversation = send_context.select_history(messages=persisted_messages) if send_context else persisted_messages conversation.append(message) normalized = await self.configuration.normalize_async( messages=conversation, normalizer_overrides=normalizer_overrides, ) if normalized: for msg in normalized: for piece in msg.message_pieces: piece.conversation_id = conversation_id return normalized def set_model_name(self, *, model_name: str) -> None: """ Set the model name for this target. Args: model_name (str): The model name to set. """ self._model_name = model_name def set_system_prompt( self, *, system_prompt: str, conversation_id: str, ) -> None: """ Inject a system prompt into memory for the given conversation. Writes a ``system``-role message so the target's normalization pipeline (or the target itself, when it natively supports system prompts) will pick it up on the next ``send_prompt_async`` call. If the target does not natively support system prompts, whether this call is ultimately honored depends on the target's ``CapabilityHandlingPolicy``: * ``ADAPT`` — the normalization pipeline (e.g. system squash) will fold the system message into user content on the wire. * ``RAISE`` — the first send after the system prompt is set will raise, because the pipeline cannot adapt the missing capability. Args: system_prompt (str): The system prompt text to set. conversation_id (str): The conversation id to attach the prompt to. Raises: ValueError: If the target does not support multi-turn or editable history. RuntimeError: If the conversation already has messages. """ if not self.capabilities.supports_multi_turn or not self.capabilities.supports_editable_history: raise ValueError( f"Target {type(self).__name__} does not support setting a system prompt. " "It must support both multi-turn conversations and editable history." ) messages = self._memory.get_conversation_messages(conversation_id=conversation_id) if messages: raise RuntimeError("Conversation already exists, system prompt needs to be set at the beginning") self._memory.add_conversation_to_memory( conversation=Conversation(conversation_id=conversation_id, target_identifier=self.get_identifier()) ) self._memory.add_message_to_memory( request=MessagePiece( role="system", conversation_id=conversation_id, original_value=system_prompt, converted_value=system_prompt, ).to_message(), ) async def reset_conversation_async(self, *, conversation_id: str) -> None: """ Release any target-side state held for a conversation. The attack execution scope calls this for objective-target conversations recorded at the common dispatch boundary. Targets that keep external state keyed by conversation (a websocket connection, a browser page, an upstream session) override this to close or discard it. Targets that are stateless between calls need not override it. This is best-effort cleanup, so implementations should not raise for a conversation id they do not recognize, and should be safe to call more than once for the same id. Args: conversation_id (str): The conversation id to release state for. """ def dispose_db_engine(self) -> None: """ Dispose database engine to release database connections and resources. """ self._memory.dispose_engine() def _create_identifier( self, *, params: dict[str, Any] | None = None, targets: list[ComponentIdentifier] | None = None, ) -> ComponentIdentifier: """ Construct the target identifier. Builds a ``TargetIdentifier`` with the base target params (endpoint, model_name, max_requests_per_minute) and the target's promoted child slot. The child slot is exposed as an explicit named parameter (mirroring ``TargetIdentifier``'s promoted field) so it cannot drift into an untyped ``children`` dict. Subclasses should call this method in their _build_identifier() implementation to set the identifier with their specific parameters. Args: params (dict[str, Any] | None): Additional behavioral parameters from the subclass (e.g., temperature, top_p). Merged into the base params. targets (list[ComponentIdentifier] | None): Inner targets of a multi-target (e.g., ``RoundRobinTarget``), promoted to ``TargetIdentifier.targets``. Returns: ComponentIdentifier: The identifier for this prompt target. """ return TargetIdentifier.of( self, params=params, endpoint=self._endpoint, model_name=self._model_name or "", underlying_model_name=self._underlying_model or "", max_requests_per_minute=self._max_requests_per_minute, targets=targets, ) @property def configuration(self) -> TargetConfiguration: """ The configuration of this target instance. Defaults to the class-level ``_DEFAULT_CONFIGURATION``. Can be overridden per instance via the ``custom_configuration`` constructor parameter, which is useful for targets whose capabilities depend on deployment configuration (e.g., Playwright, HTTP). Returns: TargetConfiguration: The configuration for this target. """ return self._configuration @property def capabilities(self) -> TargetCapabilities: """ The capabilities of this target instance. Shorthand for ``self.configuration.capabilities``. Returns: TargetCapabilities: The capabilities for this target. """ return self._configuration.capabilities def apply_capabilities(self, *, capabilities: TargetCapabilities) -> None: """ Replace this target's capabilities, preserving the existing handling policy. The normalization pipeline is rebuilt from the input capabilities and the current policy. Policy is preserved because it expresses user intent (ADAPT vs RAISE), independent of what the probe found. To change policy or normalizer overrides, build a new ``TargetConfiguration`` and pass it via ``custom_configuration`` at construction time instead. Note: This mutates the target's identifier (derived from the configuration). Args: capabilities (TargetCapabilities): The capabilities to install on this instance. """ self._configuration = TargetConfiguration( capabilities=capabilities, policy=self._configuration.policy, normalizer_overrides=self._configuration.normalizer_overrides, ) @classmethod def get_default_configuration(cls, underlying_model: str | None = None) -> TargetConfiguration: """ Return the configuration for the given underlying model, falling back to the class-level ``_DEFAULT_CONFIGURATION`` when the model is not recognized. Args: underlying_model (str | None): The underlying model name (e.g., "gpt-4o"), or None if not specified. Returns: TargetConfiguration: Known configuration for the model, or the class's own ``_DEFAULT_CONFIGURATION`` if the model is unrecognized or not provided. """ if underlying_model: known = get_known_capabilities(underlying_model) if known is not None: return TargetConfiguration(capabilities=known) logger.info( "No known capabilities for model '%s'. Falling back to %s._DEFAULT_CONFIGURATION.", underlying_model, cls.__name__, ) return cls._DEFAULT_CONFIGURATION def _build_identifier(self) -> ComponentIdentifier: """ Build the identifier for this target. Subclasses can override this method to call _create_identifier() with their specific params and children. The base implementation calls _create_identifier() with no extra parameters, which works for targets that don't have model-specific settings. Returns: ComponentIdentifier: The identifier for this prompt target. """ return self._create_identifier() def is_response_format_json(self, message_piece: MessagePiece) -> bool: """ Check if the response format is JSON and ensure the target supports it. Args: message_piece: A MessagePiece object with a `prompt_metadata` dictionary that may include a "response_format" key. Returns: bool: True if the response format is JSON, False otherwise. Raises: ValueError: If "json" response format is requested but unsupported. """ config = self._get_json_response_config(message_piece=message_piece) return config.enabled def _get_json_response_config(self, *, message_piece: MessagePiece) -> JsonResponseConfig: """ Get the JSON response configuration from the message piece metadata. Args: message_piece: A MessagePiece object with a `prompt_metadata` dictionary that may include JSON response configuration. Returns: JsonResponseConfig: The JSON response configuration. Raises: ValueError: If JSON response format is requested but unsupported. """ config = JsonResponseConfig.from_metadata(metadata=message_piece.prompt_metadata) if config.enabled and not self.capabilities.supports_json_output: target_name = self.get_identifier().class_name raise ValueError(f"This target {target_name} does not support JSON response format.") return config