[Summary view] [Print] [Text view]

   1  <?php
   2  
   3  declare (strict_types=1);
   4  namespace WordPress\AiClient\Providers\OpenAiCompatibleImplementation;
   5  
   6  use WordPress\AiClient\Common\Exception\InvalidArgumentException;
   7  use WordPress\AiClient\Files\DTO\File;
   8  use WordPress\AiClient\Files\Enums\MediaOrientationEnum;
   9  use WordPress\AiClient\Messages\DTO\Message;
  10  use WordPress\AiClient\Messages\DTO\MessagePart;
  11  use WordPress\AiClient\Messages\Enums\MessageRoleEnum;
  12  use WordPress\AiClient\Providers\ApiBasedImplementation\AbstractApiBasedModel;
  13  use WordPress\AiClient\Providers\Http\DTO\Request;
  14  use WordPress\AiClient\Providers\Http\DTO\Response;
  15  use WordPress\AiClient\Providers\Http\Enums\HttpMethodEnum;
  16  use WordPress\AiClient\Providers\Http\Exception\ResponseException;
  17  use WordPress\AiClient\Providers\Http\Util\ResponseUtil;
  18  use WordPress\AiClient\Providers\Models\ImageGeneration\Contracts\ImageGenerationModelInterface;
  19  use WordPress\AiClient\Results\DTO\Candidate;
  20  use WordPress\AiClient\Results\DTO\GenerativeAiResult;
  21  use WordPress\AiClient\Results\DTO\TokenUsage;
  22  use WordPress\AiClient\Results\Enums\FinishReasonEnum;
  23  /**
  24   * Base class for an image generation model for providers that implement OpenAI's API format.
  25   *
  26   * This abstract class is designed to work with any AI provider that offers an OpenAI-compatible
  27   * API endpoint for image generation, including but not limited to Anthropic, Google, and other
  28   * providers that have adopted OpenAI's image generation API specification as a standard interface.
  29   *
  30   * @since 0.1.0
  31   *
  32   * @phpstan-type ImageGenerationParams array{
  33   *     model: string,
  34   *     prompt: string,
  35   *     n?: int,
  36   *     response_format?: string,
  37   *     output_format?: string|null,
  38   *     size?: string,
  39   *     ...
  40   * }
  41   * @phpstan-type ChoiceData array{
  42   *     url?: string,
  43   *     b64_json?: string
  44   * }
  45   * @phpstan-type UsageData array{
  46   *     input_tokens?: int,
  47   *     output_tokens?: int,
  48   *     total_tokens?: int
  49   * }
  50   * @phpstan-type ResponseData array{
  51   *     id?: string,
  52   *     data?: list<ChoiceData>,
  53   *     usage?: UsageData
  54   * }
  55   */
  56  abstract class AbstractOpenAiCompatibleImageGenerationModel extends AbstractApiBasedModel implements ImageGenerationModelInterface
  57  {
  58      /**
  59       * {@inheritDoc}
  60       *
  61       * @since 0.1.0
  62       */
  63      public function generateImageResult(array $prompt): GenerativeAiResult
  64      {
  65          $httpTransporter = $this->getHttpTransporter();
  66          $params = $this->prepareGenerateImageParams($prompt);
  67          $request = $this->createRequest(HttpMethodEnum::POST(), 'images/generations', ['Content-Type' => 'application/json'], $params);
  68          // Add authentication credentials to the request.
  69          $request = $this->getRequestAuthentication()->authenticateRequest($request);
  70          // Send and process the request.
  71          $response = $httpTransporter->send($request);
  72          $this->throwIfNotSuccessful($response);
  73          return $this->parseResponseToGenerativeAiResult($response, isset($params['output_format']) && is_string($params['output_format']) ? "image/{$params['output_format']}" : 'image/png');
  74      }
  75      /**
  76       * Prepares the given prompt and the model configuration into parameters for the API request.
  77       *
  78       * @since 0.1.0
  79       *
  80       * @param list<Message> $prompt The prompt to generate an image for. Either a single message or a list of messages
  81       *                              from a chat. However as of today, OpenAI compatible image generation endpoints only
  82       *                              support a single user message.
  83       * @return ImageGenerationParams The parameters for the API request.
  84       */
  85      protected function prepareGenerateImageParams(array $prompt): array
  86      {
  87          $config = $this->getConfig();
  88          $params = ['model' => $this->metadata()->getId(), 'prompt' => $this->preparePromptParam($prompt)];
  89          $candidateCount = $config->getCandidateCount();
  90          if ($candidateCount !== null) {
  91              $params['n'] = $candidateCount;
  92          }
  93          $outputFileType = $config->getOutputFileType();
  94          if ($outputFileType !== null) {
  95              $params['response_format'] = $outputFileType->isRemote() ? 'url' : 'b64_json';
  96          } else {
  97              // The 'response_format' parameter is required, so we default to 'b64_json' if not set.
  98              $params['response_format'] = 'b64_json';
  99          }
 100          $outputMimeType = $config->getOutputMimeType();
 101          if ($outputMimeType !== null) {
 102              $params['output_format'] = preg_replace('/^image\//', '', $outputMimeType);
 103          }
 104          $outputMediaOrientation = $config->getOutputMediaOrientation();
 105          $outputMediaAspectRatio = $config->getOutputMediaAspectRatio();
 106          if ($outputMediaOrientation !== null || $outputMediaAspectRatio !== null) {
 107              $params['size'] = $this->prepareSizeParam($outputMediaOrientation, $outputMediaAspectRatio);
 108          }
 109          /*
 110           * Any custom options are added to the parameters as well.
 111           * This allows developers to pass other options that may be more niche or not yet supported by the SDK.
 112           */
 113          $customOptions = $config->getCustomOptions();
 114          foreach ($customOptions as $key => $value) {
 115              if (isset($params[$key])) {
 116                  throw new InvalidArgumentException(sprintf('The custom option "%s" conflicts with an existing parameter.', $key));
 117              }
 118              $params[$key] = $value;
 119          }
 120          /** @var ImageGenerationParams $params */
 121          return $params;
 122      }
 123      /**
 124       * Prepares the prompt parameter for the API request.
 125       *
 126       * @since 0.1.0
 127       *
 128       * @param list<Message> $messages The messages to prepare. However as of today, OpenAI compatible image generation
 129       *                                endpoints only support a single user message.
 130       * @return string The prepared prompt parameter.
 131       */
 132      protected function preparePromptParam(array $messages): string
 133      {
 134          if (count($messages) !== 1) {
 135              throw new InvalidArgumentException('The API requires a single user message as prompt.');
 136          }
 137          $message = $messages[0];
 138          if (!$message->getRole()->isUser()) {
 139              throw new InvalidArgumentException('The API requires a user message as prompt.');
 140          }
 141          $text = null;
 142          foreach ($message->getParts() as $part) {
 143              $text = $part->getText();
 144              if ($text !== null) {
 145                  break;
 146              }
 147          }
 148          if ($text === null) {
 149              throw new InvalidArgumentException('The API requires a single text message part as prompt.');
 150          }
 151          return $text;
 152      }
 153      /**
 154       * Prepares the size parameter for the API request.
 155       *
 156       * @since 0.1.0
 157       *
 158       * @param MediaOrientationEnum|null $orientation The desired media orientation.
 159       * @param string|null $aspectRatio The desired media aspect ratio.
 160       * @return string The prepared size parameter.
 161       */
 162      protected function prepareSizeParam(?MediaOrientationEnum $orientation, ?string $aspectRatio): string
 163      {
 164          // Use aspect ratio if set, as it is more specific.
 165          if ($aspectRatio !== null) {
 166              switch ($aspectRatio) {
 167                  case '1:1':
 168                      return '1024x1024';
 169                  case '3:2':
 170                      return '1536x1024';
 171                  case '7:4':
 172                      return '1792x1024';
 173                  case '2:3':
 174                      return '1024x1536';
 175                  case '4:7':
 176                      return '1024x1792';
 177                  default:
 178                      throw new InvalidArgumentException('The aspect ratio "' . $aspectRatio . '" is not supported.');
 179              }
 180          }
 181          // This should always have a value, as the method is only called if at least one or the other is set.
 182          if ($orientation !== null) {
 183              if ($orientation->isLandscape()) {
 184                  return '1536x1024';
 185              }
 186              if ($orientation->isPortrait()) {
 187                  return '1024x1536';
 188              }
 189          }
 190          return '1024x1024';
 191      }
 192      /**
 193       * Creates a request object for the provider's API.
 194       *
 195       * Implementations should use $this->getRequestOptions() to attach any
 196       * configured request options to the Request.
 197       *
 198       * @since 0.1.0
 199       *
 200       * @param HttpMethodEnum $method The HTTP method.
 201       * @param string $path The API endpoint path, relative to the base URI.
 202       * @param array<string, string|list<string>> $headers The request headers.
 203       * @param string|array<string, mixed>|null $data The request data.
 204       * @return Request The request object.
 205       */
 206      abstract protected function createRequest(HttpMethodEnum $method, string $path, array $headers = [], $data = null): Request;
 207      /**
 208       * Throws an exception if the response is not successful.
 209       *
 210       * @since 0.1.0
 211       *
 212       * @param Response $response The HTTP response to check.
 213       * @throws ResponseException If the response is not successful.
 214       */
 215      protected function throwIfNotSuccessful(Response $response): void
 216      {
 217          /*
 218           * While this method only calls the utility method, it's important to have it here as a protected method so
 219           * that child classes can override it if needed.
 220           */
 221          ResponseUtil::throwIfNotSuccessful($response);
 222      }
 223      /**
 224       * Parses the response from the API endpoint to a generative AI result.
 225       *
 226       * @since 0.1.0
 227       *
 228       * @param Response $response The response from the API endpoint.
 229       * @param string   $expectedMimeType The expected MIME type the response is in.
 230       * @return GenerativeAiResult The parsed generative AI result.
 231       */
 232      protected function parseResponseToGenerativeAiResult(Response $response, string $expectedMimeType = 'image/png'): GenerativeAiResult
 233      {
 234          /** @var ResponseData $responseData */
 235          $responseData = $response->getData();
 236          if (!isset($responseData['data']) || !$responseData['data']) {
 237              throw ResponseException::fromMissingData($this->providerMetadata()->getName(), 'data');
 238          }
 239          if (!is_array($responseData['data'])) {
 240              throw ResponseException::fromInvalidData($this->providerMetadata()->getName(), 'data', 'The value must be an array.');
 241          }
 242          $candidates = [];
 243          foreach ($responseData['data'] as $index => $choiceData) {
 244              if (!is_array($choiceData) || array_is_list($choiceData)) {
 245                  throw ResponseException::fromInvalidData($this->providerMetadata()->getName(), "data[{$index}]", 'The value must be an associative array.');
 246              }
 247              $candidates[] = $this->parseResponseChoiceToCandidate($choiceData, $index, $expectedMimeType);
 248          }
 249          $id = $this->getResultId($responseData);
 250          if (isset($responseData['usage']) && is_array($responseData['usage'])) {
 251              $usage = $responseData['usage'];
 252              $tokenUsage = new TokenUsage($usage['input_tokens'] ?? 0, $usage['output_tokens'] ?? 0, $usage['total_tokens'] ?? 0);
 253          } else {
 254              $tokenUsage = new TokenUsage(0, 0, 0);
 255          }
 256          // Use any other data from the response as provider-specific response metadata.
 257          $providerMetadata = $responseData;
 258          unset($providerMetadata['id'], $providerMetadata['data'], $providerMetadata['usage']);
 259          return new GenerativeAiResult($id, $candidates, $tokenUsage, $this->providerMetadata(), $this->metadata(), $providerMetadata);
 260      }
 261      /**
 262       * Parses a single choice from the API response into a Candidate object.
 263       *
 264       * @since 0.1.0
 265       *
 266       * @param ChoiceData $choiceData The choice data from the API response.
 267       * @param int $index The index of the choice in the choices array.
 268       * @param string   $expectedMimeType The expected MIME type the response is in.
 269       * @return Candidate The parsed candidate.
 270       * @throws ResponseException If the choice data is invalid.
 271       */
 272      protected function parseResponseChoiceToCandidate(array $choiceData, int $index, string $expectedMimeType = 'image/png'): Candidate
 273      {
 274          // Provider-supplied values are untrusted: only accept http(s) URLs or base64 data, never local paths.
 275          try {
 276              if (isset($choiceData['url']) && is_string($choiceData['url'])) {
 277                  $imageFile = File::fromUrl($choiceData['url'], $expectedMimeType);
 278              } elseif (isset($choiceData['b64_json']) && is_string($choiceData['b64_json'])) {
 279                  $imageFile = File::fromBase64($choiceData['b64_json'], $expectedMimeType);
 280              } else {
 281                  throw ResponseException::fromInvalidData($this->providerMetadata()->getName(), "data[{$index}]", 'The value must contain either a url or b64_json key with a string value.');
 282              }
 283          } catch (InvalidArgumentException $e) {
 284              throw ResponseException::fromInvalidData($this->providerMetadata()->getName(), "data[{$index}]", $e->getMessage());
 285          }
 286          $parts = [new MessagePart($imageFile)];
 287          $message = new Message(MessageRoleEnum::model(), $parts);
 288          return new Candidate($message, FinishReasonEnum::stop());
 289      }
 290      /**
 291       * Extracts the result ID from the API response data.
 292       *
 293       * @since 0.4.0
 294       *
 295       * @param array<string, mixed> $responseData The response data from the API.
 296       * @return string The result ID.
 297       */
 298      protected function getResultId(array $responseData): string
 299      {
 300          return isset($responseData['id']) && is_string($responseData['id']) ? $responseData['id'] : '';
 301      }
 302  }