AI Assistant: Add image index information on caption generation reque…

…st (#42274)

* Add getAllBlocks and getSerializedPostContent lib functions

* add image index for caption requests

* changelog

* treat response for common parsing error

* change indexes to positions

* change positions to block clientId

projects/js-packages/ai-client/changelog/update-jetpack-ai-image-extension-image-index

@@ -0,0 +1,5 @@
+Significance: patch
+Type: added
+Comment: Added utility functions
+
+

projects/js-packages/ai-client/src/hooks/use-post-content.ts

@@ -28,16 +28,24 @@ const usePostContent = () => {
 		};
 	}, [] );
 
+	const getSerializedPostContent = useCallback( () => {
+		const blocks = getBlocks();
+
+		if ( blocks.length === 0 ) {
+			return '';
+		}
+
+		return serialize( blocks );
+	}, [ getBlocks ] );
+
 	const getPostContent = useCallback(
 		( preprocess?: ( serialized: string ) => string ) => {
-			const blocks = getBlocks();
+			let serialized = getSerializedPostContent();
 
-			if ( blocks.length === 0 ) {
+			if ( ! serialized ) {
 				return '';
 			}
 
-			let serialized = serialize( blocks );
-
 			if ( preprocess && typeof preprocess === 'function' ) {
 				serialized = preprocess( serialized );
 			}
@@ -47,7 +55,7 @@ const usePostContent = () => {
 		[ getBlocks ]
 	);
 
-	return { getPostContent, isEditedPostEmpty };
+	return { getPostContent, isEditedPostEmpty, getSerializedPostContent };
 };
 
 export default usePostContent;

projects/js-packages/ai-client/src/libs/get-all-blocks.ts

@@ -0,0 +1,29 @@
+/**
+ * External dependencies
+ */
+import { select } from '@wordpress/data';
+/**
+ * Internal dependencies
+ */
+import type { Block } from '../types.js';
+
+/**
+ * Recursively get all blocks from the post, including nested innerBlocks
+ * @return {Array} Array of all blocks in the post
+ */
+export const getAllBlocks = (): Array< Block > => {
+	const topLevelBlocks = select( 'core/block-editor' ).getBlocks();
+	const allBlocks: Array< Block > = [];
+
+	const processBlock = ( block: Block ) => {
+		allBlocks.push( block );
+
+		if ( block.innerBlocks?.length ) {
+			block.innerBlocks.forEach( processBlock );
+		}
+	};
+
+	topLevelBlocks.forEach( processBlock );
+
+	return allBlocks;
+};

projects/js-packages/ai-client/src/libs/index.ts

@@ -13,3 +13,5 @@ export { mapActionToHumanText } from './map-action-to-human-text.js';
 export { openBlockSidebar } from './open-block-sidebar.js';
 
 export { showAiAssistantSection } from './show-ai-assistant-section.js';
+
+export { getAllBlocks } from './get-all-blocks.js';

projects/plugins/jetpack/changelog/update-jetpack-ai-image-extension-image-index

@@ -0,0 +1,4 @@
+Significance: patch
+Type: other
+
+AI Assistant: Add image index information on caption generation request

projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/image/with-ai-image-extension.tsx

@@ -137,10 +137,16 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 				startLoading( type );
 
 				try {
+					const context: { positions?: number[] } = {};
+
 					if ( type === TYPE_ALT_TEXT ) {
 						openBlockSidebar( props.clientId );
 					}
 
+					if ( type === TYPE_CAPTION ) {
+						context.positions = [ props.clientId ];
+					}
+
 					dequeueAsyncRequest();
 
 					const response = await askQuestionSync(
@@ -150,8 +156,7 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 								context: {
 									type: type,
 									content: getPostContent( preprocessImageContent ),
-									// URL of the image for the AI to find where the image is in the post.
-									urls: [ props.attributes.url ],
+									...context,
 									images: [
 										{
 											// We convert the image to a base64 string to avoid inaccesible URLs for private images.
@@ -169,7 +174,11 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 
 					increaseRequestsCount();
 
-					const parsedResponse: { texts?: string[]; captions?: string[] } = JSON.parse( response );
+					const parsedResponse: { texts?: string[]; captions?: string[] } = JSON.parse(
+						response
+							?.replace?.( /^```json\s*/, '' ) // Remove the markdown code block if it exists.
+							?.replace( /```$/, '' )
+					);
 
 					if ( type === TYPE_ALT_TEXT ) {
 						const alt = parsedResponse.texts?.[ 0 ];

projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/lib/preprocess-image-content.ts

@@ -1,20 +1,32 @@
 /**
- * Preprocess the serialized post content to remove image URLs and alt text, adding a unique count to each image.
- * This is used to ensure that the AI is not confused by the URLs in the post content.
+ * External dependencies
+ */
+import { getAllBlocks } from '@automattic/jetpack-ai-client';
+
+/**
+ * Preprocess the serialized post content to remove image captions and alt text, changing the image urls to their clientId.
+ * This is used to ensure that the AI is not confused by URLs in the post content, while still being able to identify the image blocks.
  * @param {string} content - The content to preprocess.
  * @return {string} The preprocessed content.
  */
 export const preprocessImageContent = ( content: string ): string => {
-	let imageCounter = 0;
+	const imageBlocks = getAllBlocks().filter( block => block.name === 'core/image' );
+	let currentImageIndex = 0;
 
 	// Remove figcaption elements
 	content = content.replace( /<figcaption[^>]*>.*?<\/figcaption>/g, '' );
 
-	// Replace image URLs with a unique count and remove alt text
-	content = content.replace( /<img[^>]*>/g, () => {
-		imageCounter++;
+	// Replace the urls of images within image blocks with their clientId
+	content = content.replace( /<!-- wp:image[^>]*>.*?<img[^>]*>.*?<!-- \/wp:image -->/gs, () => {
+		// The assumption here is that the image blocks are always in the same order as the images in the post's HTML content
+		const imageBlock = imageBlocks[ currentImageIndex ];
+		currentImageIndex++;
+
+		if ( ! imageBlock ) {
+			return '';
+		}
 
-		return `<img src="image-${ imageCounter }">`;
+		return `<img src="${ imageBlock.clientId }">`;
 	} );
 
 	return content;