fix

Author: felixfeng33

apps/www/src/registry/default/components/editor/plugins/ai-plugins.tsx

@@ -2,13 +2,19 @@
 
 import React from 'react';
 
-import { AIChatPlugin, AIPlugin } from '@udecode/plate-ai/react';
+import type { AIChatPluginConfig } from '@udecode/plate-ai/react';
+
+import { PathApi } from '@udecode/plate';
+import { streamInsertChunk, withAIBatch } from '@udecode/plate-ai';
+import { AIChatPlugin, AIPlugin, useChatChunk } from '@udecode/plate-ai/react';
+import { useEditorPlugin, usePluginOption } from '@udecode/plate/react';
 
 import { markdownPlugin } from '@/registry/default/components/editor/plugins/markdown-plugin';
 import { AILoadingBar } from '@/registry/default/plate-ui/ai-loading-bar';
 import { AIMenu } from '@/registry/default/plate-ui/ai-menu';
 
 import { cursorOverlayPlugin } from './cursor-overlay-plugin';
+
 const systemCommon = `\
 You are an advanced AI-powered note-taking assistant, designed to enhance productivity and creativity in note management.
 Respond directly to user prompts with clear, concise, and relevant content. Maintain a neutral, helpful tone.
@@ -113,5 +119,55 @@ export const aiPlugins = [
       afterContainer: () => <AILoadingBar />,
       afterEditable: () => <AIMenu />,
     },
+  }).extend({
+    useHooks: () => {
+      const { editor, getOption } = useEditorPlugin(AIChatPlugin);
+
+      const mode = usePluginOption(
+        { key: 'aiChat' } as AIChatPluginConfig,
+        'mode'
+      );
+
+      useChatChunk({
+        onChunk: ({ chunk, isFirst, nodes, text }) => {
+          if (isFirst && mode == 'insert') {
+            editor.tf.withoutSaving(() => {
+              editor.tf.insertNodes(
+                {
+                  children: [{ text: '' }],
+                  type: AIChatPlugin.key,
+                },
+                {
+                  at: PathApi.next(editor.selection!.focus.path.slice(0, 1)),
+                }
+              );
+            });
+            editor.setOption(AIChatPlugin, 'streaming', true);
+          }
+
+          if (mode === 'insert' && nodes.length > 0) {
+            withAIBatch(
+              editor,
+              () => {
+                if (!getOption('streaming')) return;
+                editor.tf.withScrolling(() => {
+                  streamInsertChunk(editor, chunk, {
+                    textProps: {
+                      ai: true,
+                    },
+                  });
+                });
+              },
+              { split: isFirst }
+            );
+          }
+        },
+        onFinish: ({ content }) => {
+          editor.setOption(AIChatPlugin, 'streaming', false);
+          editor.setOption(AIChatPlugin, '_blockChunks', '');
+          editor.setOption(AIChatPlugin, '_blockPath', null);
+        },
+      });
+    },
   }),
 ] as const;

packages/ai/src/react/ai-chat/AIChatPlugin.ts

@@ -108,11 +108,9 @@ export const AIChatPlugin = createTPlatePlugin<AIChatPluginConfig>({
     promptTemplate: () => '{prompt}',
     systemTemplate: () => {},
   },
+  useHooks: useAIChatHooks,
 })
   .overrideEditor(withAIChat)
-  .extend(() => ({
-    useHooks: useAIChatHooks,
-  }))
   .extendApi<
     Pick<
       AIChatPluginConfig['api']['aiChat'],

packages/ai/src/react/ai-chat/useAIChatHook.ts

@@ -5,6 +5,7 @@ import { streamInsertChunk, withAIBatch } from '../../lib';
 import { type AIChatPluginConfig, AIChatPlugin } from './AIChatPlugin';
 import { useChatChunk } from './hooks/useChatChunk';
 
+/** @deprecated Already moved to registry ai-plugins.tsx */
 export const useAIChatHooks = () => {
   const { editor, getOption } = useEditorPlugin(AIChatPlugin);