@automattic/jetpack-ai-client 0.25.6 → 0.26.0

package/src/components/ai-image/general-purpose-image.tsx

@@ -0,0 +1,303 @@
+/**
+ * External dependencies
+ */
+import {
+	useAnalytics,
+	PLAN_TYPE_UNLIMITED,
+	usePlanType,
+} from '@automattic/jetpack-shared-extension-utils';
+import { Button } from '@wordpress/components';
+import { useCallback, useState } from '@wordpress/element';
+import { __, sprintf } from '@wordpress/i18n';
+import debugFactory from 'debug';
+/**
+ * Internal dependencies
+ */
+import './style.scss';
+import useAiFeature from '../../hooks/use-ai-feature/index.js';
+import usePostContent from '../../hooks/use-post-content.js';
+import useSaveToMediaLibrary from '../../hooks/use-save-to-media-library.js';
+import AiImageModal from './components/ai-image-modal.js';
+import useAiImage from './hooks/use-ai-image.js';
+import useSiteType from './hooks/use-site-type.js';
+import {
+	IMAGE_GENERATION_MODEL_STABLE_DIFFUSION,
+	IMAGE_GENERATION_MODEL_DALL_E_3,
+	GENERAL_IMAGE_FEATURE_NAME,
+} from './types.js';
+
+/**
+ * The type for the callback function that is called when the user selects an image.
+ */
+type SetImageCallbackProps = {
+	id: number;
+	url: string;
+};
+
+type GeneralPurposeImageProps = {
+	placement: string;
+	onClose?: () => void;
+	onSetImage?: ( image: SetImageCallbackProps ) => void;
+};
+
+const debug = debugFactory( 'jetpack-ai:general-purpose-image' );
+
+/**
+ * GeneralPurposeImage component
+ * @param {GeneralPurposeImageProps} props - The component properties.
+ * @return {React.ReactElement} - rendered component.
+ */
+export default function GeneralPurposeImage( {
+	placement,
+	onClose = () => {},
+	onSetImage = () => {},
+}: GeneralPurposeImageProps ) {
+	const [ isFeaturedImageModalVisible, setIsFeaturedImageModalVisible ] = useState( true );
+	const siteType = useSiteType();
+	const postContent = usePostContent();
+	const { saveToMediaLibrary } = useSaveToMediaLibrary();
+	const { tracks } = useAnalytics();
+	const { recordEvent } = tracks;
+	const [ prompt, setPrompt ] = useState( '' );
+
+	// Get feature data
+	const { requireUpgrade, requestsCount, requestsLimit, currentTier, costs } = useAiFeature();
+	const planType = usePlanType( currentTier );
+	const generalImageCost = costs?.[ GENERAL_IMAGE_FEATURE_NAME ]?.activeModel ?? 10;
+	const generalImageActiveModel =
+		generalImageCost === costs?.[ GENERAL_IMAGE_FEATURE_NAME ]?.stableDiffusion
+			? IMAGE_GENERATION_MODEL_STABLE_DIFFUSION
+			: IMAGE_GENERATION_MODEL_DALL_E_3;
+	const isUnlimited = planType === PLAN_TYPE_UNLIMITED;
+	const requestsBalance = requestsLimit - requestsCount;
+	const notEnoughRequests = requestsBalance < generalImageCost;
+
+	const {
+		current,
+		setCurrent,
+		processImageGeneration,
+		handlePreviousImage,
+		handleNextImage,
+		currentImage,
+		currentPointer,
+		images,
+		pointer,
+		imageStyles,
+		guessStyle,
+	} = useAiImage( {
+		cost: generalImageCost,
+		autoStart: false,
+		type: 'general-image-generation',
+		feature: GENERAL_IMAGE_FEATURE_NAME,
+	} );
+
+	const hasPrompt = prompt.length >= 3;
+	const disableInput = notEnoughRequests || currentPointer?.generating || requireUpgrade;
+	const disableAction = disableInput || ! hasPrompt;
+
+	const handleModalClose = useCallback( () => {
+		setIsFeaturedImageModalVisible( false );
+		onClose?.();
+	}, [ onClose ] );
+
+	const handleGenerate = useCallback(
+		async ( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleGenerate', userPrompt, style );
+
+			// track the generate image event
+			recordEvent( 'jetpack_ai_general_image_generation_generate_image', {
+				placement,
+				model: generalImageActiveModel,
+				site_type: siteType,
+				style,
+			} );
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+						style,
+					} );
+				}
+			);
+		},
+		[
+			recordEvent,
+			placement,
+			generalImageActiveModel,
+			siteType,
+			processImageGeneration,
+			postContent,
+			notEnoughRequests,
+		]
+	);
+
+	const handleRegenerate = useCallback(
+		( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleRegenerate', userPrompt );
+			// track the regenerate image event
+			recordEvent( 'jetpack_ai_general_image_generation_generate_another_image', {
+				placement,
+				model: generalImageActiveModel,
+				site_type: siteType,
+				style,
+			} );
+
+			setCurrent( crrt => crrt + 1 );
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+					} );
+				}
+			);
+		},
+		[
+			recordEvent,
+			placement,
+			generalImageActiveModel,
+			siteType,
+			processImageGeneration,
+			postContent,
+			notEnoughRequests,
+			setCurrent,
+		]
+	);
+
+	const handleTryAgain = useCallback(
+		( { userPrompt, style }: { userPrompt?: string; style?: string } ) => {
+			debug( 'handleTryAgain', userPrompt );
+			// track the try again event
+			recordEvent( 'jetpack_ai_general_image_generation_try_again', {
+				placement,
+				model: generalImageActiveModel,
+				site_type: siteType,
+				style,
+			} );
+
+			processImageGeneration( { userPrompt, postContent, notEnoughRequests, style } ).catch(
+				error => {
+					recordEvent( 'jetpack_ai_general_image_generation_error', {
+						placement,
+						error: error?.message,
+						model: generalImageActiveModel,
+						site_type: siteType,
+					} );
+				}
+			);
+		},
+		[
+			recordEvent,
+			placement,
+			generalImageActiveModel,
+			siteType,
+			processImageGeneration,
+			postContent,
+			notEnoughRequests,
+		]
+	);
+
+	const handleAccept = useCallback( () => {
+		// track the accept/use image event
+		recordEvent( 'jetpack_ai_general_image_generation_use_image', {
+			placement,
+			model: generalImageActiveModel,
+			site_type: siteType,
+		} );
+
+		const setImage = image => {
+			onSetImage?.( { id: image.id, url: image.url } );
+			handleModalClose();
+		};
+
+		// If the image is already in the media library, use it directly, if it failed for some reason
+		// save it to the media library and then use it.
+		if ( currentImage?.libraryId ) {
+			setImage( {
+				id: currentImage?.libraryId,
+				url: currentImage?.libraryUrl,
+			} );
+		} else {
+			saveToMediaLibrary( currentImage?.image ).then( image => {
+				setImage( image );
+			} );
+		}
+	}, [
+		recordEvent,
+		placement,
+		generalImageActiveModel,
+		siteType,
+		currentImage?.libraryId,
+		currentImage?.libraryUrl,
+		currentImage?.image,
+		onSetImage,
+		handleModalClose,
+		saveToMediaLibrary,
+	] );
+
+	const generateAgainText = __( 'Generate another image', 'jetpack-ai-client' );
+	const generateText = __( 'Generate', 'jetpack-ai-client' );
+
+	const upgradeDescription = notEnoughRequests
+		? sprintf(
+				// Translators: %d is the cost of generating a featured image.
+				__(
+					"Image generation costs %d requests per image. You don't have enough requests to generate another image.",
+					'jetpack-ai-client'
+				),
+				generalImageCost
+		  )
+		: null;
+
+	const acceptButton = (
+		<Button
+			onClick={ handleAccept }
+			variant="primary"
+			disabled={ ! currentImage?.image || currentImage?.generating }
+		>
+			{ __( 'Insert image', 'jetpack-ai-client' ) }
+		</Button>
+	);
+
+	return (
+		<AiImageModal
+			postContent={ true }
+			images={ images }
+			currentIndex={ current }
+			title={ __( 'Generate an image with AI', 'jetpack-ai-client' ) }
+			cost={ generalImageCost }
+			open={ isFeaturedImageModalVisible }
+			placement={ placement }
+			onClose={ handleModalClose }
+			onTryAgain={ handleTryAgain }
+			onGenerate={ pointer?.current > 0 ? handleRegenerate : handleGenerate }
+			generating={ currentPointer?.generating }
+			notEnoughRequests={ notEnoughRequests }
+			requireUpgrade={ requireUpgrade }
+			upgradeDescription={ upgradeDescription }
+			currentLimit={ requestsLimit }
+			currentUsage={ requestsCount }
+			isUnlimited={ isUnlimited }
+			hasError={ Boolean( currentPointer?.error ) }
+			handlePreviousImage={ handlePreviousImage }
+			handleNextImage={ handleNextImage }
+			acceptButton={ acceptButton }
+			generateButtonLabel={ pointer?.current > 0 ? generateAgainText : generateText }
+			instructionsPlaceholder={ __(
+				"Describe the image you'd like to create and select a style.",
+				'jetpack-ai-client'
+			) }
+			imageStyles={ imageStyles }
+			onGuessStyle={ guessStyle }
+			prompt={ prompt }
+			setPrompt={ setPrompt }
+			inputDisabled={ disableInput }
+			actionDisabled={ disableAction }
+		/>
+	);
+}

package/src/components/ai-image/hooks/use-ai-image.ts

@@ -0,0 +1,339 @@
+/**
+ * External dependencies
+ */
+import { useDispatch, useSelect } from '@wordpress/data';
+import { useCallback, useEffect, useRef, useState } from '@wordpress/element';
+import { __ } from '@wordpress/i18n';
+import { cleanForSlug } from '@wordpress/url';
+import React from 'react';
+/**
+ * Internal dependencies
+ */
+import askQuestionSync from '../../../ask-question/sync.js';
+import useAiFeature from '../../../hooks/use-ai-feature/index.js';
+import { ImageStyleObject, ImageStyle } from '../../../hooks/use-image-generator/constants.js';
+import useImageGenerator from '../../../hooks/use-image-generator/index.js';
+import useSaveToMediaLibrary from '../../../hooks/use-save-to-media-library.js';
+/**
+ * Types
+ */
+import {
+	CoreSelectors,
+	FEATURED_IMAGE_FEATURE_NAME,
+	GENERAL_IMAGE_FEATURE_NAME,
+} from '../types.js';
+import type { RoleType } from '../../../types.js';
+import type { CarrouselImageData, CarrouselImages } from '../components/carrousel.js';
+import type { FeatureControl } from '@automattic/jetpack-shared-extension-utils/store/wordpress-com/types';
+
+type ImageFeatureControl = FeatureControl & {
+	styles: Array< ImageStyleObject > | [];
+};
+
+type AiImageType = 'featured-image-generation' | 'general-image-generation';
+type AiImageFeature = typeof FEATURED_IMAGE_FEATURE_NAME | typeof GENERAL_IMAGE_FEATURE_NAME;
+export type ImageResponse = {
+	image?: string;
+	libraryId?: string;
+	libraryUrl?: string;
+	revisedPrompt?: string;
+};
+
+type ProcessImageGenerationProps = {
+	userPrompt?: string | null;
+	postContent?: string | null;
+	notEnoughRequests: boolean;
+	style?: string;
+};
+
+type GuessStyleFunction = (
+	prompt: string,
+	requestType: string,
+	content: string
+) => Promise< ImageStyle | null >;
+
+type UseAiImageProps = {
+	feature: AiImageFeature;
+	type: AiImageType;
+	cost: number;
+	autoStart?: boolean;
+	previousMediaId?: number;
+};
+
+type UseAiImageReturn = {
+	current: number;
+	setCurrent: ( value: number ) => void;
+	processImageGeneration: ( props: ProcessImageGenerationProps ) => Promise< ImageResponse >;
+	handlePreviousImage: () => void;
+	handleNextImage: () => void;
+	currentImage: CarrouselImageData;
+	currentPointer: CarrouselImageData;
+	images: CarrouselImages;
+	pointer: React.RefObject< number >;
+	imageStyles: Array< ImageStyleObject >;
+	guessStyle: GuessStyleFunction;
+};
+
+/**
+ * Hook to get properties for AiImage
+ *
+ * @param {UseAiImageProps} props - The component properties.
+ * @return {UseAiImageReturn} - Object containing properties for AiImage.
+ */
+export default function useAiImage( {
+	feature,
+	type,
+	cost,
+	autoStart = true,
+	previousMediaId,
+}: UseAiImageProps ) {
+	const { generateImageWithParameters } = useImageGenerator();
+	const { increaseRequestsCount, featuresControl } = useAiFeature();
+	const { saveToMediaLibrary } = useSaveToMediaLibrary();
+	const { createNotice } = useDispatch( 'core/notices' );
+
+	/* Images Control */
+	// pointer keeps track of request/generation iteration
+	const pointer = useRef( 0 );
+	// and current keeps track of what is the image exposed at the moment
+	// TODO: should current be any relevant here? It's just modal/carrousel logic after all
+	const [ current, setCurrent ] = useState( 0 );
+	const [ images, setImages ] = useState< CarrouselImages >( [ { generating: autoStart } ] );
+
+	// map feature-to-control prop, if this goes over 2 options, make a hook for it
+	const featureControl = feature === FEATURED_IMAGE_FEATURE_NAME ? 'featured-image' : 'image';
+	const imageFeatureControl = featuresControl?.[ featureControl ] as ImageFeatureControl;
+	const imageStyles: Array< ImageStyleObject > = imageFeatureControl?.styles;
+
+	/* Merge the image data with the new data. */
+	const updateImages = useCallback( ( data: CarrouselImageData, index ) => {
+		setImages( currentImages => {
+			const newImages = [ ...currentImages ];
+			newImages[ index ] = {
+				...newImages[ index ],
+				...data,
+			};
+			return newImages;
+		} );
+	}, [] );
+
+	// the selec/useEffect combo...
+	const loadedMedia = useSelect(
+		( select: ( store ) => CoreSelectors ) => select( 'core' )?.getMedia?.( previousMediaId ),
+		[ previousMediaId ]
+	);
+	useEffect( () => {
+		if ( loadedMedia ) {
+			updateImages(
+				{
+					image: loadedMedia.source_url,
+					libraryId: loadedMedia.id,
+					libraryUrl: loadedMedia.source_url,
+					generating: false,
+				},
+				pointer.current
+			);
+		}
+	}, [ loadedMedia, updateImages ] );
+
+	/*
+	 * Function to show a snackbar notice on the editor.
+	 */
+	const showSnackbarNotice = useCallback(
+		( message: string ) => {
+			createNotice( 'success', message, {
+				type: 'snackbar',
+				isDismissible: true,
+			} );
+		},
+		[ createNotice ]
+	);
+
+	/*
+	 * Function to update the requests count after a featured image generation.
+	 */
+	const updateRequestsCount = useCallback( () => {
+		increaseRequestsCount( cost );
+	}, [ increaseRequestsCount, cost ] );
+
+	/*
+	 * Function to suggest a name for the image based on the user prompt.
+	 */
+	const getImageNameSuggestion = useCallback( ( userPrompt: string ) => {
+		if ( ! userPrompt ) {
+			return 'image.png';
+		}
+
+		const truncatedPrompt = userPrompt.split( ' ' ).slice( 0, 10 ).join( ' ' );
+		return cleanForSlug( truncatedPrompt ) + '.png';
+	}, [] );
+
+	/*
+	 * Function to generate a new image with the current value of the post content.
+	 */
+	const processImageGeneration = useCallback(
+		( {
+			userPrompt,
+			postContent,
+			notEnoughRequests,
+			style = null,
+		}: ProcessImageGenerationProps ) => {
+			return new Promise< ImageResponse >( ( resolve, reject ) => {
+				if ( previousMediaId && pointer.current === 0 ) {
+					pointer.current++;
+				}
+				updateImages( { generating: true, error: null }, pointer.current );
+
+				// Ensure the site has enough requests to generate the image.
+				if ( notEnoughRequests ) {
+					updateImages(
+						{
+							generating: false,
+							error: new Error(
+								__(
+									"You don't have enough requests to generate another image.",
+									'jetpack-ai-client'
+								)
+							),
+						},
+						pointer.current
+					);
+					resolve( {} );
+					return;
+				}
+
+				/**
+				 * Make a generic call to backend and let it decide the model.
+				 */
+				const generateImagePromise = generateImageWithParameters( {
+					feature,
+					size: '1792x1024', // the size, when the generation happens with DALL-E-3
+					responseFormat: 'b64_json', // the response format, when the generation happens with DALL-E-3
+					messages: [
+						{
+							role: 'jetpack-ai',
+							context: {
+								type,
+								request: userPrompt ? userPrompt : null,
+								content: postContent,
+								style,
+							},
+						},
+					],
+					style: style || '',
+				} );
+
+				const name = getImageNameSuggestion( userPrompt );
+
+				generateImagePromise
+					.then( result => {
+						if ( result.data.length > 0 ) {
+							const image = 'data:image/png;base64,' + result.data[ 0 ].b64_json;
+							const prompt = userPrompt || null;
+							const revisedPrompt = result.data[ 0 ].revised_prompt || null;
+							updateImages( { image, prompt, revisedPrompt }, pointer.current );
+							updateRequestsCount();
+							saveToMediaLibrary( image, name )
+								.then( savedImage => {
+									showSnackbarNotice( __( 'Image saved to media library.', 'jetpack-ai-client' ) );
+									updateImages(
+										{ libraryId: savedImage?.id, libraryUrl: savedImage?.url, generating: false },
+										pointer.current
+									);
+									pointer.current += 1;
+									resolve( {
+										image,
+										libraryId: savedImage?.id,
+										libraryUrl: savedImage?.url,
+										revisedPrompt,
+									} );
+								} )
+								.catch( () => {
+									updateImages( { generating: false }, pointer.current );
+									pointer.current += 1;
+									resolve( { image } );
+								} );
+						}
+					} )
+					.catch( e => {
+						updateImages( { generating: false, error: e }, pointer.current );
+						reject( e );
+					} );
+			} );
+		},
+		[
+			updateImages,
+			generateImageWithParameters,
+			feature,
+			type,
+			updateRequestsCount,
+			saveToMediaLibrary,
+			showSnackbarNotice,
+			getImageNameSuggestion,
+			previousMediaId,
+		]
+	);
+
+	const handlePreviousImage = useCallback( () => {
+		setCurrent( Math.max( current - 1, 0 ) );
+	}, [ current ] );
+
+	const handleNextImage = useCallback( () => {
+		setCurrent( Math.min( current + 1, images.length - 1 ) );
+	}, [ current, images.length ] );
+
+	const guessStyle = useCallback(
+		async function (
+			prompt: string,
+			requestType: string = '',
+			content: string = ''
+		): Promise< ImageStyle | null > {
+			if ( ! imageStyles || ! imageStyles.length ) {
+				return null;
+			}
+
+			const messages = [
+				{
+					role: 'jetpack-ai' as RoleType,
+					context: {
+						type: requestType || 'general-image-guess-style',
+						request: prompt,
+						content,
+					},
+				},
+			];
+
+			try {
+				const style = await askQuestionSync( messages, { feature: 'jetpack-ai-image-generator' } );
+
+				if ( ! style ) {
+					return null;
+				}
+				const styleObject = imageStyles.find( ( { value } ) => value === style );
+
+				if ( ! styleObject ) {
+					return null;
+				}
+
+				return styleObject.value;
+			} catch ( error ) {
+				Promise.reject( error );
+			}
+		},
+		[ imageStyles ]
+	);
+
+	return {
+		current,
+		setCurrent,
+		processImageGeneration,
+		handlePreviousImage,
+		handleNextImage,
+		currentImage: images[ current ],
+		currentPointer: images[ pointer.current ],
+		images,
+		pointer,
+		imageStyles,
+		guessStyle,
+	};
+}

package/src/components/ai-image/hooks/use-site-type.ts

@@ -0,0 +1,26 @@
+/**
+ * External dependencies
+ */
+import { isAtomicSite, isSimpleSite } from '@automattic/jetpack-shared-extension-utils';
+import { useState } from '@wordpress/element';
+
+/**
+ * Hook to get the type of site.
+ *
+ * @return {string} - The type of site.
+ */
+export default function useSiteType() {
+	const getSiteType = () => {
+		if ( isAtomicSite() ) {
+			return 'atomic';
+		}
+		if ( isSimpleSite() ) {
+			return 'simple';
+		}
+		return 'jetpack';
+	};
+
+	const [ siteType ] = useState( getSiteType() );
+
+	return siteType;
+}

package/src/components/ai-image/index.ts

@@ -0,0 +1,10 @@
+import FeaturedImage from './featured-image.js';
+import GeneralPurposeImage from './general-purpose-image.js';
+import { PLACEMENT_MEDIA_SOURCE_DROPDOWN, PLACEMENT_BLOCK_PLACEHOLDER_BUTTON } from './types.js';
+
+export {
+	FeaturedImage,
+	PLACEMENT_MEDIA_SOURCE_DROPDOWN,
+	PLACEMENT_BLOCK_PLACEHOLDER_BUTTON,
+	GeneralPurposeImage,
+};