AI Extension: handle multiple blocks editing (#31491)

* [not verified] introduce getBlockTextContent() fn

* [not verified] pick block content by using getBlockTextContent

* [not verified] try to pick content from attr first

* [not verified] get block content only when requesting

* [not verified] jsdoc improvement

* [not verified] add a function to get the content of selected blocks

* [not verified] pull content from selected blocks

* [not verified] change the API of helper function

* [not verified] expose clientIds from getTextContentFromBlocks()

* [not verified] combine multi selected blocks into only one

* [not verified] changelog

* use ref to deal with selected client ids

* improve jsdoc. fix typo

projects/plugins/jetpack/changelog/update-ai-extension-edit-with-multiple-selection

@@ -0,0 +1,4 @@
+Significance: patch
+Type: enhancement
+
+AI Extension: handle multiple blocks editing

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

@@ -5,7 +5,7 @@ import { BlockControls } from '@wordpress/block-editor';
 import { store as blockEditorStore } from '@wordpress/block-editor';
 import { createHigherOrderComponent } from '@wordpress/compose';
 import { useDispatch } from '@wordpress/data';
-import { useCallback, useState } from '@wordpress/element';
+import { useCallback, useState, useRef } from '@wordpress/element';
 import React from 'react';
 /**
  * Internal dependencies
@@ -15,6 +15,7 @@ import AiAssistantDropdown, {
 } from '../../components/ai-assistant-controls';
 import useSuggestionsFromAI from '../../hooks/use-suggestions-from-ai';
 import { getPrompt } from '../../lib/prompt';
+import { getTextContentFromBlocks } from '../../lib/utils/block-content';
 /*
  * Types
  */
@@ -30,18 +31,13 @@ type StoredPromptProps = {
  */
 export const withAIAssistant = createHigherOrderComponent(
 	BlockEdit => props => {
-		const { clientId } = props;
 		const [ storedPrompt, setStoredPrompt ] = useState< StoredPromptProps >( {
 			messages: [],
 		} );
 
-		const { updateBlockAttributes } = useDispatch( blockEditorStore );
+		const clientIdsRef = useRef< Array< string > >();
 
-		/*
-		 * Pick the content from the block attribute from now.
-		 * @todo: it doesn't scale well, we need to find a better way to get the content.
-		 */
-		const content: string = props?.attributes?.content;
+		const { updateBlockAttributes, removeBlocks } = useDispatch( blockEditorStore );
 
 		/**
 		 * Set the content of the block.
@@ -51,16 +47,29 @@ export const withAIAssistant = createHigherOrderComponent(
 		 */
 		const setContent = useCallback(
 			( newContent: string ) => {
+				/*
+				 * Pick the first item of the array,
+				 * to be udpated with the new content.
+				 * The rest of the items will be removed.
+				 */
+				const [ firstClientId, ...restClientIds ] = clientIdsRef.current;
 				/*
 				 * Update the content of the block
 				 * by calling the setAttributes function,
 				 * updating the `content` attribute.
 				 * It doesn't scale for other blocks.
 				 * @todo: find a better way to update the content.
 				 */
-				updateBlockAttributes( clientId, { content: newContent } );
+				updateBlockAttributes( firstClientId, { content: newContent } );
+
+				// Remove the rest of the block in case there are more than one.
+				if ( restClientIds.length ) {
+					removeBlocks( restClientIds ).then( () => {
+						clientIdsRef.current = [ firstClientId ];
+					} );
+				}
 			},
-			[ clientId, updateBlockAttributes ]
+			[ removeBlocks, updateBlockAttributes ]
 		);
 
 		const addAssistantMessage = useCallback(
@@ -102,6 +111,15 @@ export const withAIAssistant = createHigherOrderComponent(
 
 		const requestSuggestion = useCallback(
 			( promptType: PromptTypeProp, options: AiAssistantDropdownOnChangeOptionsArgProps ) => {
+				const { content, clientIds } = getTextContentFromBlocks();
+
+				/*
+				 * Store the selected clientIds when the user requests a suggestion.
+				 * The client Ids will be used to update the content of the block,
+				 * when suggestions are received from the AI.
+				 */
+				clientIdsRef.current = clientIds;
+
 				setStoredPrompt( prevPrompt => {
 					const freshPrompt = {
 						...prevPrompt,
@@ -115,11 +133,11 @@ export const withAIAssistant = createHigherOrderComponent(
 					// Request the suggestion from the AI.
 					request( freshPrompt.messages );
 
-					// Update the stored prompt.
+					// Update the stored prompt locally.
 					return freshPrompt;
 				} );
 			},
-			[ content, request ]
+			[ request ]
 		);
 
 		return (

projects/plugins/jetpack/extensions/blocks/ai-assistant/lib/utils/block-content.ts

@@ -1,13 +1,17 @@
 /**
  * External dependencies
  */
+import { store as blockEditorStore } from '@wordpress/block-editor';
+import { getBlockContent } from '@wordpress/blocks';
 import { serialize } from '@wordpress/blocks';
 import { select } from '@wordpress/data';
 import TurndownService from 'turndown';
 
 // Turndown instance
 const turndownService = new TurndownService();
 
+const HTML_JOIN_CHARACTERS = '<br />';
+
 /**
  * Returns partial content from the beginning of the post
  * to the current block, based on the given block clientId.
@@ -47,3 +51,74 @@ export function getContentFromBlocks(): string {
 
 	return turndownService.turndown( serialize( blocks ) );
 }
+
+type GetTextContentFromBlocksProps = {
+	count: number;
+	clientIds: string[];
+	content: string;
+};
+
+/**
+ * Returns the text content from all selected blocks.
+ *
+ * @returns {GetTextContentFromBlocksProps} The text content.
+ */
+export function getTextContentFromBlocks(): GetTextContentFromBlocksProps {
+	const clientIds = select( blockEditorStore ).getSelectedBlockClientIds();
+	const defaultContent = {
+		count: 0,
+		clientIds: [],
+		content: '',
+	};
+
+	if ( ! clientIds?.length ) {
+		return defaultContent;
+	}
+
+	const blocks = select( blockEditorStore ).getBlocksByClientId( clientIds );
+	if ( ! blocks?.length ) {
+		return defaultContent;
+	}
+
+	return {
+		count: blocks.length,
+		clientIds,
+		content: blocks
+			.map( block => getBlockTextContent( block.clientId ) )
+			.join( HTML_JOIN_CHARACTERS ),
+	};
+}
+
+/**
+ * Return the block content from the given block clientId.
+ *
+ * It will try to get the content from the block `content` attribute.
+ * Otherwise, it will try to get the content
+ * by using the `getBlockContent` function.
+ *
+ * @param {string} clientId   - The block clientId.
+ * @returns {string}            The block content.
+ */
+export function getBlockTextContent( clientId: string ): string {
+	if ( ! clientId ) {
+		return '';
+	}
+
+	const editor = select( blockEditorStore );
+	const block = editor.getBlock( clientId );
+
+	/*
+	 * In some context, the block can be undefined,
+	 * for instance, when previewing the block.
+	 */
+	if ( ! block ) {
+		return '';
+	}
+
+	// Attempt to pick the content from the block `content` attribute.
+	if ( block?.attributes?.content ) {
+		return block.attributes.content;
+	}
+
+	return getBlockContent( block );
+}