Explore having the LLM return content in JSON schema instead of WordPress block markup

includes/Classifai/Features/Prompts/ContentGeneration/return-format.php

@@ -2,48 +2,70 @@
 /**
  * Return-format instruction appended to the system message for the Content Generation feature.
  *
- * Describes the WordPress block markup that the model should emit.
+ * Instructs the model to emit a constrained JSON "BlockTree" structure, which
+ * ClassifAI converts to valid WordPress block markup client-side (rather than
+ * asking the model to hand-author fragile `<!-- wp:… -->` markup).
  *
  * @package Classifai
  */
 
 // phpcs:disable Squiz.PHP.Heredoc.NotAllowed, PluginCheck.CodeAnalysis.Heredoc.NotAllowed
 return <<<'INSTRUCTION'
-The content returned should be valid WordPress block markup as described below, using elements like paragraphs and headings where appropriate. Be selective on the elements you use, defaulting to paragraphs. Please check the content before returning to ensure each element has proper opening and closing block markup and HTML tags and any required block attributes. Ensure elements don't nest inside each other, i.e. don't put a paragraph inside another paragraph or a list within a paragraph. Don't start the content with a heading, start with a paragraph.
-
-Markup available to use; don't use any other blocks, even if requested:
-<!-- wp:paragraph -->
-<p>CONTENT</p>
-<!-- /wp:paragraph -->
-
-<!-- wp:heading -->
-<h2 class="wp-block-heading">CONTENT</h2>
-<!-- /wp:heading -->
-
-<!-- wp:table -->
-<figure class="wp-block-table"><table class="has-fixed-layout"><tbody><tr><td>CONTENT</td></tr><tr><td>CONTENT</td></tr></tbody></table></figure>
-<!-- /wp:table -->
-
-<!-- wp:quote -->
-<blockquote class="wp-block-quote">
-<p>CONTENT</p>
-</blockquote>
-<!-- /wp:quote -->
-
-<!-- wp:pullquote -->
-<figure class="wp-block-pullquote"><blockquote><p>QUOTE</p><cite>AUTHOR</cite></blockquote></figure>
-<!-- /wp:pullquote -->
-
-<!-- wp:list -->
-<ul class="wp-block-list">
-<li>CONTENT</li>
-</ul>
-<!-- /wp:list -->
-
-<!-- wp:list {"ordered":true} -->
-<ol class="wp-block-list">
-<li>CONTENT</li>
-</ol>
-<!-- /wp:list -->
+Return the content as a single JSON object describing a flat WordPress "block tree". Do not return HTML, Markdown, block comment markup, code fences, or any prose. Output only the JSON object.
+
+## JSON Structure
+
+Output valid JSON matching this structure:
+
+interface BlockTree {
+  root: string;                            // Key of the root element
+  elements: Record<string, BlockElement>;  // Map of key -> element
+}
+
+interface BlockElement {
+  key: string;                     // Unique identifier, matching its key in `elements`
+  type: string;                    // Block name, e.g. "core/paragraph"
+  props: Record<string, unknown>;  // Block attributes (use {} when there are none)
+  children?: string[];             // Ordered keys of child elements (for blocks that hold inner blocks)
+  parentKey?: string;              // Key of the parent element
+}
+
+Rules:
+- Every element's `key` must match its key in `elements`.
+- The root must be a single element. To return multiple top-level blocks, make the root an element of type "fragment" with no props and list the top-level block keys in its `children`. "fragment" is a virtual wrapper only; it produces no markup of its own.
+- Omit `children` for blocks that hold no inner blocks.
+- Default to paragraphs; be selective with other blocks. Do not start the content with a heading; start with a paragraph.
+- Use only the block types listed below; do not use any other blocks, even if requested.
+
+## Available blocks
+
+- core/paragraph — props: { "content": string }.
+- core/heading — props: { "content": string, "level": 2 or 3 }.
+- core/list — props: { "ordered": boolean (optional, default false) }. Supports children: core/list-item.
+- core/list-item — props: { "content": string }. Must be inside core/list.
+- core/quote — props: { "citation": string (optional) }. Supports children: core/paragraph.
+- core/pullquote — props: { "value": string, "citation": string (optional) }.
+- core/table — props: { "body": [ { "cells": [ { "content": string, "tag": "td" } ] } ] }. Each row is an object with a "cells" array; each cell has "content" and "tag" ("td").
+- core/separator — props: {}.
+- core/image — props: { "url": string, "alt": string (optional), "caption": string (optional) }.
+- core/group — props: { "layout": { "type": "constrained" } }. Supports children: any blocks. Use to group related blocks.
+- core/columns — props: {}. Supports children: core/column (two or more).
+- core/column — props: { "width": string (optional, e.g. "50%") }. Must be inside core/columns. Supports children: any blocks.
+- core/buttons — props: {}. Supports children: core/button.
+- core/button — props: { "text": string, "url": string (optional) }. Must be inside core/buttons.
+
+## Block requirements
+
+- core/list: must use core/list-item children for list items; the deprecated "values" attribute is not supported.
+- core/quote: place the quoted text in one or more core/paragraph children; use the optional "citation" prop for attribution.
+- core/button: must be a direct child of a core/buttons wrapper.
+- core/column: must be a direct child of a core/columns wrapper.
+
+## Example
+
+Input: "A short intro about a topic, a section heading, and two key points."
+
+Output:
+{"root":"r","elements":{"r":{"key":"r","type":"fragment","props":{},"children":["p1","h1","l1"]},"p1":{"key":"p1","type":"core/paragraph","props":{"content":"An opening paragraph that introduces the topic."},"parentKey":"r"},"h1":{"key":"h1","type":"core/heading","props":{"content":"A section heading","level":2},"parentKey":"r"},"l1":{"key":"l1","type":"core/list","props":{"ordered":false},"children":["li1","li2"],"parentKey":"r"},"li1":{"key":"li1","type":"core/list-item","props":{"content":"First point"},"parentKey":"l1"},"li2":{"key":"li2","type":"core/list-item","props":{"content":"Second point"},"parentKey":"l1"}}}
 INSTRUCTION;
 // phpcs:enable

includes/Classifai/Features/QuickDraftIntegration.php

@@ -206,19 +206,9 @@ public function endpoint_callback( WP_REST_Request $request ) {
 			return $result;
 		}
 
-		// Update the post with generated content.
-		$updated_post = array(
-			'ID'           => $post_id,
-			'post_content' => $result,
-			'post_status'  => 'draft',
-		);
-
-		$update_result = wp_update_post( $updated_post );
-
-		if ( is_wp_error( $update_result ) ) {
-			return new WP_Error( 'post_update_failed', esc_html__( 'Failed to update post with generated content.', 'classifai' ) );
-		}
-
+		// $result is a JSON BlockTree. The draft is left empty here; the client
+		// renders it to block markup (using the editor's block registry) and
+		// saves it back via the core REST API before redirecting the user.
 		return rest_ensure_response(
 			array(
 				'post_id'  => $post_id,

includes/Classifai/Helpers.php

@@ -988,3 +988,52 @@ function get_temperature( float $temperature, int $results = 1 ): float {
 
 	return (float) min( 2.0, $temperature + ( $results / 10 ) );
 }
+
+/**
+ * Recursively sanitize the values of an AI-generated block tree.
+ *
+ * The Content Generation feature receives a JSON "block tree" from the model
+ * and renders its string props (paragraph and heading content, captions, list
+ * items, table cells, etc.) as HTML in the editor and saves them to post
+ * content. Sanitize every string value with wp_kses_post() and treat `url`
+ * props as URLs so untrusted markup (e.g. script tags or javascript: URLs)
+ * can't reach the browser.
+ *
+ * Decode the tree with json_decode( $json ) (objects, not associative arrays)
+ * so that empty objects such as `"props":{}` survive re-encoding; decoding to
+ * associative arrays would turn them into `[]`, which the client-side block
+ * tree schema rejects.
+ *
+ * @param mixed $value Decoded block tree, or a nested value within it.
+ * @return mixed Sanitized value.
+ */
+function sanitize_generated_block_tree( $value ) {
+	if ( is_object( $value ) ) {
+		foreach ( get_object_vars( $value ) as $key => $item ) {
+			if ( 'url' === $key && is_string( $item ) ) {
+				$value->$key = esc_url_raw( $item );
+			} else {
+				$value->$key = sanitize_generated_block_tree( $item );
+			}
+		}
+		return $value;
+	}
+
+	if ( is_array( $value ) ) {
+		$sanitized = array();
+		foreach ( $value as $key => $item ) {
+			if ( 'url' === $key && is_string( $item ) ) {
+				$sanitized[ $key ] = esc_url_raw( $item );
+			} else {
+				$sanitized[ $key ] = sanitize_generated_block_tree( $item );
+			}
+		}
+		return $sanitized;
+	}
+
+	if ( is_string( $value ) ) {
+		return wp_kses_post( $value );
+	}
+
+	return $value;
+}

includes/Classifai/Providers/Azure/OpenAI.php

@@ -18,6 +18,7 @@
 use function Classifai\sanitize_number_of_responses_field;
 use function Classifai\safe_wp_remote_post;
 use function Classifai\get_temperature;
+use function Classifai\sanitize_generated_block_tree;
 
 class OpenAI extends Provider {
 
@@ -954,8 +955,9 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 		$body = apply_filters(
 			'classifai_azure_openai_content_request_body',
 			array(
-				'messages'    => $messages,
-				'temperature' => 0.9,
+				'messages'        => $messages,
+				'temperature'     => 0.9,
+				'response_format' => array( 'type' => 'json_object' ),
 			),
 			$post_id
 		);
@@ -977,17 +979,28 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 			return $response;
 		}
 
-		// If we have a message, return it.
-		$return = '';
+		// Pull the message content out of the response.
+		$content = '';
 		if ( ! empty( $response['choices'] ) ) {
 			foreach ( $response['choices'] as $choice ) {
 				if ( isset( $choice['message'], $choice['message']['content'] ) ) {
-					$return = wp_kses_post( trim( $choice['message']['content'], ' "\'' ) );
+					$content = trim( $choice['message']['content'] );
 				}
 			}
 		}
 
-		return $return;
+		// The response should be a JSON BlockTree; validate before returning.
+		// Decode to objects (not arrays) so empty objects like "props":{} are
+		// preserved when re-encoded rather than becoming "props":[].
+		$decoded = json_decode( $content );
+		if ( null === $decoded || JSON_ERROR_NONE !== json_last_error() ) {
+			return new WP_Error( 'invalid_content_response', esc_html__( 'The generated content was not in the expected format. Please try again.', 'classifai' ) );
+		}
+
+		// Sanitize the block tree's string values before they are rendered/saved.
+		$decoded = sanitize_generated_block_tree( $decoded );
+
+		return wp_json_encode( $decoded );
 	}
 
 	/**

includes/Classifai/Providers/Localhost/Ollama.php

@@ -17,6 +17,7 @@
 
 use function Classifai\get_default_prompt;
 use function Classifai\sanitize_number_of_responses_field;
+use function Classifai\sanitize_generated_block_tree;
 
 /**
  * Ollama class
@@ -787,6 +788,7 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 				'model'    => $settings[ static::ID ]['model'] ?? '',
 				'messages' => $messages,
 				'stream'   => false,
+				'format'   => 'json',
 			),
 			$post_id
 		);
@@ -804,13 +806,24 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 			return $response;
 		}
 
-		// If we have a message, return it.
-		$return = '';
+		// Pull the message content out of the response.
+		$content = '';
 		if ( isset( $response['message'], $response['message']['content'] ) ) {
-			$return = wp_kses_post( trim( $response['message']['content'], ' "\'' ) );
+			$content = trim( $response['message']['content'] );
 		}
 
-		return $return;
+		// The response should be a JSON BlockTree; validate before returning.
+		// Decode to objects (not arrays) so empty objects like "props":{} are
+		// preserved when re-encoded rather than becoming "props":[].
+		$decoded = json_decode( $content );
+		if ( null === $decoded || JSON_ERROR_NONE !== json_last_error() ) {
+			return new WP_Error( 'invalid_content_response', esc_html__( 'The generated content was not in the expected format. Please try again.', 'classifai' ) );
+		}
+
+		// Sanitize the block tree's string values before they are rendered/saved.
+		$decoded = sanitize_generated_block_tree( $decoded );
+
+		return wp_json_encode( $decoded );
 	}
 
 	/**

includes/Classifai/Providers/OpenAI/ChatGPT.php

@@ -22,6 +22,7 @@
 use function Classifai\get_modified_image_source_url;
 use function Classifai\get_largest_size_and_dimensions_image_url;
 use function Classifai\get_temperature;
+use function Classifai\sanitize_generated_block_tree;
 
 class ChatGPT extends Provider {
 
@@ -1216,9 +1217,10 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 		$body = apply_filters(
 			'classifai_chatgpt_content_request_body',
 			array(
-				'model'       => $this->chatgpt_model,
-				'messages'    => $messages,
-				'temperature' => 0.9,
+				'model'           => $this->chatgpt_model,
+				'messages'        => $messages,
+				'temperature'     => 0.9,
+				'response_format' => array( 'type' => 'json_object' ),
 			),
 			$post_id
 		);
@@ -1235,17 +1237,28 @@ public function generate_content( int $post_id = 0, array $args = array() ) {
 			return $response;
 		}
 
-		// If we have a message, return it.
-		$return = '';
+		// Pull the message content out of the response.
+		$content = '';
 		if ( ! empty( $response['choices'] ) ) {
 			foreach ( $response['choices'] as $choice ) {
 				if ( isset( $choice['message'], $choice['message']['content'] ) ) {
-					$return = wp_kses_post( trim( $choice['message']['content'], ' "\'' ) );
+					$content = trim( $choice['message']['content'] );
 				}
 			}
 		}
 
-		return $return;
+		// The response should be a JSON BlockTree; validate before returning.
+		// Decode to objects (not arrays) so empty objects like "props":{} are
+		// preserved when re-encoded rather than becoming "props":[].
+		$decoded = json_decode( $content );
+		if ( null === $decoded || JSON_ERROR_NONE !== json_last_error() ) {
+			return new WP_Error( 'invalid_content_response', esc_html__( 'The generated content was not in the expected format. Please try again.', 'classifai' ) );
+		}
+
+		// Sanitize the block tree's string values before they are rendered/saved.
+		$decoded = sanitize_generated_block_tree( $decoded );
+
+		return wp_json_encode( $decoded );
 	}
 
 	/**