@ai-sdk/react 0.0.29 → 0.0.31

package/src/use-chat.ui.test.tsx

@@ -1,9 +1,12 @@
 /* eslint-disable @next/next/no-img-element */
 import { withTestServer } from '@ai-sdk/provider-utils/test';
-import { formatStreamPart, getTextFromDataUrl } from '@ai-sdk/ui-utils';
+import {
+  formatStreamPart,
+  getTextFromDataUrl,
+  Message,
+} from '@ai-sdk/ui-utils';
 import '@testing-library/jest-dom/vitest';
 import {
-  RenderResult,
   cleanup,
   findByText,
   render,
@@ -15,9 +18,26 @@ import React, { useRef, useState } from 'react';
 import { useChat } from './use-chat';
 
 describe('stream data stream', () => {
+  let onFinishCalls: Array<{
+    message: Message;
+    options: {
+      finishReason: string;
+      usage: {
+        completionTokens: number;
+        promptTokens: number;
+        totalTokens: number;
+      };
+    };
+  }> = [];
+
   const TestComponent = () => {
     const [id, setId] = React.useState<string>('first-id');
-    const { messages, append, error, data, isLoading } = useChat({ id });
+    const { messages, append, error, data, isLoading } = useChat({
+      id,
+      onFinish: (message, options) => {
+        onFinishCalls.push({ message, options });
+      },
+    });
 
     return (
       <div>
@@ -49,11 +69,13 @@ describe('stream data stream', () => {
 
   beforeEach(() => {
     render(<TestComponent />);
+    onFinishCalls = [];
   });
 
   afterEach(() => {
     vi.restoreAllMocks();
     cleanup();
+    onFinishCalls = [];
   });
 
   it(
@@ -148,6 +170,50 @@ describe('stream data stream', () => {
     );
   });
 
+  it(
+    'should invoke onFinish when the stream finishes',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: [
+          formatStreamPart('text', 'Hello'),
+          formatStreamPart('text', ','),
+          formatStreamPart('text', ' world'),
+          formatStreamPart('text', '.'),
+          formatStreamPart('finish_message', {
+            finishReason: 'stop',
+            usage: { completionTokens: 1, promptTokens: 3 },
+          }),
+        ],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        await screen.findByTestId('message-1');
+
+        expect(onFinishCalls).toStrictEqual([
+          {
+            message: {
+              id: expect.any(String),
+              createdAt: expect.any(Date),
+              role: 'assistant',
+              content: 'Hello, world.',
+            },
+            options: {
+              finishReason: 'stop',
+              usage: {
+                completionTokens: 1,
+                promptTokens: 3,
+                totalTokens: 4,
+              },
+            },
+          },
+        ]);
+      },
+    ),
+  );
+
   describe('id', () => {
     it(
       'should clear out messages when the id changes',
@@ -175,8 +241,25 @@ describe('stream data stream', () => {
 });
 
 describe('text stream', () => {
+  let onFinishCalls: Array<{
+    message: Message;
+    options: {
+      finishReason: string;
+      usage: {
+        completionTokens: number;
+        promptTokens: number;
+        totalTokens: number;
+      };
+    };
+  }> = [];
+
   const TestComponent = () => {
-    const { messages, append } = useChat({ streamMode: 'text' });
+    const { messages, append } = useChat({
+      streamProtocol: 'text',
+      onFinish: (message, options) => {
+        onFinishCalls.push({ message, options });
+      },
+    });
 
     return (
       <div>
@@ -199,11 +282,13 @@ describe('text stream', () => {
 
   beforeEach(() => {
     render(<TestComponent />);
+    onFinishCalls = [];
   });
 
   afterEach(() => {
     vi.restoreAllMocks();
     cleanup();
+    onFinishCalls = [];
   });
 
   it(
@@ -229,12 +314,47 @@ describe('text stream', () => {
       },
     ),
   );
+
+  it(
+    'should invoke onFinish when the stream finishes',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['Hello', ',', ' world', '.'],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append-text-stream'));
+
+        await screen.findByTestId('message-1-text-stream');
+
+        expect(onFinishCalls).toStrictEqual([
+          {
+            message: {
+              id: expect.any(String),
+              createdAt: expect.any(Date),
+              role: 'assistant',
+              content: 'Hello, world.',
+            },
+            options: {
+              finishReason: 'unknown',
+              usage: {
+                completionTokens: NaN,
+                promptTokens: NaN,
+                totalTokens: NaN,
+              },
+            },
+          },
+        ]);
+      },
+    ),
+  );
 });
 
 describe('form actions', () => {
   const TestComponent = () => {
     const { messages, handleSubmit, handleInputChange, isLoading, input } =
-      useChat({ streamMode: 'text' });
+      useChat({ streamProtocol: 'text' });
 
     return (
       <div>
@@ -307,7 +427,7 @@ describe('form actions', () => {
 describe('form actions (with options)', () => {
   const TestComponent = () => {
     const { messages, handleSubmit, handleInputChange, isLoading, input } =
-      useChat({ streamMode: 'text' });
+      useChat({ streamProtocol: 'text' });
 
     return (
       <div>
@@ -1192,3 +1312,92 @@ describe('file attachments with url', () => {
     ),
   );
 });
+
+describe('reload', () => {
+  const TestComponent = () => {
+    const { messages, append, reload } = useChat();
+
+    return (
+      <div>
+        {messages.map((m, idx) => (
+          <div data-testid={`message-${idx}`} key={m.id}>
+            {m.role === 'user' ? 'User: ' : 'AI: '}
+            {m.content}
+          </div>
+        ))}
+
+        <button
+          data-testid="do-append"
+          onClick={() => {
+            append({ role: 'user', content: 'hi' });
+          }}
+        />
+
+        <button
+          data-testid="do-reload"
+          onClick={() => {
+            reload({
+              data: { 'test-data-key': 'test-data-value' },
+              body: { 'request-body-key': 'request-body-value' },
+              headers: { 'header-key': 'header-value' },
+            });
+          }}
+        />
+      </div>
+    );
+  };
+
+  beforeEach(() => {
+    render(<TestComponent />);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    cleanup();
+  });
+
+  it(
+    'should show streamed response',
+    withTestServer(
+      [
+        {
+          url: '/api/chat',
+          type: 'stream-values',
+          content: ['0:"first response"\n'],
+        },
+        {
+          url: '/api/chat',
+          type: 'stream-values',
+          content: ['0:"second response"\n'],
+        },
+      ],
+      async ({ call }) => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
+
+        await screen.findByTestId('message-1');
+
+        // setup done, click reload:
+        await userEvent.click(screen.getByTestId('do-reload'));
+
+        expect(await call(1).getRequestBodyJson()).toStrictEqual({
+          messages: [{ content: 'hi', role: 'user' }],
+          data: { 'test-data-key': 'test-data-value' },
+          'request-body-key': 'request-body-value',
+        });
+
+        expect(call(1).getRequestHeaders()).toStrictEqual({
+          'content-type': 'application/json',
+          'header-key': 'header-value',
+        });
+
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: second response',
+        );
+      },
+    ),
+  );
+});

package/src/use-completion.ts

@@ -72,11 +72,17 @@ export function useCompletion({
   headers,
   body,
   streamMode,
+  streamProtocol,
   fetch,
   onResponse,
   onFinish,
   onError,
 }: UseCompletionOptions = {}): UseCompletionHelpers {
+  // streamMode is deprecated, use streamProtocol instead.
+  if (streamMode) {
+    streamProtocol ??= streamMode === 'text' ? 'text' : undefined;
+  }
+
   // Generate an unique id for the completion if not provided.
   const hookId = useId();
   const completionId = id || hookId;
@@ -126,7 +132,7 @@ export function useCompletion({
           ...extraMetadataRef.current.body,
           ...options?.body,
         },
-        streamMode,
+        streamProtocol,
         fetch,
         setCompletion: completion => mutate(completion, false),
         setLoading: mutateLoading,
@@ -150,7 +156,7 @@ export function useCompletion({
       onError,
       setError,
       streamData,
-      streamMode,
+      streamProtocol,
       fetch,
       mutateStreamData,
     ],

package/src/use-completion.ui.test.tsx

@@ -124,9 +124,7 @@ describe('stream data stream', () => {
 describe('text stream', () => {
   const TestComponent = () => {
     const { completion, handleSubmit, handleInputChange, input } =
-      useCompletion({
-        streamMode: 'text',
-      });
+      useCompletion({ streamProtocol: 'text' });
 
     return (
       <div>