@automattic/jetpack-ai-client 0.13.0 → 0.14.0

package/CHANGELOG.md

@@ -5,6 +5,23 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.14.0] - 2024-05-20
+### Added
+- AI Client: Expose HTML render rules type. [#37386]
+- AI Featured Image: Support Stable Diffusion image generation. [#37413]
+
+### Changed
+- AI Client: Change default behavior of Message components [#37365]
+- Updated package dependencies. [#37379] [#37380]
+
+## [0.13.1] - 2024-05-13
+### Added
+- AI Client: Add className to AI Control component. [#37322]
+- AI Client: Add "try again" prop on Extension AI Control. [#37250]
+
+### Changed
+- AI Client: Add event to upgrade handler function of Extension AI Control. [#37224]
+
 ## [0.13.0] - 2024-05-06
 ### Added
 - AI Client: Add wrapper ref to AI Control. [#37145]
@@ -306,6 +323,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Updated package dependencies. [#31659]
 - Updated package dependencies. [#31785]
 
+[0.14.0]: https://github.com/Automattic/jetpack-ai-client/compare/v0.13.1...v0.14.0
+[0.13.1]: https://github.com/Automattic/jetpack-ai-client/compare/v0.13.0...v0.13.1
 [0.13.0]: https://github.com/Automattic/jetpack-ai-client/compare/v0.12.4...v0.13.0
 [0.12.4]: https://github.com/Automattic/jetpack-ai-client/compare/v0.12.3...v0.12.4
 [0.12.3]: https://github.com/Automattic/jetpack-ai-client/compare/v0.12.2...v0.12.3

package/build/ask-question/sync.d.ts

@@ -0,0 +1,30 @@
+import { AskQuestionOptionsArgProps } from './index.js';
+import type { PromptProp } from '../types.js';
+/**
+ * The response data from the AI assistant when doing a sync, not-streamed question.
+ */
+export type ResponseData = {
+    choices: Array<{
+        message: {
+            content: string;
+        };
+    }>;
+};
+/**
+ * A function that asks a question without streaming.
+ *
+ * @param {PromptProp} question - The question to ask. It can be a simple string or an array of PromptMessageItemProps objects.
+ * @param {AskQuestionOptionsArgProps} options - An optional object for additional configuration: postId, feature, model.
+ * @returns {Promise<ResponseData>} - A promise that resolves to an instance of the ResponseData
+ * @example
+ * const question = "What is the meaning of life?";
+ * const options = {
+ *   feature: 'ai-featured-image',
+ *   model: 'gpt-4-turbo'
+ * }
+ * askQuestionSync( question, options ).then( responseData => {
+ *    // access the choices array on the response data
+ *    const content = responseData.choices[ 0 ].message.content;
+ * } );
+ */
+export default function askQuestionSync(question: PromptProp, { postId, feature, model }?: AskQuestionOptionsArgProps): Promise<ResponseData>;

package/build/ask-question/sync.js

@@ -0,0 +1,66 @@
+/**
+ * External dependencies
+ */
+import debugFactory from 'debug';
+/*
+ * Types & constants
+ */
+import requestJwt from '../jwt/index.js';
+const debug = debugFactory('jetpack-ai-client:ask-question-sync');
+/**
+ * A function that asks a question without streaming.
+ *
+ * @param {PromptProp} question - The question to ask. It can be a simple string or an array of PromptMessageItemProps objects.
+ * @param {AskQuestionOptionsArgProps} options - An optional object for additional configuration: postId, feature, model.
+ * @returns {Promise<ResponseData>} - A promise that resolves to an instance of the ResponseData
+ * @example
+ * const question = "What is the meaning of life?";
+ * const options = {
+ *   feature: 'ai-featured-image',
+ *   model: 'gpt-4-turbo'
+ * }
+ * askQuestionSync( question, options ).then( responseData => {
+ *    // access the choices array on the response data
+ *    const content = responseData.choices[ 0 ].message.content;
+ * } );
+ */
+export default async function askQuestionSync(question, { postId = null, feature, model } = {}) {
+    debug('Asking question with no streaming: %o. options: %o', question, {
+        postId,
+        feature,
+        model,
+    });
+    /**
+     * The URL to the AI assistant query endpoint.
+     */
+    const URL = 'https://public-api.wordpress.com/wpcom/v2/jetpack-ai-query';
+    let token = null;
+    try {
+        token = (await requestJwt()).token;
+    }
+    catch (error) {
+        debug('Error getting token: %o', error);
+        return Promise.reject(error);
+    }
+    const body = {
+        question: question,
+        stream: false,
+        postId,
+        feature,
+        model,
+    };
+    const headers = {
+        Authorization: `Bearer ${token}`,
+        'Content-Type': 'application/json',
+    };
+    const data = await fetch(URL, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body),
+    }).then(response => response.json());
+    if (data?.data?.status && data?.data?.status > 200) {
+        debug('Error generating prompt: %o', data);
+        return Promise.reject(data);
+    }
+    return data;
+}

package/build/components/ai-control/ai-control.d.ts

@@ -6,6 +6,7 @@ import './style.scss';
 import type { RequestingStateProp } from '../../types.js';
 import type { ReactElement } from 'react';
 type AIControlProps = {
+    className?: string;
     disabled?: boolean;
     value: string;
     placeholder?: string;
@@ -25,5 +26,5 @@ type AIControlProps = {
  * @param {AIControlProps} props - Component props
  * @returns {ReactElement}       Rendered component
  */
-export default function AIControl({ disabled, value, placeholder, isTransparent, state, onChange, banner, error, actions, message, promptUserInputRef, wrapperRef, }: AIControlProps): ReactElement;
+export default function AIControl({ className, disabled, value, placeholder, isTransparent, state, onChange, banner, error, actions, message, promptUserInputRef, wrapperRef, }: AIControlProps): ReactElement;
 export {};

package/build/components/ai-control/ai-control.js

@@ -15,8 +15,8 @@ import './style.scss';
  * @param {AIControlProps} props - Component props
  * @returns {ReactElement}       Rendered component
  */
-export default function AIControl({ disabled = false, value = '', placeholder = '', isTransparent = false, state = 'init', onChange, banner = null, error = null, actions = null, message = null, promptUserInputRef = null, wrapperRef = null, }) {
-    return (_jsxs("div", { className: "jetpack-components-ai-control__container-wrapper", ref: wrapperRef, children: [error, _jsxs("div", { className: "jetpack-components-ai-control__container", children: [banner, _jsxs("div", { className: classNames('jetpack-components-ai-control__wrapper', {
+export default function AIControl({ className, disabled = false, value = '', placeholder = '', isTransparent = false, state = 'init', onChange, banner = null, error = null, actions = null, message = null, promptUserInputRef = null, wrapperRef = null, }) {
+    return (_jsxs("div", { className: classNames('jetpack-components-ai-control__container-wrapper', className), ref: wrapperRef, children: [error, _jsxs("div", { className: "jetpack-components-ai-control__container", children: [banner, _jsxs("div", { className: classNames('jetpack-components-ai-control__wrapper', {
                             'is-transparent': isTransparent,
                         }), children: [_jsx(AiStatusIndicator, { state: state }), _jsx("div", { className: "jetpack-components-ai-control__input-wrapper", children: _jsx(PlainText, { value: value, onChange: onChange, placeholder: placeholder, className: "jetpack-components-ai-control__input", disabled: disabled, ref: promptUserInputRef }) }), actions] }), message] })] }));
 }

package/build/components/ai-control/extension-ai-control.d.ts

@@ -3,9 +3,10 @@ import './style.scss';
 /**
  * Types
  */
-import type { RequestingStateProp } from '../../types.js';
-import type { ReactElement } from 'react';
+import type { RequestingErrorProps, RequestingStateProp } from '../../types.js';
+import type { ReactElement, MouseEvent } from 'react';
 type ExtensionAIControlProps = {
+    className?: string;
     disabled?: boolean;
     value: string;
     placeholder?: string;
@@ -13,7 +14,7 @@ type ExtensionAIControlProps = {
     isTransparent?: boolean;
     state?: RequestingStateProp;
     showGuideLine?: boolean;
-    error?: string;
+    error?: RequestingErrorProps;
     requestsRemaining?: number;
     showUpgradeMessage?: boolean;
     wrapperRef?: React.MutableRefObject<HTMLDivElement | null>;
@@ -22,7 +23,8 @@ type ExtensionAIControlProps = {
     onStop?: () => void;
     onClose?: () => void;
     onUndo?: () => void;
-    onUpgrade?: () => void;
+    onUpgrade?: (event: MouseEvent<HTMLButtonElement>) => void;
+    onTryAgain?: () => void;
 };
 /**
  * ExtensionAIControl component. Used by the AI Assistant inline extensions, adding logic and components to the base AIControl component.
@@ -31,6 +33,6 @@ type ExtensionAIControlProps = {
  * @param {React.MutableRefObject} ref     - Ref to the component
  * @returns {ReactElement}                 Rendered component
  */
-export declare function ExtensionAIControl({ disabled, value, placeholder, showButtonLabels, isTransparent, state, showGuideLine, error, requestsRemaining, showUpgradeMessage, wrapperRef, onChange, onSend, onStop, onClose, onUndo, onUpgrade, }: ExtensionAIControlProps, ref: React.MutableRefObject<HTMLInputElement>): ReactElement;
+export declare function ExtensionAIControl({ className, disabled, value, placeholder, showButtonLabels, isTransparent, state, showGuideLine, error, requestsRemaining, showUpgradeMessage, wrapperRef, onChange, onSend, onStop, onClose, onUndo, onUpgrade, onTryAgain, }: ExtensionAIControlProps, ref: React.MutableRefObject<HTMLInputElement>): ReactElement;
 declare const _default: React.ForwardRefExoticComponent<ExtensionAIControlProps & React.RefAttributes<HTMLInputElement>>;
 export default _default;

package/build/components/ai-control/extension-ai-control.js

@@ -21,7 +21,7 @@ import './style.scss';
  * @param {React.MutableRefObject} ref     - Ref to the component
  * @returns {ReactElement}                 Rendered component
  */
-export function ExtensionAIControl({ disabled = false, value = '', placeholder = '', showButtonLabels = true, isTransparent = false, state = 'init', showGuideLine = false, error, requestsRemaining, showUpgradeMessage = false, wrapperRef, onChange, onSend, onStop, onClose, onUndo, onUpgrade, }, ref) {
+export function ExtensionAIControl({ className, disabled = false, value = '', placeholder = '', showButtonLabels = true, isTransparent = false, state = 'init', showGuideLine = false, error, requestsRemaining, showUpgradeMessage = false, wrapperRef, onChange, onSend, onStop, onClose, onUndo, onUpgrade, onTryAgain, }, ref) {
     const loading = state === 'requesting' || state === 'suggesting';
     const [editRequest, setEditRequest] = useState(false);
     const [lastValue, setLastValue] = useState(value || null);
@@ -61,9 +61,12 @@ export function ExtensionAIControl({ disabled = false, value = '', placeholder =
     const undoHandler = useCallback(() => {
         onUndo?.();
     }, [onUndo]);
-    const upgradeHandler = useCallback(() => {
-        onUpgrade?.();
+    const upgradeHandler = useCallback((event) => {
+        onUpgrade?.(event);
     }, [onUpgrade]);
+    const tryAgainHandler = useCallback(() => {
+        onTryAgain?.();
+    }, [onTryAgain]);
     useKeyboardShortcut('enter', e => {
         e.preventDefault();
         sendHandler();
@@ -72,8 +75,8 @@ export function ExtensionAIControl({ disabled = false, value = '', placeholder =
     });
     const actions = (_jsx(_Fragment, { children: loading ? (_jsx(Button, { className: "jetpack-components-ai-control__controls-prompt_button", onClick: stopHandler, variant: "secondary", label: __('Stop request', 'jetpack-ai-client'), children: showButtonLabels ? __('Stop', 'jetpack-ai-client') : _jsx(Icon, { icon: closeSmall }) })) : (_jsxs(_Fragment, { children: [value?.length > 0 && (_jsx("div", { className: "jetpack-components-ai-control__controls-prompt_button_wrapper", children: _jsx(Button, { className: "jetpack-components-ai-control__controls-prompt_button", onClick: sendHandler, variant: "primary", disabled: !value?.length || disabled, label: __('Send request', 'jetpack-ai-client'), children: showButtonLabels ? (__('Generate', 'jetpack-ai-client')) : (_jsx(Icon, { icon: arrowUp })) }) })), value?.length <= 0 && state === 'done' && (_jsx("div", { className: "jetpack-components-ai-control__controls-prompt_button_wrapper", children: _jsxs(ButtonGroup, { children: [_jsx(Button, { className: "jetpack-components-ai-control__controls-prompt_button", label: __('Undo', 'jetpack-ai-client'), onClick: undoHandler, tooltipPosition: "top", children: _jsx(Icon, { icon: undo }) }), _jsx(Button, { className: "jetpack-components-ai-control__controls-prompt_button", label: __('Close', 'jetpack-ai-client'), onClick: closeHandler, variant: "tertiary", children: __('Close', 'jetpack-ai-client') })] }) }))] })) }));
     let message = null;
-    if (error) {
-        message = _jsx(ErrorMessage, { error: error, onTryAgainClick: sendHandler });
+    if (error?.message) {
+        message = (_jsx(ErrorMessage, { error: error.message, code: error.code, onTryAgainClick: tryAgainHandler, onUpgradeClick: upgradeHandler }));
     }
     else if (showUpgradeMessage) {
         message = (_jsx(UpgradeMessage, { requestsRemaining: requestsRemaining, onUpgradeClick: upgradeHandler }));
@@ -81,6 +84,6 @@ export function ExtensionAIControl({ disabled = false, value = '', placeholder =
     else if (showGuideLine) {
         message = _jsx(GuidelineMessage, {});
     }
-    return (_jsx(AIControl, { disabled: disabled || loading, value: value, placeholder: placeholder, isTransparent: isTransparent, state: state, onChange: changeHandler, actions: actions, message: message, promptUserInputRef: promptUserInputRef, wrapperRef: wrapperRef }));
+    return (_jsx(AIControl, { className: className, disabled: disabled || loading, value: value, placeholder: placeholder, isTransparent: isTransparent, state: state, onChange: changeHandler, actions: actions, message: message, promptUserInputRef: promptUserInputRef, wrapperRef: wrapperRef }));
 }
 export default forwardRef(ExtensionAIControl);

package/build/components/message/index.d.ts

@@ -5,6 +5,7 @@ import './style.scss';
 /**
  * Types
  */
+import type { SuggestionErrorCode } from '../../types.js';
 import type React from 'react';
 export declare const MESSAGE_SEVERITY_WARNING = "warning";
 export declare const MESSAGE_SEVERITY_ERROR = "error";
@@ -19,13 +20,17 @@ export type MessageProps = {
     onSidebarIconClick?: () => void;
     children: React.ReactNode;
 };
+export type OnUpgradeClick = (event?: React.MouseEvent<HTMLButtonElement>) => void;
 export type UpgradeMessageProps = {
     requestsRemaining: number;
-    onUpgradeClick: () => void;
+    severity?: MessageSeverityProp;
+    onUpgradeClick: OnUpgradeClick;
 };
 export type ErrorMessageProps = {
     error?: string;
+    code?: SuggestionErrorCode;
     onTryAgainClick: () => void;
+    onUpgradeClick: OnUpgradeClick;
 };
 /**
  * React component to render a block message.
@@ -46,12 +51,12 @@ export declare function GuidelineMessage(): React.ReactElement;
  * @param {number} requestsRemaining - Number of requests remaining.
  * @returns {React.ReactElement } - Message component.
  */
-export declare function UpgradeMessage({ requestsRemaining, onUpgradeClick, }: UpgradeMessageProps): React.ReactElement;
+export declare function UpgradeMessage({ requestsRemaining, severity, onUpgradeClick, }: UpgradeMessageProps): React.ReactElement;
 /**
  * React component to render an error message
  *
  * @param {number} requestsRemaining - Number of requests remaining.
  * @returns {React.ReactElement } - Message component.
  */
-export declare function ErrorMessage({ error, onTryAgainClick }: ErrorMessageProps): React.ReactElement;
+export declare function ErrorMessage({ error, code, onTryAgainClick, onUpgradeClick, }: ErrorMessageProps): React.ReactElement;
 export {};

package/build/components/message/index.js

@@ -11,6 +11,7 @@ import classNames from 'classnames';
  */
 import './style.scss';
 import errorExclamation from '../../icons/error-exclamation.js';
+import { ERROR_QUOTA_EXCEEDED } from '../../types.js';
 export const MESSAGE_SEVERITY_WARNING = 'warning';
 export const MESSAGE_SEVERITY_ERROR = 'error';
 export const MESSAGE_SEVERITY_SUCCESS = 'success';
@@ -50,8 +51,12 @@ export function GuidelineMessage() {
  * @param {number} requestsRemaining - Number of requests remaining.
  * @returns {React.ReactElement } - Message component.
  */
-export function UpgradeMessage({ requestsRemaining, onUpgradeClick, }) {
-    return (_jsxs(Message, { severity: MESSAGE_SEVERITY_WARNING, children: [_jsx("span", { children: sprintf(
+export function UpgradeMessage({ requestsRemaining, severity, onUpgradeClick, }) {
+    let messageSeverity = severity;
+    if (messageSeverity == null) {
+        messageSeverity = requestsRemaining > 0 ? MESSAGE_SEVERITY_INFO : MESSAGE_SEVERITY_WARNING;
+    }
+    return (_jsxs(Message, { severity: messageSeverity, children: [_jsx("span", { children: sprintf(
                 // translators: %1$d: number of requests remaining
                 __('You have %1$d free requests remaining.', 'jetpack-ai-client'), requestsRemaining) }), _jsx(Button, { variant: "link", onClick: onUpgradeClick, children: __('Upgrade now', 'jetpack-ai-client') })] }));
 }
@@ -61,9 +66,9 @@ export function UpgradeMessage({ requestsRemaining, onUpgradeClick, }) {
  * @param {number} requestsRemaining - Number of requests remaining.
  * @returns {React.ReactElement } - Message component.
  */
-export function ErrorMessage({ error, onTryAgainClick }) {
+export function ErrorMessage({ error, code, onTryAgainClick, onUpgradeClick, }) {
     const errorMessage = error || __('Something went wrong', 'jetpack-ai-client');
     return (_jsxs(Message, { severity: MESSAGE_SEVERITY_ERROR, children: [_jsx("span", { children: sprintf(
                 // translators: %1$d: A dynamic error message
-                __('Error: %1$s.', 'jetpack-ai-client'), errorMessage) }), _jsx(Button, { variant: "link", onClick: onTryAgainClick, children: __('Try Again', 'jetpack-ai-client') })] }));
+                __('Error: %1$s', 'jetpack-ai-client'), errorMessage) }), code === ERROR_QUOTA_EXCEEDED ? (_jsx(Button, { variant: "link", onClick: onUpgradeClick, children: __('Upgrade now', 'jetpack-ai-client') })) : (_jsx(Button, { variant: "link", onClick: onTryAgainClick, children: __('Try again', 'jetpack-ai-client') }))] }));
 }

package/build/hooks/use-image-generator/index.d.ts

@@ -9,5 +9,14 @@ declare const useImageGenerator: () => {
             [key: string]: string;
         }[];
     }>;
+    generateImageWithStableDiffusion: ({ feature, postContent, userPrompt, }: {
+        feature: string;
+        postContent: string;
+        userPrompt?: string;
+    }) => Promise<{
+        data: {
+            [key: string]: string;
+        }[];
+    }>;
 };
 export default useImageGenerator;

package/build/hooks/use-image-generator/index.js

@@ -1,10 +1,12 @@
 /**
  * External dependencies
  */
+import { __ } from '@wordpress/i18n';
 import debugFactory from 'debug';
 /**
  * Internal dependencies
  */
+import askQuestionSync from '../../ask-question/sync.js';
 import requestJwt from '../../jwt/index.js';
 const debug = debugFactory('ai-client:use-image-generator');
 /**
@@ -27,7 +29,7 @@ const truncateContent = (content, currentPromptLength) => {
  * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated.
  * @returns {string} the prompt string
  */
-const getImageGenerationPrompt = (postContent, userPrompt) => {
+const getDalleImageGenerationPrompt = (postContent, userPrompt) => {
     /**
      * If the user provide some custom prompt for the image generation,
      * we will use it, add the post content as additional context and
@@ -73,7 +75,124 @@ This is the post content:
     // truncating the content so the whole prompt is not longer than 4000 characters, the model limit.
     return imageGenerationPrompt + truncateContent(postContent, imageGenerationPrompt.length);
 };
+/**
+ * Create the Stable Diffusion pre-processing prompt based on the provided context.
+ * @param {string} postContent - the content of the post.
+ * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated.
+ * @returns {string} the prompt string to be fed to the AI Assistant model.
+ */
+const getStableDiffusionPreProcessingPrompt = (postContent, userPrompt) => {
+    /**
+     * If the user provide some custom prompt for the image generation,
+     * we will use it and add the post content as additional context.
+     */
+    if (userPrompt) {
+        const preProcessingPrompt = `I need a Stable Diffusion prompt to generate a featured image for a blog post based on this user-provided image description:
+
+${userPrompt.length > 1000 ? userPrompt.substring(0, 1000) : userPrompt}
+
+The image should be a photo. Make sure you highlight the main suject of the image description, and include brief details about the light and style of the image.
+Include a request to use high resolution and produce a highly detailed image, with sharp focus.
+Return just the prompt, without comments.
+
+For additional context, this is the post content:
+
+`;
+        // truncating the content so the whole prompt is not longer than 4000 characters, the model limit.
+        return preProcessingPrompt + truncateContent(postContent, preProcessingPrompt.length);
+    }
+    /**
+     * When the user does not provide a custom prompt, we will use the
+     * standard one, based solely on the post content.
+     */
+    const preProcessingPrompt = `I need a Stable Diffusion prompt to generate a featured image for a blog post with the following content.
+The image should be a photo. Make sure you highlight the main suject of the content, and include brief details about the light and style of the image.
+Include a request to use high resolution and produce a highly detailed image, with sharp focus.
+Return just the prompt, without comments. The content is:
+
+`;
+    // truncating the content so the whole prompt is not longer than 4000 characters, the model limit.
+    return preProcessingPrompt + truncateContent(postContent, preProcessingPrompt.length);
+};
+/**
+ * Uses the Jetpack AI query endpoint to produce a prompt for the stable diffusion model.
+ * @param {string} postContent - the content of the post.
+ * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated
+ * @param {string} feature - the feature to be used for the image generation.
+ * @returns {string} the prompt string to be used on stable diffusion image generation.
+ */
+const getStableDiffusionImageGenerationPrompt = async (postContent, userPrompt, feature) => {
+    const prompt = getStableDiffusionPreProcessingPrompt(postContent, userPrompt);
+    /**
+     * Request the prompt on the AI Assistant endpoint
+     */
+    const data = await askQuestionSync(prompt, { feature });
+    return data.choices?.[0]?.message?.content;
+};
 const useImageGenerator = () => {
+    const generateImageWithStableDiffusion = async function ({ feature, postContent, userPrompt, }) {
+        let token = null;
+        try {
+            token = await requestJwt();
+        }
+        catch (error) {
+            debug('Error getting token: %o', error);
+            return Promise.reject(error);
+        }
+        try {
+            debug('Generating image with Stable Diffusion');
+            const prompt = await getStableDiffusionImageGenerationPrompt(postContent, userPrompt, feature);
+            const data = {
+                prompt,
+                style: 'photographic',
+                token: token.token,
+                width: 1024,
+                height: 768,
+            };
+            const response = await fetch(`https://public-api.wordpress.com/wpcom/v2/sites/${token.blogId}/ai-image`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify(data),
+            });
+            if (!response?.ok) {
+                debug('Error generating image: %o', response);
+                return Promise.reject({
+                    data: {
+                        status: response.status,
+                    },
+                    message: __('Error generating image. Please try again later.', 'jetpack-ai-client'),
+                });
+            }
+            const blob = await response.blob();
+            /**
+             * Convert the blob to base64 to keep the same format as the Dalle API.
+             */
+            const base64 = await new Promise((resolve, reject) => {
+                const reader = new FileReader();
+                reader.onloadend = () => {
+                    const base64data = reader.result;
+                    return resolve(base64data.replace(/^data:image\/(png|jpg);base64,/, ''));
+                };
+                reader.onerror = reject;
+                reader.readAsDataURL(blob);
+            });
+            // Return the Dalle API format
+            return {
+                data: [
+                    {
+                        b64_json: base64,
+                        revised_prompt: prompt,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            debug('Error generating image: %o', error);
+            return Promise.reject(error);
+        }
+    };
     const generateImage = async function ({ feature, postContent, responseFormat = 'url', userPrompt, }) {
         let token = '';
         try {
@@ -85,7 +204,7 @@ const useImageGenerator = () => {
         }
         try {
             debug('Generating image');
-            const imageGenerationPrompt = getImageGenerationPrompt(postContent, userPrompt);
+            const imageGenerationPrompt = getDalleImageGenerationPrompt(postContent, userPrompt);
             const URL = 'https://public-api.wordpress.com/wpcom/v2/jetpack-ai-image';
             const body = {
                 prompt: imageGenerationPrompt,
@@ -115,6 +234,7 @@ const useImageGenerator = () => {
     };
     return {
         generateImage,
+        generateImageWithStableDiffusion,
     };
 };
 export default useImageGenerator;

package/build/libs/index.d.ts

@@ -1 +1,2 @@
 export { MarkdownToHTML, HTMLToMarkdown, renderHTMLFromMarkdown, renderMarkdownFromHTML, } from './markdown/index.js';
+export type { RenderHTMLRules } from './markdown/index.js';

package/build/libs/markdown/index.d.ts

@@ -7,9 +7,10 @@ import MarkdownToHTML from './markdown-to-html.js';
  * Types
  */
 import type { Fix as HTMLFix } from './markdown-to-html.js';
+export type RenderHTMLRules = 'all' | Array<HTMLFix>;
 declare const renderHTMLFromMarkdown: ({ content, rules, }: {
     content: string;
-    rules?: Array<HTMLFix> | 'all';
+    rules?: RenderHTMLRules;
 }) => string;
 declare const renderMarkdownFromHTML: ({ content }: {
     content: string;

package/build/types.d.ts

@@ -39,3 +39,4 @@ export type Block = {
     originalContent?: string;
 };
 export type TranscriptionState = RecordingState | 'validating' | 'processing' | 'error';
+export type { RenderHTMLRules } from './libs/index.js';

package/package.json

@@ -1,7 +1,7 @@
 {
 	"private": false,
 	"name": "@automattic/jetpack-ai-client",
-	"version": "0.13.0",
+	"version": "0.14.0",
 	"description": "A JS client for consuming Jetpack AI services",
 	"homepage": "https://github.com/Automattic/jetpack/tree/HEAD/projects/js-packages/ai-client/#readme",
 	"bugs": {
@@ -42,24 +42,24 @@
 	"main": "./build/index.js",
 	"types": "./build/index.d.ts",
 	"dependencies": {
-		"@automattic/jetpack-base-styles": "^0.6.24",
-		"@automattic/jetpack-connection": "^0.33.9",
-		"@automattic/jetpack-shared-extension-utils": "^0.14.11",
+		"@automattic/jetpack-base-styles": "^0.6.25",
+		"@automattic/jetpack-connection": "^0.33.11",
+		"@automattic/jetpack-shared-extension-utils": "^0.14.13",
 		"@microsoft/fetch-event-source": "2.0.1",
 		"@types/react": "18.3.1",
-		"@wordpress/api-fetch": "6.53.0",
-		"@wordpress/block-editor": "12.24.0",
-		"@wordpress/components": "27.4.0",
-		"@wordpress/compose": "6.33.0",
-		"@wordpress/data": "9.26.0",
-		"@wordpress/element": "5.33.0",
-		"@wordpress/i18n": "4.56.0",
-		"@wordpress/icons": "9.47.0",
+		"@wordpress/api-fetch": "6.54.0",
+		"@wordpress/block-editor": "12.25.0",
+		"@wordpress/components": "27.5.0",
+		"@wordpress/compose": "6.34.0",
+		"@wordpress/data": "9.27.0",
+		"@wordpress/element": "5.34.0",
+		"@wordpress/i18n": "4.57.0",
+		"@wordpress/icons": "9.48.0",
 		"classnames": "2.3.2",
 		"debug": "4.3.4",
 		"markdown-it": "14.0.0",
-		"react": "18.2.0",
-		"react-dom": "18.2.0",
+		"react": "18.3.1",
+		"react-dom": "18.3.1",
 		"turndown": "7.1.2"
 	}
 }

package/src/ask-question/sync.ts

@@ -0,0 +1,91 @@
+/**
+ * External dependencies
+ */
+import debugFactory from 'debug';
+/*
+ * Types & constants
+ */
+import requestJwt from '../jwt/index.js';
+import { AskQuestionOptionsArgProps } from './index.js';
+import type { PromptProp } from '../types.js';
+
+/**
+ * The response data from the AI assistant when doing a sync, not-streamed question.
+ */
+export type ResponseData = {
+	choices: Array< {
+		message: {
+			content: string;
+		};
+	} >;
+};
+
+const debug = debugFactory( 'jetpack-ai-client:ask-question-sync' );
+
+/**
+ * A function that asks a question without streaming.
+ *
+ * @param {PromptProp} question - The question to ask. It can be a simple string or an array of PromptMessageItemProps objects.
+ * @param {AskQuestionOptionsArgProps} options - An optional object for additional configuration: postId, feature, model.
+ * @returns {Promise<ResponseData>} - A promise that resolves to an instance of the ResponseData
+ * @example
+ * const question = "What is the meaning of life?";
+ * const options = {
+ *   feature: 'ai-featured-image',
+ *   model: 'gpt-4-turbo'
+ * }
+ * askQuestionSync( question, options ).then( responseData => {
+ *    // access the choices array on the response data
+ *    const content = responseData.choices[ 0 ].message.content;
+ * } );
+ */
+export default async function askQuestionSync(
+	question: PromptProp,
+	{ postId = null, feature, model }: AskQuestionOptionsArgProps = {}
+): Promise< ResponseData > {
+	debug( 'Asking question with no streaming: %o. options: %o', question, {
+		postId,
+		feature,
+		model,
+	} );
+
+	/**
+	 * The URL to the AI assistant query endpoint.
+	 */
+	const URL = 'https://public-api.wordpress.com/wpcom/v2/jetpack-ai-query';
+
+	let token = null;
+
+	try {
+		token = ( await requestJwt() ).token;
+	} catch ( error ) {
+		debug( 'Error getting token: %o', error );
+		return Promise.reject( error );
+	}
+
+	const body = {
+		question: question,
+		stream: false,
+		postId,
+		feature,
+		model,
+	};
+
+	const headers = {
+		Authorization: `Bearer ${ token }`,
+		'Content-Type': 'application/json',
+	};
+
+	const data = await fetch( URL, {
+		method: 'POST',
+		headers,
+		body: JSON.stringify( body ),
+	} ).then( response => response.json() );
+
+	if ( data?.data?.status && data?.data?.status > 200 ) {
+		debug( 'Error generating prompt: %o', data );
+		return Promise.reject( data );
+	}
+
+	return data as ResponseData;
+}

package/src/components/ai-control/ai-control.tsx

@@ -16,6 +16,7 @@ import type { RequestingStateProp } from '../../types.js';
 import type { ReactElement } from 'react';
 
 type AIControlProps = {
+	className?: string;
 	disabled?: boolean;
 	value: string;
 	placeholder?: string;
@@ -37,6 +38,7 @@ type AIControlProps = {
  * @returns {ReactElement}       Rendered component
  */
 export default function AIControl( {
+	className,
 	disabled = false,
 	value = '',
 	placeholder = '',
@@ -51,7 +53,10 @@ export default function AIControl( {
 	wrapperRef = null,
 }: AIControlProps ): ReactElement {
 	return (
-		<div className="jetpack-components-ai-control__container-wrapper" ref={ wrapperRef }>
+		<div
+			className={ classNames( 'jetpack-components-ai-control__container-wrapper', className ) }
+			ref={ wrapperRef }
+		>
 			{ error }
 			<div className="jetpack-components-ai-control__container">
 				{ banner }

package/src/components/ai-control/extension-ai-control.tsx

@@ -16,10 +16,11 @@ import './style.scss';
 /**
  * Types
  */
-import type { RequestingStateProp } from '../../types.js';
-import type { ReactElement } from 'react';
+import type { RequestingErrorProps, RequestingStateProp } from '../../types.js';
+import type { ReactElement, MouseEvent } from 'react';
 
 type ExtensionAIControlProps = {
+	className?: string;
 	disabled?: boolean;
 	value: string;
 	placeholder?: string;
@@ -27,7 +28,7 @@ type ExtensionAIControlProps = {
 	isTransparent?: boolean;
 	state?: RequestingStateProp;
 	showGuideLine?: boolean;
-	error?: string;
+	error?: RequestingErrorProps;
 	requestsRemaining?: number;
 	showUpgradeMessage?: boolean;
 	wrapperRef?: React.MutableRefObject< HTMLDivElement | null >;
@@ -36,7 +37,8 @@ type ExtensionAIControlProps = {
 	onStop?: () => void;
 	onClose?: () => void;
 	onUndo?: () => void;
-	onUpgrade?: () => void;
+	onUpgrade?: ( event: MouseEvent< HTMLButtonElement > ) => void;
+	onTryAgain?: () => void;
 };
 
 /**
@@ -48,6 +50,7 @@ type ExtensionAIControlProps = {
  */
 export function ExtensionAIControl(
 	{
+		className,
 		disabled = false,
 		value = '',
 		placeholder = '',
@@ -65,6 +68,7 @@ export function ExtensionAIControl(
 		onClose,
 		onUndo,
 		onUpgrade,
+		onTryAgain,
 	}: ExtensionAIControlProps,
 	ref: React.MutableRefObject< HTMLInputElement >
 ): ReactElement {
@@ -118,9 +122,16 @@ export function ExtensionAIControl(
 		onUndo?.();
 	}, [ onUndo ] );
 
-	const upgradeHandler = useCallback( () => {
-		onUpgrade?.();
-	}, [ onUpgrade ] );
+	const upgradeHandler = useCallback(
+		( event: MouseEvent< HTMLButtonElement > ) => {
+			onUpgrade?.( event );
+		},
+		[ onUpgrade ]
+	);
+
+	const tryAgainHandler = useCallback( () => {
+		onTryAgain?.();
+	}, [ onTryAgain ] );
 
 	useKeyboardShortcut(
 		'enter',
@@ -191,8 +202,16 @@ export function ExtensionAIControl(
 	);
 
 	let message = null;
-	if ( error ) {
-		message = <ErrorMessage error={ error } onTryAgainClick={ sendHandler } />;
+
+	if ( error?.message ) {
+		message = (
+			<ErrorMessage
+				error={ error.message }
+				code={ error.code }
+				onTryAgainClick={ tryAgainHandler }
+				onUpgradeClick={ upgradeHandler }
+			/>
+		);
 	} else if ( showUpgradeMessage ) {
 		message = (
 			<UpgradeMessage requestsRemaining={ requestsRemaining } onUpgradeClick={ upgradeHandler } />
@@ -203,6 +222,7 @@ export function ExtensionAIControl(
 
 	return (
 		<AIControl
+			className={ className }
 			disabled={ disabled || loading }
 			value={ value }
 			placeholder={ placeholder }

package/src/components/message/index.tsx

@@ -10,9 +10,11 @@ import classNames from 'classnames';
  */
 import './style.scss';
 import errorExclamation from '../../icons/error-exclamation.js';
+import { ERROR_QUOTA_EXCEEDED } from '../../types.js';
 /**
  * Types
  */
+import type { SuggestionErrorCode } from '../../types.js';
 import type React from 'react';
 
 export const MESSAGE_SEVERITY_WARNING = 'warning';
@@ -37,14 +39,19 @@ export type MessageProps = {
 	children: React.ReactNode;
 };
 
+export type OnUpgradeClick = ( event?: React.MouseEvent< HTMLButtonElement > ) => void;
+
 export type UpgradeMessageProps = {
 	requestsRemaining: number;
-	onUpgradeClick: () => void;
+	severity?: MessageSeverityProp;
+	onUpgradeClick: OnUpgradeClick;
 };
 
 export type ErrorMessageProps = {
 	error?: string;
+	code?: SuggestionErrorCode;
 	onTryAgainClick: () => void;
+	onUpgradeClick: OnUpgradeClick;
 };
 
 const messageIconsMap = {
@@ -113,10 +120,17 @@ export function GuidelineMessage(): React.ReactElement {
  */
 export function UpgradeMessage( {
 	requestsRemaining,
+	severity,
 	onUpgradeClick,
 }: UpgradeMessageProps ): React.ReactElement {
+	let messageSeverity = severity;
+
+	if ( messageSeverity == null ) {
+		messageSeverity = requestsRemaining > 0 ? MESSAGE_SEVERITY_INFO : MESSAGE_SEVERITY_WARNING;
+	}
+
 	return (
-		<Message severity={ MESSAGE_SEVERITY_WARNING }>
+		<Message severity={ messageSeverity }>
 			<span>
 				{ sprintf(
 					// translators: %1$d: number of requests remaining
@@ -137,7 +151,12 @@ export function UpgradeMessage( {
  * @param {number} requestsRemaining - Number of requests remaining.
  * @returns {React.ReactElement } - Message component.
  */
-export function ErrorMessage( { error, onTryAgainClick }: ErrorMessageProps ): React.ReactElement {
+export function ErrorMessage( {
+	error,
+	code,
+	onTryAgainClick,
+	onUpgradeClick,
+}: ErrorMessageProps ): React.ReactElement {
 	const errorMessage = error || __( 'Something went wrong', 'jetpack-ai-client' );
 
 	return (
@@ -145,13 +164,19 @@ export function ErrorMessage( { error, onTryAgainClick }: ErrorMessageProps ): R
 			<span>
 				{ sprintf(
 					// translators: %1$d: A dynamic error message
-					__( 'Error: %1$s.', 'jetpack-ai-client' ),
+					__( 'Error: %1$s', 'jetpack-ai-client' ),
 					errorMessage
 				) }
 			</span>
-			<Button variant="link" onClick={ onTryAgainClick }>
-				{ __( 'Try Again', 'jetpack-ai-client' ) }
-			</Button>
+			{ code === ERROR_QUOTA_EXCEEDED ? (
+				<Button variant="link" onClick={ onUpgradeClick }>
+					{ __( 'Upgrade now', 'jetpack-ai-client' ) }
+				</Button>
+			) : (
+				<Button variant="link" onClick={ onTryAgainClick }>
+					{ __( 'Try again', 'jetpack-ai-client' ) }
+				</Button>
+			) }
 		</Message>
 	);
 }

package/src/hooks/use-image-generator/index.ts

@@ -1,10 +1,12 @@
 /**
  * External dependencies
  */
+import { __ } from '@wordpress/i18n';
 import debugFactory from 'debug';
 /**
  * Internal dependencies
  */
+import askQuestionSync from '../../ask-question/sync.js';
 import requestJwt from '../../jwt/index.js';
 
 const debug = debugFactory( 'ai-client:use-image-generator' );
@@ -30,7 +32,7 @@ const truncateContent = ( content: string, currentPromptLength: number ): string
  * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated.
  * @returns {string} the prompt string
  */
-const getImageGenerationPrompt = ( postContent: string, userPrompt?: string ): string => {
+const getDalleImageGenerationPrompt = ( postContent: string, userPrompt?: string ): string => {
 	/**
 	 * If the user provide some custom prompt for the image generation,
 	 * we will use it, add the post content as additional context and
@@ -78,7 +80,160 @@ This is the post content:
 	return imageGenerationPrompt + truncateContent( postContent, imageGenerationPrompt.length );
 };
 
+/**
+ * Create the Stable Diffusion pre-processing prompt based on the provided context.
+ * @param {string} postContent - the content of the post.
+ * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated.
+ * @returns {string} the prompt string to be fed to the AI Assistant model.
+ */
+const getStableDiffusionPreProcessingPrompt = (
+	postContent: string,
+	userPrompt?: string
+): string => {
+	/**
+	 * If the user provide some custom prompt for the image generation,
+	 * we will use it and add the post content as additional context.
+	 */
+	if ( userPrompt ) {
+		const preProcessingPrompt = `I need a Stable Diffusion prompt to generate a featured image for a blog post based on this user-provided image description:
+
+${ userPrompt.length > 1000 ? userPrompt.substring( 0, 1000 ) : userPrompt }
+
+The image should be a photo. Make sure you highlight the main suject of the image description, and include brief details about the light and style of the image.
+Include a request to use high resolution and produce a highly detailed image, with sharp focus.
+Return just the prompt, without comments.
+
+For additional context, this is the post content:
+
+`;
+		// truncating the content so the whole prompt is not longer than 4000 characters, the model limit.
+		return preProcessingPrompt + truncateContent( postContent, preProcessingPrompt.length );
+	}
+
+	/**
+	 * When the user does not provide a custom prompt, we will use the
+	 * standard one, based solely on the post content.
+	 */
+	const preProcessingPrompt = `I need a Stable Diffusion prompt to generate a featured image for a blog post with the following content.
+The image should be a photo. Make sure you highlight the main suject of the content, and include brief details about the light and style of the image.
+Include a request to use high resolution and produce a highly detailed image, with sharp focus.
+Return just the prompt, without comments. The content is:
+
+`;
+
+	// truncating the content so the whole prompt is not longer than 4000 characters, the model limit.
+	return preProcessingPrompt + truncateContent( postContent, preProcessingPrompt.length );
+};
+
+/**
+ * Uses the Jetpack AI query endpoint to produce a prompt for the stable diffusion model.
+ * @param {string} postContent - the content of the post.
+ * @param {string} userPrompt - the user prompt for the image generation, if provided. Max length is 1000 characters, will be truncated
+ * @param {string} feature - the feature to be used for the image generation.
+ * @returns {string} the prompt string to be used on stable diffusion image generation.
+ */
+const getStableDiffusionImageGenerationPrompt = async (
+	postContent: string,
+	userPrompt?: string,
+	feature?: string
+): Promise< string > => {
+	const prompt = getStableDiffusionPreProcessingPrompt( postContent, userPrompt );
+
+	/**
+	 * Request the prompt on the AI Assistant endpoint
+	 */
+	const data = await askQuestionSync( prompt, { feature } );
+
+	return data.choices?.[ 0 ]?.message?.content;
+};
+
 const useImageGenerator = () => {
+	const generateImageWithStableDiffusion = async function ( {
+		feature,
+		postContent,
+		userPrompt,
+	}: {
+		feature: string;
+		postContent: string;
+		userPrompt?: string;
+	} ): Promise< { data: Array< { [ key: string ]: string } > } > {
+		let token = null;
+
+		try {
+			token = await requestJwt();
+		} catch ( error ) {
+			debug( 'Error getting token: %o', error );
+			return Promise.reject( error );
+		}
+
+		try {
+			debug( 'Generating image with Stable Diffusion' );
+
+			const prompt = await getStableDiffusionImageGenerationPrompt(
+				postContent,
+				userPrompt,
+				feature
+			);
+
+			const data = {
+				prompt,
+				style: 'photographic',
+				token: token.token,
+				width: 1024,
+				height: 768,
+			};
+
+			const response = await fetch(
+				`https://public-api.wordpress.com/wpcom/v2/sites/${ token.blogId }/ai-image`,
+				{
+					method: 'POST',
+					headers: {
+						'Content-Type': 'application/json',
+					},
+					body: JSON.stringify( data ),
+				}
+			);
+
+			if ( ! response?.ok ) {
+				debug( 'Error generating image: %o', response );
+				return Promise.reject( {
+					data: {
+						status: response.status,
+					},
+					message: __( 'Error generating image. Please try again later.', 'jetpack-ai-client' ),
+				} );
+			}
+
+			const blob = await response.blob();
+
+			/**
+			 * Convert the blob to base64 to keep the same format as the Dalle API.
+			 */
+			const base64 = await new Promise( ( resolve, reject ) => {
+				const reader = new FileReader();
+				reader.onloadend = () => {
+					const base64data = reader.result as string;
+					return resolve( base64data.replace( /^data:image\/(png|jpg);base64,/, '' ) );
+				};
+				reader.onerror = reject;
+				reader.readAsDataURL( blob );
+			} );
+
+			// Return the Dalle API format
+			return {
+				data: [
+					{
+						b64_json: base64 as string,
+						revised_prompt: prompt,
+					},
+				],
+			};
+		} catch ( error ) {
+			debug( 'Error generating image: %o', error );
+			return Promise.reject( error );
+		}
+	};
+
 	const generateImage = async function ( {
 		feature,
 		postContent,
@@ -102,7 +257,7 @@ const useImageGenerator = () => {
 		try {
 			debug( 'Generating image' );
 
-			const imageGenerationPrompt = getImageGenerationPrompt( postContent, userPrompt );
+			const imageGenerationPrompt = getDalleImageGenerationPrompt( postContent, userPrompt );
 
 			const URL = 'https://public-api.wordpress.com/wpcom/v2/jetpack-ai-image';
 
@@ -138,6 +293,7 @@ const useImageGenerator = () => {
 
 	return {
 		generateImage,
+		generateImageWithStableDiffusion,
 	};
 };
 

package/src/libs/index.ts

@@ -4,3 +4,5 @@ export {
 	renderHTMLFromMarkdown,
 	renderMarkdownFromHTML,
 } from './markdown/index.js';
+
+export type { RenderHTMLRules } from './markdown/index.js';

package/src/libs/markdown/index.ts

@@ -11,12 +11,14 @@ import type { Fix as HTMLFix } from './markdown-to-html.js';
 const defaultMarkdownConverter = new MarkdownToHTML();
 const defaultHTMLConverter = new HTMLToMarkdown();
 
+export type RenderHTMLRules = 'all' | Array< HTMLFix >;
+
 const renderHTMLFromMarkdown = ( {
 	content,
 	rules = 'all',
 }: {
 	content: string;
-	rules?: Array< HTMLFix > | 'all';
+	rules?: RenderHTMLRules;
 } ) => {
 	return defaultMarkdownConverter.render( { content, rules } );
 };

package/src/types.ts

@@ -108,3 +108,8 @@ export type Block = {
  * Transcription types
  */
 export type TranscriptionState = RecordingState | 'validating' | 'processing' | 'error';
+
+/*
+ * Lib types
+ */
+export type { RenderHTMLRules } from './libs/index.js';