WIP: Start work on integrating the new WordPress Abilities API

composer.json

@@ -13,7 +13,8 @@
     "php": ">=7.4",
     "yahnis-elsts/plugin-update-checker": "5.1",
     "aws/aws-sdk-php": "^3.300",
-    "woocommerce/action-scheduler": "3.8.1"
+    "woocommerce/action-scheduler": "3.8.1",
+    "wordpress/abilities-api": "^0.1"
   },
   "autoload": {
     "psr-4": {

includes/Classifai/Features/Feature.php

@@ -74,6 +74,10 @@ public function setup() {
 		if ( $this->is_feature_enabled() ) {
 			$this->feature_setup();
 		}
+
+		if ( $this->is_enabled() ) {
+			add_action( 'abilities_api_init', [ $this, 'abilities_api_init' ] );
+		}
 	}
 
 	/**
@@ -84,6 +88,25 @@ public function setup() {
 	public function feature_setup() {
 	}
 
+	/**
+	 * Register an ability after the abilities API is initialized.
+	 *
+	 * Only fires if the Feature is enabled and configured.
+	 */
+	public function abilities_api_init() {
+		if ( function_exists( 'wp_register_ability' ) ) {
+			$this->register_ability();
+		}
+	}
+
+	/**
+	 * Register the ability for the Feature.
+	 *
+	 * Override this method in the Feature to register an ability.
+	 */
+	public function register_ability() {
+	}
+
 	/**
 	 * Assigns user roles to the $roles array.
 	 */

includes/Classifai/Features/ImageGeneration.php

@@ -62,6 +62,81 @@ public function feature_setup() {
 		add_action( 'print_media_templates', [ $this, 'print_media_templates' ] );
 	}
 
+	/**
+	 * Register the ability for the Feature.
+	 */
+	public function register_ability() {
+		/**
+		 * Filter the input schema for the ability.
+		 *
+		 * This allows for adding or modifying the arguments for the ability.
+		 * TODO: If we get rid of our custom REST endpoint, we can change
+		 * how this filter works. Right now we're trying to match what the
+		 * REST endpoint uses but that means we have some unnecessary code here,
+		 * particularly the args part.
+		 *
+		 * @since x.x.x
+		 * @hook classifai_{feature}_ability_input_schema
+		 *
+		 * @param array $schema Array of arguments for the input schema.
+		 *
+		 * @return array Modified array of arguments.
+		 */
+		$input_schema = apply_filters(
+			'classifai_' . static::ID . '_ability_input_schema',
+			[
+				'args' => [
+					'prompt' => [
+						'type'              => 'string',
+						'sanitize_callback' => 'sanitize_text_field',
+						'description'       => esc_html__( 'Prompt to use to generate one or more images.', 'classifai' ),
+					],
+				],
+			]
+		);
+
+		wp_register_ability(
+			'classifai/generate-image',
+			[
+				'label'               => esc_html__( 'Generate an image', 'classifai' ),
+				'description'         => esc_html__( 'Use AI to generate one or more images based on a prompt. Will return either a URL or a base64 encoded image.', 'classifai' ),
+				'input_schema'        => [
+					'type'       => 'object',
+					'properties' => $input_schema['args'],
+					'required'   => [
+						'prompt',
+					],
+				],
+				'output_schema'       => [
+					'type'       => 'object',
+					'properties' => [
+						'images' => [
+							'type'        => 'array',
+							'items'       => [
+								'type'       => 'object',
+								'properties' => [
+									'image' => [
+										'type'        => 'string',
+										'description' => esc_html__( 'The image, either a URL or a base64 encoded image.', 'classifai' ),
+									],
+								],
+								'required'   => [
+									'image',
+								],
+							],
+							'description' => esc_html__( 'The generated images.', 'classifai' ),
+						],
+					],
+				],
+				'execute_callback'    => [ $this, 'abilities_api_callback' ],
+				'permission_callback' => [ $this, 'generate_image_permissions_check' ],
+				'meta'                => [
+					'type' => 'tool',
+				],
+			],
+		);
+	}
+
 	/**
 	 * Register any needed endpoints.
 	 */
@@ -147,6 +222,124 @@ public function rest_endpoint_callback( WP_REST_Request $request ) {
 		return parent::rest_endpoint_callback( $request );
 	}
 
+	/**
+	 * Request handler for the abilities API.
+	 *
+	 * @param array $input The input array.
+	 * @return \WP_REST_Response
+	 */
+	public function abilities_api_callback( array $input ) {
+		$prompt = $input['prompt'] ?? null;
+		$format = $input['format'] ?? 'b64_json';
+
+		unset( $input['prompt'] );
+		unset( $input['format'] );
+
+		$output = $this->run(
+			$prompt,
+			'image_gen',
+			$input,
+		);
+
+		if ( is_wp_error( $output ) ) {
+			return $output;
+		}
+
+		// If the format is not url, return the output as-is.
+		if ( 'url' !== $format ) {
+			return [ 'images' => $output ];
+		}
+
+		// If the format is url, import the images into the Media Library so we have URLs.
+		$cleaned_output = [];
+
+		foreach ( $output as $image ) {
+			$image_url = $this->import_base64_image( $image['image'], $prompt );
+
+			if ( ! is_wp_error( $image_url ) ) {
+				$cleaned_output[] = [ 'image' => $image_url ];
+			} else {
+				return $image_url;
+			}
+		}
+
+		return [ 'images' => $cleaned_output ];
+	}
+
+	/**
+	 * Import a base64 encoded image into the Media Library.
+	 *
+	 * @param string $image The base64 encoded image.
+	 * @param string $title The title of the image. Defaults to using the prompt.
+	 * @return string|WP_Error
+	 */
+	public function import_base64_image( string $image, string $title = '' ) {
+		require_once ABSPATH . 'wp-admin/includes/image.php';
+		require_once ABSPATH . 'wp-admin/includes/file.php';
+		require_once ABSPATH . 'wp-admin/includes/media.php';
+
+		// We expect a base64 encoded image so we decode that here.
+		$image = base64_decode( $image ); // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_decode
+
+		if ( ! $image ) {
+			return new WP_Error( 'image_creation_failed', esc_html__( 'Failed to create image from base64 string.', 'classifai' ) );
+		}
+
+		// Ensure the filename isn't too long.
+		$filename = $this->sanitize_filename( $title );
+
+		// Upload the image to the Media Library.
+		$image = wp_upload_bits( $filename . '.png', null, $image );
+
+		if ( isset( $image['error'] ) && false !== $image['error'] ) {
+			/* translators: %s is the error message. */
+			return new WP_Error( 'image_upload_failed', sprintf( esc_html__( 'Failed to upload image: %s', 'classifai' ), $image['error'] ) );
+		}
+
+		// Ensure the uploaded image is treated as an attachment.
+		$image_id = wp_insert_attachment(
+			[
+				'post_mime_type' => 'image/png',
+				'post_title'     => $filename,
+			],
+			$image['file']
+		);
+
+		if ( ! $image_id ) {
+			return new WP_Error( 'image_insert_failed', esc_html__( 'Failed to insert image into Media Library.', 'classifai' ) );
+		}
+
+		// Generate the attachment metadata.
+		wp_generate_attachment_metadata( $image_id, $image['file'] );
+
+		return $image['url'];
+	}
+
+	/**
+	 * Create a safe filename from a longer string.
+	 *
+	 * Try and truncate at a word boundary.
+	 *
+	 * @param string $filename The filename to sanitize.
+	 * @param int    $max_length The maximum length of the filename.
+	 * @return string
+	 */
+	private function sanitize_filename( string $filename, int $max_length = 80 ): string {
+		if ( strlen( $filename ) <= $max_length ) {
+			return $filename;
+		}
+
+		$filename = substr( $filename, 0, $max_length );
+		$filename = sanitize_file_name( $filename );
+
+		$last_hyphen = strrpos( $filename, '-' );
+		if ( $last_hyphen > $max_length * 0.5 ) {
+			$filename = substr( $filename, 0, $last_hyphen );
+		}
+
+		return $filename;
+	}
+
 	/**
 	 * Registers a Media > Generate Image submenu.
 	 */
@@ -238,7 +431,7 @@ public function enqueue_admin_scripts( string $hook_suffix = '' ) {
 			'classifai-plugin-image-generation-media-modal-js',
 			'classifaiDalleData',
 			[
-				'endpoint'   => 'classifai/v1/generate-image',
+				'endpoint'   => 'wp/v2/abilities/classifai/generate-image/run/',
 				'tabText'    => $number_of_images > 1 ? esc_html__( 'Generate images', 'classifai' ) : esc_html__( 'Generate image', 'classifai' ),
 				'errorText'  => esc_html__( 'Something went wrong. No results found', 'classifai' ),
 				'buttonText' => esc_html__( 'Select image', 'classifai' ),
@@ -419,7 +612,7 @@ public function print_media_templates() {
 		?>
 		<script type="text/html" id="tmpl-dalle-image">
 			<div class="generated-image">
-				<img src="data:image/png;base64,{{{ data.url }}}" />
+				<img src="data:image/png;base64,{{{ data.image }}}" />
 				<button type="button" class="components-button button-secondary button-import"><?php esc_html_e( 'Import into Media Library', 'classifai' ); ?></button>
 				<button type="button" class="components-button is-tertiary button-import-insert"><?php esc_html_e( 'Import and Insert', 'classifai' ); ?></button>
 				<span class="spinner"></span>