diff --git a/ai-blocks.php b/ai-blocks.php
index 8ebe70833dbdb49611d4ad06091db9f5063fa139..43f7b0b14bda2c4a820ec42e4da638d3c6346fed 100644
--- a/ai-blocks.php
+++ b/ai-blocks.php
@@ -8,18 +8,17 @@
 * Author URI: https://autistici.org
 **/
 
+function wpdocs_allowed_block_types ( $block_editor_context, $editor_context ) {
+    if ( ! empty( $editor_context->post ) ) {
+        return array(
+            'core/paragraph',
+            'core/heading',
+            'core/list',
+        );
+    }
  
-function blacklist_blocks( $allowed_blocks ) {
-	// get all the registered blocks
-	$blocks = WP_Block_Type_Registry::get_instance()->get_all_registered();
-
-	// then disable some of them
-	unset( $blocks[ 'core/embed' ] );
-
-	// return the new list of allowed blocks
-	return array_keys( $blocks );
-	
+    return $block_editor_context;
 }
-
-add_filter( 'allowed_block_types_all', 'blacklist_blocks' );
+ 
+add_filter( 'allowed_block_types_all', 'wpdocs_allowed_block_types', 10, 2 );