AI Client: add style guessing and prompt processing (#39712)

* add type definitions as required by linter * add guessStyle on use-logo-generator hook * use guessStyle on logo generator prompt * changelog Committed via a GitHub action: https://github.com/Automattic/jetpack/actions/runs/11293480841 Upstream-Ref: Automattic/jetpack@8362bde
Automattic · Oct 11, 2024 · 4b58174 · 4b58174
1 parent a2d51ee
commit 4b58174
Show file tree

Hide file tree

Showing 12 changed files with 205 additions and 39 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ This is an alpha version! The changes listed here are not final.
 
 ### Added
 - AI Client - Logo generator: add image styles 'auto' and 'none'. Order styles so those are on top in the dropdown selector
+- AI Client: add prompt processing and style guess function for logo generator
 
 ### Changed
 - AI Client: change plans limit to use and accept new 3000 value

diff --git a/build/ai-client/src/logo-generator/components/prompt.js b/build/ai-client/src/logo-generator/components/prompt.js
@@ -13,7 +13,7 @@ import { useCallback, useEffect, useState, useRef } from 'react';
  */
 import { IMAGE_STYLE_NONE, IMAGE_STYLE_AUTO } from '../../hooks/use-image-generator/constants.js';
 import AiIcon from '../assets/icons/ai.js';
-import { EVENT_GENERATE, MINIMUM_PROMPT_LENGTH, EVENT_UPGRADE, EVENT_PLACEMENT_INPUT_FOOTER, EVENT_SWITCH_STYLE, } from '../constants.js';
+import { EVENT_GENERATE, MINIMUM_PROMPT_LENGTH, EVENT_UPGRADE, EVENT_PLACEMENT_INPUT_FOOTER, EVENT_SWITCH_STYLE, EVENT_GUESS_STYLE, } from '../constants.js';
 import { useCheckout } from '../hooks/use-checkout.js';
 import useLogoGenerator from '../hooks/use-logo-generator.js';
 import useRequestErrors from '../hooks/use-request-errors.js';
@@ -32,7 +32,7 @@ export const Prompt = ({ initialPrompt = '' }) => {
     const [showStyleSelector, setShowStyleSelector] = useState(false);
     const [style, setStyle] = useState(null);
     const [styles, setStyles] = useState([]);
-    const { generateLogo, enhancePrompt, setIsEnhancingPrompt, isBusy, isEnhancingPrompt, site, getAiAssistantFeature, requireUpgrade, context, tierPlansEnabled, imageStyles, } = useLogoGenerator();
+    const { generateLogo, enhancePrompt, setIsEnhancingPrompt, isBusy, isEnhancingPrompt, site, getAiAssistantFeature, requireUpgrade, context, tierPlansEnabled, imageStyles, guessStyle, } = useLogoGenerator();
     const enhancingLabel = __('Enhancing…', 'jetpack-ai-client');
     const enhanceLabel = __('Enhance prompt', 'jetpack-ai-client');
     const enhanceButtonLabel = isEnhancingPrompt ? enhancingLabel : enhanceLabel;
@@ -88,9 +88,20 @@ export const Prompt = ({ initialPrompt = '' }) => {
         }
     }, [imageStyles]);
     const onGenerate = useCallback(async () => {
-        // shouldn't tool be "logo-generator" to be more specific?
-        recordTracksEvent(EVENT_GENERATE, { context, tool: 'image', style });
-        generateLogo({ prompt, style });
+        debug(context);
+        if (style === IMAGE_STYLE_AUTO) {
+            setIsEnhancingPrompt(true);
+            recordTracksEvent(EVENT_GUESS_STYLE, { context, tool: 'image' });
+            const guessedStyle = (await guessStyle(prompt)) || IMAGE_STYLE_NONE;
+            setStyle(guessedStyle);
+            recordTracksEvent(EVENT_GENERATE, { context, tool: 'image', style: guessedStyle });
+            setIsEnhancingPrompt(false);
+            generateLogo({ prompt, style: guessedStyle });
+        }
+        else {
+            recordTracksEvent(EVENT_GENERATE, { context, tool: 'image', style });
+            generateLogo({ prompt, style });
+        }
     }, [context, generateLogo, prompt, style]);
     const onPromptInput = (event) => {
         setPrompt(event.target.textContent || '');
@@ -123,7 +134,7 @@ export const Prompt = ({ initialPrompt = '' }) => {
             onGenerate();
         }
     };
-    return (_jsxs("div", { className: "jetpack-ai-logo-generator__prompt", children: [_jsxs("div", { className: "jetpack-ai-logo-generator__prompt-header", children: [_jsx("div", { className: "jetpack-ai-logo-generator__prompt-label", children: __('Describe your site:', 'jetpack-ai-client') }), _jsx("div", { className: "jetpack-ai-logo-generator__prompt-actions", children: _jsxs(Button, { variant: "link", disabled: isBusy || requireUpgrade || !hasPrompt, onClick: onEnhance, children: [_jsx(AiIcon, {}), enhanceButtonLabel] }) }), showStyleSelector && (_jsx(SelectControl, { __nextHasNoMarginBottom: true, value: style, options: styles, onChange: updateStyle }))] }), _jsxs("div", { className: "jetpack-ai-logo-generator__prompt-query", children: [_jsx("div", { role: "textbox", tabIndex: 0, ref: inputRef, contentEditable: !isBusy && !requireUpgrade, 
+    return (_jsxs("div", { className: "jetpack-ai-logo-generator__prompt", children: [_jsxs("div", { className: "jetpack-ai-logo-generator__prompt-header", children: [_jsx("div", { className: "jetpack-ai-logo-generator__prompt-label", children: __('Describe your site:', 'jetpack-ai-client') }), _jsx("div", { className: "jetpack-ai-logo-generator__prompt-actions", children: _jsxs(Button, { variant: "link", disabled: isBusy || requireUpgrade || !hasPrompt, onClick: onEnhance, children: [_jsx(AiIcon, {}), enhanceButtonLabel] }) }), showStyleSelector && (_jsx(SelectControl, { __nextHasNoMarginBottom: true, value: style, options: styles, onChange: updateStyle, disabled: isBusy || requireUpgrade }))] }), _jsxs("div", { className: "jetpack-ai-logo-generator__prompt-query", children: [_jsx("div", { role: "textbox", tabIndex: 0, ref: inputRef, contentEditable: !isBusy && !requireUpgrade, 
                         // The content editable div is expected to be updated by the enhance prompt, so warnings are suppressed
                         suppressContentEditableWarning: true, className: "prompt-query__input", onInput: onPromptInput, onPaste: onPromptPaste, onKeyDown: onKeyDown, "data-placeholder": __('Describe your site or simply ask for a logo specifying some details about it', 'jetpack-ai-client') }), _jsx(Button, { variant: "primary", className: "jetpack-ai-logo-generator__prompt-submit", onClick: onGenerate, disabled: isBusy || requireUpgrade || !hasPrompt, children: __('Generate', 'jetpack-ai-client') })] }), _jsxs("div", { className: "jetpack-ai-logo-generator__prompt-footer", children: [!isUnlimited && !requireUpgrade && (_jsxs("div", { className: "jetpack-ai-logo-generator__prompt-requests", children: [_jsx("div", { children: sprintf(
                                 // translators: %u is the number of requests

diff --git a/build/ai-client/src/logo-generator/constants.d.ts b/build/ai-client/src/logo-generator/constants.d.ts
@@ -9,6 +9,7 @@ export declare const EVENT_NAVIGATE = "jetpack_ai_logo_generator_navigate";
 export declare const EVENT_FEEDBACK = "jetpack_ai_logo_generator_feedback";
 export declare const EVENT_UPGRADE = "jetpack_ai_upgrade_button";
 export declare const EVENT_SWITCH_STYLE = "jetpack_ai_logo_generator_switch_style";
+export declare const EVENT_GUESS_STYLE = "jetpack_ai_logo_generator_guess_style";
 export declare const EVENT_PLACEMENT_QUICK_LINKS = "quick_links";
 export declare const EVENT_PLACEMENT_INPUT_FOOTER = "input_footer";
 export declare const EVENT_PLACEMENT_FREE_USER_SCREEN = "free_user_screen";

diff --git a/build/ai-client/src/logo-generator/constants.js b/build/ai-client/src/logo-generator/constants.js
@@ -10,6 +10,7 @@ export const EVENT_NAVIGATE = 'jetpack_ai_logo_generator_navigate';
 export const EVENT_FEEDBACK = 'jetpack_ai_logo_generator_feedback';
 export const EVENT_UPGRADE = 'jetpack_ai_upgrade_button';
 export const EVENT_SWITCH_STYLE = 'jetpack_ai_logo_generator_switch_style';
+export const EVENT_GUESS_STYLE = 'jetpack_ai_logo_generator_guess_style';
 // Event placement constants
 export const EVENT_PLACEMENT_QUICK_LINKS = 'quick_links';
 export const EVENT_PLACEMENT_INPUT_FOOTER = 'input_footer';

diff --git a/build/ai-client/src/logo-generator/hooks/use-logo-generator.d.ts b/build/ai-client/src/logo-generator/hooks/use-logo-generator.d.ts
@@ -49,5 +49,6 @@ declare const useLogoGenerator: () => {
     isLoadingHistory: boolean;
     setIsLoadingHistory: any;
     imageStyles: ImageStyleObject[];
+    guessStyle: (prompt: string) => Promise<ImageStyle | null>;
 };
 export default useLogoGenerator;
diff --git a/build/ai-client/src/logo-generator/hooks/use-logo-generator.js b/build/ai-client/src/logo-generator/hooks/use-logo-generator.js
@@ -7,6 +7,7 @@ import { useCallback } from 'react';
 /**
  * Internal dependencies
  */
+import askQuestionSync from '../../ask-question/sync.js';
 import useImageGenerator from '../../hooks/use-image-generator/index.js';
 import useSaveToMediaLibrary from '../../hooks/use-save-to-media-library/index.js';
 import requestJwt from '../../jwt/index.js';
@@ -129,13 +130,48 @@ For example: user's prompt: A logo for an ice cream shop. Returned prompt: A log
             throw error;
         }
     };
+    const guessStyle = useCallback(async function (prompt) {
+        setLogoFetchError(null);
+        if (!imageStyles || !imageStyles.length) {
+            return null;
+        }
+        const messages = [
+            {
+                role: 'jetpack-ai',
+                context: {
+                    type: 'ai-assistant-guess-logo-style',
+                    request: prompt,
+                    name,
+                    description,
+                },
+            },
+        ];
+        try {
+            const style = await askQuestionSync(messages, { feature: 'jetpack-ai-logo-generator' });
+            if (!style) {
+                return null;
+            }
+            const styleObject = imageStyles.find(({ value }) => value === style);
+            if (!styleObject) {
+                return null;
+            }
+            return styleObject.value;
+        }
+        catch (error) {
+            debug('Error guessing style', error);
+            Promise.reject(error);
+        }
+    }, [imageStyles, name, description]);
     const generateImage = useCallback(async function ({ prompt, style = null, }) {
         setLogoFetchError(null);
         try {
             const tokenData = await requestJwt();
             if (!tokenData || !tokenData.token) {
                 throw new Error('No token provided');
             }
+            if (style === 'auto') {
+                throw new Error('Auto style is not supported');
+            }
             debug('Generating image with prompt', prompt);
             const imageGenerationPrompt = `I NEED to test how the tool works with extremely simple prompts. DO NOT add any detail, just use it AS-IS:
 Create a single text-free iconic vector logo that symbolically represents the user request, using abstract or symbolic imagery.
@@ -148,6 +184,21 @@ The image should contain a single icon, without variations, color palettes or di
 User request:${prompt}`;
             const body = {
                 prompt: imageGenerationPrompt,
+                // if style is set prompt is reworked at backend with messages
+                messages: style
+                    ? [
+                        {
+                            role: 'jetpack-ai',
+                            context: {
+                                type: 'ai-assistant-generate-logo',
+                                request: prompt,
+                                name,
+                                description,
+                                style,
+                            },
+                        },
+                    ]
+                    : [],
                 feature: 'jetpack-ai-logo-generator',
                 response_format: 'b64_json',
                 style: style || '', // backend expects an empty string if no style is provided
@@ -159,7 +210,7 @@ User request:${prompt}`;
             setLogoFetchError(error);
             throw error;
         }
-    }, []);
+    }, [name, description]);
     const saveLogo = useCallback(async (logo) => {
         setSaveToLibraryError(null);
         try {
@@ -290,6 +341,7 @@ User request:${prompt}`;
         isLoadingHistory,
         setIsLoadingHistory,
         imageStyles,
+        guessStyle,
     };
 };
 export default useLogoGenerator;
diff --git a/build/ai-client/src/types.d.ts b/build/ai-client/src/types.d.ts
@@ -7,8 +7,13 @@ export declare const ERROR_NETWORK: "error_network";
 export declare const ERROR_UNCLEAR_PROMPT: "error_unclear_prompt";
 export declare const ERROR_RESPONSE: "error_response";
 export type SuggestionErrorCode = typeof ERROR_SERVICE_UNAVAILABLE | typeof ERROR_QUOTA_EXCEEDED | typeof ERROR_MODERATION | typeof ERROR_CONTEXT_TOO_LARGE | typeof ERROR_NETWORK | typeof ERROR_UNCLEAR_PROMPT | typeof ERROR_RESPONSE;
+export declare const ROLE_SYSTEM: "system";
+export declare const ROLE_USER: "user";
+export declare const ROLE_ASSISTANT: "assistant";
+export declare const ROLE_JETPACK_AI: "jetpack-ai";
+export type RoleType = typeof ROLE_SYSTEM | typeof ROLE_USER | typeof ROLE_ASSISTANT | typeof ROLE_JETPACK_AI;
 export type PromptItemProps = {
-    role: 'system' | 'user' | 'assistant' | 'jetpack-ai';
+    role: RoleType;
     content?: string;
     context?: object;
 };

diff --git a/build/ai-client/src/types.js b/build/ai-client/src/types.js
@@ -5,6 +5,10 @@ export const ERROR_CONTEXT_TOO_LARGE = 'error_context_too_large';
 export const ERROR_NETWORK = 'error_network';
 export const ERROR_UNCLEAR_PROMPT = 'error_unclear_prompt';
 export const ERROR_RESPONSE = 'error_response';
+export const ROLE_SYSTEM = 'system';
+export const ROLE_USER = 'user';
+export const ROLE_ASSISTANT = 'assistant';
+export const ROLE_JETPACK_AI = 'jetpack-ai';
 /*
  * Hook constants
  */

diff --git a/src/logo-generator/components/prompt.tsx b/src/logo-generator/components/prompt.tsx
@@ -18,6 +18,7 @@ import {
 	EVENT_UPGRADE,
 	EVENT_PLACEMENT_INPUT_FOOTER,
 	EVENT_SWITCH_STYLE,
+	EVENT_GUESS_STYLE,
 } from '../constants.js';
 import { useCheckout } from '../hooks/use-checkout.js';
 import useLogoGenerator from '../hooks/use-logo-generator.js';
@@ -60,6 +61,7 @@ export const Prompt = ( { initialPrompt = '' }: PromptProps ) => {
 		context,
 		tierPlansEnabled,
 		imageStyles,
+		guessStyle,
 	} = useLogoGenerator();
 
 	const enhancingLabel = __( 'Enhancing…', 'jetpack-ai-client' );
@@ -126,9 +128,19 @@ export const Prompt = ( { initialPrompt = '' }: PromptProps ) => {
 	}, [ imageStyles ] );
 
 	const onGenerate = useCallback( async () => {
-		// shouldn't tool be "logo-generator" to be more specific?
-		recordTracksEvent( EVENT_GENERATE, { context, tool: 'image', style } );
-		generateLogo( { prompt, style } );
+		debug( context );
+		if ( style === IMAGE_STYLE_AUTO ) {
+			setIsEnhancingPrompt( true );
+			recordTracksEvent( EVENT_GUESS_STYLE, { context, tool: 'image' } );
+			const guessedStyle = ( await guessStyle( prompt ) ) || IMAGE_STYLE_NONE;
+			setStyle( guessedStyle );
+			recordTracksEvent( EVENT_GENERATE, { context, tool: 'image', style: guessedStyle } );
+			setIsEnhancingPrompt( false );
+			generateLogo( { prompt, style: guessedStyle } );
+		} else {
+			recordTracksEvent( EVENT_GENERATE, { context, tool: 'image', style } );
+			generateLogo( { prompt, style } );
+		}
 	}, [ context, generateLogo, prompt, style ] );
 
 	const onPromptInput = ( event: React.ChangeEvent< HTMLInputElement > ) => {
@@ -196,6 +208,7 @@ export const Prompt = ( { initialPrompt = '' }: PromptProps ) => {
 						value={ style }
 						options={ styles }
 						onChange={ updateStyle }
+						disabled={ isBusy || requireUpgrade }
 					/>
 				) }
 			</div>

diff --git a/src/logo-generator/constants.ts b/src/logo-generator/constants.ts
@@ -11,6 +11,7 @@ export const EVENT_NAVIGATE = 'jetpack_ai_logo_generator_navigate';
 export const EVENT_FEEDBACK = 'jetpack_ai_logo_generator_feedback';
 export const EVENT_UPGRADE = 'jetpack_ai_upgrade_button';
 export const EVENT_SWITCH_STYLE = 'jetpack_ai_logo_generator_switch_style';
+export const EVENT_GUESS_STYLE = 'jetpack_ai_logo_generator_guess_style';
 
 // Event placement constants
 export const EVENT_PLACEMENT_QUICK_LINKS = 'quick_links';