ai

apps/www/src/app/api/stream/route.ts

@@ -0,0 +1,19 @@
+export function POST() {
+  const streams = [
+    { delay: 100, texts: 'Hello' },
+    { delay: 200, texts: 'World' },
+  ];
+
+  const stream = new ReadableStream({
+    async start(controller) {
+      for (const stream of streams) {
+        await new Promise((resolve) => setTimeout(resolve, stream.delay)); // Increasing delay
+        controller.enqueue(stream.texts);
+      }
+
+      controller.close();
+    },
+  });
+
+  return new Response(stream);
+}

apps/www/src/registry/default/example/playground-demo.tsx

@@ -171,6 +171,25 @@ export const usePlaygroundEditor = (id: any = '', scrollSelector?: string) => {
         AIPlugin.configure({
           options: {
             createAIEditor: createAIEditor,
+            // eslint-disable-next-line @typescript-eslint/require-await, @typescript-eslint/no-unused-vars
+            fetchSuggestion: async ({ abortSignal, prompt, system }) => {
+              const response = await fetch('/api/stream', {
+                body: JSON.stringify({ prompt, system }),
+                headers: {
+                  'Content-Type': 'application/json',
+                },
+                method: 'POST',
+                signal: abortSignal.signal,
+              }).catch((error) => {
+                console.error(error);
+              });
+
+              if (!response || !response.body) {
+                throw new Error('Response or response body is null or abort');
+              }
+
+              return response.body;
+            },
             scrollContainerSelector: `#${scrollSelector}`,
           },
           render: { aboveEditable: AIMenu },

apps/www/src/registry/default/plate-ui/ai-actions.tsx

@@ -195,3 +195,10 @@ export const SelectionSuggestionActions = {
     value: ACTION_SELECTION_SUGGESTION_TRY_AGAIN,
   },
 };
+
+export const aiActions = {
+  CursorCommandsActions,
+  CursorSuggestionActions,
+  SelectionCommandsActions,
+  SelectionSuggestionActions,
+};

apps/www/src/registry/default/plate-ui/ai-menu-items.tsx

@@ -86,3 +86,10 @@ export const SelectionSuggestions = () => {
     </>
   );
 };
+
+export const aiCommands = {
+  CursorCommands,
+  CursorSuggestions,
+  SelectionCommands,
+  SelectionSuggestions,
+};

apps/www/src/registry/default/plate-ui/ai-menu.tsx

@@ -9,19 +9,8 @@ import { useAI } from '@udecode/plate-ai/react';
 import { Icons } from '@/components/icons';
 
 import { useActionHandler } from './action-handler';
-import {
-  CursorCommandsActions,
-  CursorSuggestionActions,
-  SelectionCommandsActions,
-  SelectionSuggestionActions,
-  defaultValues,
-} from './ai-actions';
-import {
-  CursorCommands,
-  CursorSuggestions,
-  SelectionCommands,
-  SelectionSuggestions,
-} from './ai-menu-items';
+import { aiActions, defaultValues } from './ai-actions';
+import { aiCommands } from './ai-menu-items';
 import { AIPreviewEditor } from './ai-previdew-editor';
 import { Button } from './button';
 import { Menu, comboboxVariants, renderSearchMenuItems } from './menu';
@@ -40,25 +29,13 @@ export const AIMenu = memo(({ children }: React.PropsWithChildren) => {
     submitButtonProps,
     onCloseMenu,
   } = useAI({
-    aiActions: {
-      CursorCommandsActions,
-      CursorSuggestionActions,
-      SelectionCommandsActions,
-      SelectionSuggestionActions,
-    },
-    aiCommands: {
-      CursorCommands,
-      CursorSuggestions,
-      SelectionCommands,
-      SelectionSuggestions,
-    },
+    aiActions: aiActions,
+    aiCommands: aiCommands,
     defaultValues,
   });
 
   useActionHandler(action, aiEditor!);
 
-  /** IME */
-
   return (
     <>
       <Menu

packages/ai/src/react/ai/AIPlugin.ts

@@ -13,9 +13,16 @@ import { useAIHooks } from './useAIHook';
 
 export const KEY_AI = 'ai';
 
+export interface FetchAISuggestionProps {
+  abortSignal: AbortController;
+  prompt: string;
+  system?: string;
+}
+
 interface ExposeOptions {
   createAIEditor: () => PlateEditor;
   scrollContainerSelector: string;
+  fetchSuggestion?: (props: FetchAISuggestionProps) => Promise<ReadableStream>;
   trigger?: RegExp | string[] | string;
 
   triggerPreviousCharPattern?: RegExp;

packages/ai/src/react/ai/hook/useAI.ts

@@ -97,7 +97,6 @@ export const useAI = ({
     if (isHotkey('enter')(e)) await streamInsert();
   };
 
-  // TODO: move to API
   const onCloseMenu = useCallback(() => {
     //  close menu if ai is not generating
     if (aiState === 'idle' || aiState === 'done') {

packages/ai/src/react/ai/stream/streamTraversal.ts

@@ -17,28 +17,15 @@ export const streamTraversal = async (
   const abortController = new AbortController();
   editor.setOptions(AIPlugin, { abortController });
 
-  const response = await fetch('/api/ai/command', {
-    body: JSON.stringify({ prompt, system }),
-    headers: {
-      'Content-Type': 'application/json',
-    },
-    method: 'POST',
-    signal: abortController.signal,
-  }).catch((error) => {
-    console.error(error);
-  });
+  const fetchSuggestion = editor.getOptions(AIPlugin).fetchSuggestion!;
 
-  if (response?.status === 429) {
-    return fn(
-      'Rate limit exceeded. You have made too many requests. Please try again later.',
-      true
-    );
-  }
-  if (!response || !response.body) {
-    throw new Error('Response or response body is null or abort');
-  }
+  const response = await fetchSuggestion({
+    abortSignal: abortController,
+    prompt,
+    system,
+  });
 
-  const reader = response.body.getReader();
+  const reader = response.getReader();
   const decoder = new TextDecoder();
 
   // eslint-disable-next-line no-constant-condition