# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. import logging from typing import Any from pyrit.common import forward_init_parameters from pyrit.exceptions.exception_classes import ( pyrit_target_retry, ) from pyrit.models import ComponentIdentifier, Message, construct_response_from_request from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import limit_requests_per_minute from pyrit.prompt_target.openai._response_adapter import CompletionsResponseAdapter from pyrit.prompt_target.openai.openai_target import OpenAITarget logger = logging.getLogger(__name__) class OpenAICompletionTarget(OpenAITarget): """A prompt target for OpenAI completion endpoints.""" _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration(capabilities=TargetCapabilities()) _response_adapter = CompletionsResponseAdapter() @forward_init_parameters def __init__( self, *, max_tokens: int | None = None, temperature: float | None = None, top_p: float | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, n: int | None = None, custom_configuration: TargetConfiguration | None = None, **kwargs: Any, ) -> None: """ Initialize the OpenAICompletionTarget with the given parameters. Args: model_name (str, Optional): The name of the model (or deployment name in Azure). If no value is provided, the OPENAI_COMPLETION_MODEL environment variable will be used. endpoint (str, Optional): The target URL for the OpenAI service. api_key (str | Callable[[], str], Optional): The API key for accessing the OpenAI service, or a callable that returns an access token. For Azure endpoints with Entra authentication, pass a token provider from pyrit.auth (e.g., get_azure_openai_auth(endpoint)). Defaults to the `OPENAI_CHAT_KEY` environment variable. headers (str, Optional): Headers of the endpoint (JSON). max_requests_per_minute (int, Optional): Number of requests the target can handle per minute before hitting a rate limit. The number of requests sent to the target will be capped at the value provided. max_tokens (int, Optional): The maximum number of tokens that can be generated in the completion. The token count of your prompt plus `max_tokens` cannot exceed the model's context length. temperature (float, Optional): What sampling temperature to use, between 0 and 2. Values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. top_p (float, Optional): An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. presence_penalty (float, Optional): Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics. frequency_penalty (float, Optional): Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim. n (int, Optional): How many completions to generate for each prompt. custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target instance. Defaults to None. **kwargs: Additional keyword arguments passed to the parent OpenAITarget class. httpx_client_kwargs (dict, Optional): Additional kwargs to be passed to the ``httpx.AsyncClient()`` constructor. For example, to specify a 3 minute timeout: ``httpx_client_kwargs={"timeout": 180}`` """ super().__init__(custom_configuration=custom_configuration, **kwargs) self._max_tokens = max_tokens self._temperature = temperature self._top_p = top_p self._frequency_penalty = frequency_penalty self._presence_penalty = presence_penalty self._n = n def _build_identifier(self) -> ComponentIdentifier: """ Build the identifier with OpenAI completion-specific parameters. Returns: ComponentIdentifier: The identifier for this target instance. """ return self._create_identifier( params={ "temperature": self._temperature, "top_p": self._top_p, "max_tokens": self._max_tokens, "frequency_penalty": self._frequency_penalty, "presence_penalty": self._presence_penalty, "n": self._n, }, ) def _set_openai_env_configuration_vars(self) -> None: self.model_name_environment_variable = "OPENAI_COMPLETION_MODEL" self.endpoint_environment_variable = "OPENAI_COMPLETION_ENDPOINT" self.api_key_environment_variable = "OPENAI_COMPLETION_API_KEY" def _get_target_api_paths(self) -> list[str]: """Return API paths that should not be in the URL.""" return ["/completions", "/v1/completions"] def _get_provider_examples(self) -> dict[str, str]: """Return provider-specific example URLs.""" return { ".openai.azure.com": "https://{resource}.openai.azure.com/openai/v1", "api.openai.com": "https://api.openai.com/v1", } @limit_requests_per_minute @pyrit_target_retry async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ Asynchronously send a message to the OpenAI completion target. 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]: A list containing the response from the prompt target. """ message = normalized_conversation[-1] message_piece = message.message_pieces[0] logger.info(f"Sending the following prompt to the prompt target: {message_piece.converted_value}") # Build request parameters body_parameters: dict[str, Any] = { "model": self._model_name, "prompt": message_piece.converted_value, "top_p": self._top_p, "temperature": self._temperature, "frequency_penalty": self._frequency_penalty, "presence_penalty": self._presence_penalty, "max_tokens": self._max_tokens, "n": self._n, } # Filter out None values request_params: dict[str, Any] = {k: v for k, v in body_parameters.items() if v is not None} # Use unified error handler - automatically detects Completion and validates response = await self._handle_openai_request_async( api_call=lambda: self._client.completions.create(**request_params), request=message, ) return [response] async def _construct_message_from_response_async(self, response: Any, request: Any) -> Message: """ Construct a Message from a Completion response. Args: response: The Completion response from OpenAI SDK. request: The original request MessagePiece. Returns: Message: Constructed message with extracted text. """ logger.info(f"Received response from the prompt target with {len(response.choices)} choices") # Extract response text from validated choices extracted_response = [choice.text for choice in response.choices] message = construct_response_from_request(request=request, response_text_pieces=extracted_response) self._capture_response_metadata(response=response, pieces=message.message_pieces) return message