@ai-sdk/react 3.0.47 → 3.0.49

package/src/use-completion.ts

@@ -0,0 +1,208 @@
+import {
+  CompletionRequestOptions,
+  UseCompletionOptions,
+  callCompletionApi,
+} from 'ai';
+import { useCallback, useEffect, useId, useRef, useState } from 'react';
+import useSWR from 'swr';
+import { throttle } from './throttle';
+
+export type { UseCompletionOptions };
+
+export type UseCompletionHelpers = {
+  /** The current completion result */
+  completion: string;
+  /**
+   * Send a new prompt to the API endpoint and update the completion state.
+   */
+  complete: (
+    prompt: string,
+    options?: CompletionRequestOptions,
+  ) => Promise<string | null | undefined>;
+  /** The error object of the API request */
+  error: undefined | Error;
+  /**
+   * Abort the current API request but keep the generated tokens.
+   */
+  stop: () => void;
+  /**
+   * Update the `completion` state locally.
+   */
+  setCompletion: (completion: string) => void;
+  /** The current value of the input */
+  input: string;
+  /** setState-powered method to update the input value */
+  setInput: React.Dispatch<React.SetStateAction<string>>;
+  /**
+   * An input/textarea-ready onChange handler to control the value of the input
+   * @example
+   * ```jsx
+   * <input onChange={handleInputChange} value={input} />
+   * ```
+   */
+  handleInputChange: (
+    event:
+      | React.ChangeEvent<HTMLInputElement>
+      | React.ChangeEvent<HTMLTextAreaElement>,
+  ) => void;
+
+  /**
+   * Form submission handler to automatically reset input and append a user message
+   * @example
+   * ```jsx
+   * <form onSubmit={handleSubmit}>
+   *  <input onChange={handleInputChange} value={input} />
+   * </form>
+   * ```
+   */
+  handleSubmit: (event?: { preventDefault?: () => void }) => void;
+
+  /** Whether the API request is in progress */
+  isLoading: boolean;
+};
+
+export function useCompletion({
+  api = '/api/completion',
+  id,
+  initialCompletion = '',
+  initialInput = '',
+  credentials,
+  headers,
+  body,
+  streamProtocol = 'data',
+  fetch,
+  onFinish,
+  onError,
+  experimental_throttle: throttleWaitMs,
+}: UseCompletionOptions & {
+  /**
+   * Custom throttle wait in ms for the completion and data updates.
+   * Default is undefined, which disables throttling.
+   */
+  experimental_throttle?: number;
+} = {}): UseCompletionHelpers {
+  // Generate an unique id for the completion if not provided.
+  const hookId = useId();
+  const completionId = id || hookId;
+
+  // Store the completion state in SWR, using the completionId as the key to share states.
+  const { data, mutate } = useSWR<string>([api, completionId], null, {
+    fallbackData: initialCompletion,
+  });
+
+  const { data: isLoading = false, mutate: mutateLoading } = useSWR<boolean>(
+    [completionId, 'loading'],
+    null,
+  );
+
+  const [error, setError] = useState<undefined | Error>(undefined);
+  const completion = data!;
+
+  // Abort controller to cancel the current API call.
+  const [abortController, setAbortController] =
+    useState<AbortController | null>(null);
+
+  const extraMetadataRef = useRef({
+    credentials,
+    headers,
+    body,
+  });
+
+  useEffect(() => {
+    extraMetadataRef.current = {
+      credentials,
+      headers,
+      body,
+    };
+  }, [credentials, headers, body]);
+
+  const triggerRequest = useCallback(
+    async (prompt: string, options?: CompletionRequestOptions) =>
+      callCompletionApi({
+        api,
+        prompt,
+        credentials: extraMetadataRef.current.credentials,
+        headers: { ...extraMetadataRef.current.headers, ...options?.headers },
+        body: {
+          ...extraMetadataRef.current.body,
+          ...options?.body,
+        },
+        streamProtocol,
+        fetch,
+        // throttle streamed ui updates:
+        setCompletion: throttle(
+          (completion: string) => mutate(completion, false),
+          throttleWaitMs,
+        ),
+        setLoading: mutateLoading,
+        setError,
+        setAbortController,
+        onFinish,
+        onError,
+      }),
+    [
+      mutate,
+      mutateLoading,
+      api,
+      extraMetadataRef,
+      setAbortController,
+      onFinish,
+      onError,
+      setError,
+      streamProtocol,
+      fetch,
+      throttleWaitMs,
+    ],
+  );
+
+  const stop = useCallback(() => {
+    if (abortController) {
+      abortController.abort();
+      setAbortController(null);
+    }
+  }, [abortController]);
+
+  const setCompletion = useCallback(
+    (completion: string) => {
+      mutate(completion, false);
+    },
+    [mutate],
+  );
+
+  const complete = useCallback<UseCompletionHelpers['complete']>(
+    async (prompt, options) => {
+      return triggerRequest(prompt, options);
+    },
+    [triggerRequest],
+  );
+
+  const [input, setInput] = useState(initialInput);
+
+  const handleSubmit = useCallback(
+    (event?: { preventDefault?: () => void }) => {
+      event?.preventDefault?.();
+      return input ? complete(input) : undefined;
+    },
+    [input, complete],
+  );
+
+  const handleInputChange = useCallback(
+    (e: any) => {
+      setInput(e.target.value);
+    },
+    [setInput],
+  );
+
+  return {
+    completion,
+    complete,
+    error,
+    setCompletion,
+    stop,
+    input,
+    setInput,
+    handleInputChange,
+    handleSubmit,
+    isLoading,
+  };
+}

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

@@ -0,0 +1,165 @@
+import {
+  createTestServer,
+  TestResponseController,
+} from '@ai-sdk/test-server/with-vitest';
+import '@testing-library/jest-dom/vitest';
+import { screen, waitFor } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import { UIMessageChunk } from 'ai';
+import { setupTestComponent } from './setup-test-component';
+import { useCompletion } from './use-completion';
+import { describe, it, expect, beforeEach } from 'vitest';
+
+function formatChunk(part: UIMessageChunk) {
+  return `data: ${JSON.stringify(part)}\n\n`;
+}
+
+const server = createTestServer({
+  '/api/completion': {},
+});
+
+describe('stream data stream', () => {
+  let onFinishResult: { prompt: string; completion: string } | undefined;
+
+  setupTestComponent(() => {
+    const {
+      completion,
+      handleSubmit,
+      error,
+      handleInputChange,
+      input,
+      isLoading,
+    } = useCompletion({
+      onFinish(prompt, completion) {
+        onFinishResult = { prompt, completion };
+      },
+    });
+
+    return (
+      <div>
+        <div data-testid="loading">{isLoading.toString()}</div>
+        <div data-testid="error">{error?.toString()}</div>
+        <div data-testid="completion">{completion}</div>
+        <form onSubmit={handleSubmit}>
+          <input
+            data-testid="input"
+            value={input}
+            placeholder="Say something..."
+            onChange={handleInputChange}
+          />
+        </form>
+      </div>
+    );
+  });
+
+  beforeEach(() => {
+    onFinishResult = undefined;
+  });
+
+  describe('render simple stream', () => {
+    beforeEach(async () => {
+      server.urls['/api/completion'].response = {
+        type: 'stream-chunks',
+        chunks: [
+          formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }),
+          formatChunk({ type: 'text-delta', id: '0', delta: ',' }),
+          formatChunk({ type: 'text-delta', id: '0', delta: ' world' }),
+          formatChunk({ type: 'text-delta', id: '0', delta: '.' }),
+        ],
+      };
+      await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
+    });
+
+    it('should render stream', async () => {
+      await waitFor(() => {
+        expect(screen.getByTestId('completion')).toHaveTextContent(
+          'Hello, world.',
+        );
+      });
+    });
+
+    it("should call 'onFinish' callback", async () => {
+      await waitFor(() => {
+        expect(onFinishResult).toEqual({
+          prompt: 'hi',
+          completion: 'Hello, world.',
+        });
+      });
+    });
+  });
+
+  describe('loading state', () => {
+    it('should show loading state', async () => {
+      const controller = new TestResponseController();
+
+      server.urls['/api/completion'].response = {
+        type: 'controlled-stream',
+        controller,
+      };
+
+      await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
+
+      controller.write(
+        formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }),
+      );
+
+      await waitFor(() => {
+        expect(screen.getByTestId('loading')).toHaveTextContent('true');
+      });
+
+      await controller.close();
+
+      await waitFor(() => {
+        expect(screen.getByTestId('loading')).toHaveTextContent('false');
+      });
+    });
+
+    it('should reset loading state on error', async () => {
+      server.urls['/api/completion'].response = {
+        type: 'error',
+        status: 404,
+        body: 'Not found',
+      };
+
+      await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
+
+      await screen.findByTestId('loading');
+      expect(screen.getByTestId('loading')).toHaveTextContent('false');
+    });
+  });
+});
+
+describe('text stream', () => {
+  setupTestComponent(() => {
+    const { completion, handleSubmit, handleInputChange, input } =
+      useCompletion({ streamProtocol: 'text' });
+
+    return (
+      <div>
+        <div data-testid="completion-text-stream">{completion}</div>
+        <form onSubmit={handleSubmit}>
+          <input
+            data-testid="input-text-stream"
+            value={input}
+            placeholder="Say something..."
+            onChange={handleInputChange}
+          />
+        </form>
+      </div>
+    );
+  });
+
+  it('should render stream', async () => {
+    server.urls['/api/completion'].response = {
+      type: 'stream-chunks',
+      chunks: ['Hello', ',', ' world', '.'],
+    };
+
+    await userEvent.type(screen.getByTestId('input-text-stream'), 'hi{enter}');
+
+    await screen.findByTestId('completion-text-stream');
+    expect(screen.getByTestId('completion-text-stream')).toHaveTextContent(
+      'Hello, world.',
+    );
+  });
+});

package/src/use-object.ts

@@ -0,0 +1,269 @@
+import {
+  FetchFunction,
+  FlexibleSchema,
+  InferSchema,
+  isAbortError,
+  Resolvable,
+  resolve,
+  normalizeHeaders,
+  safeValidateTypes,
+} from '@ai-sdk/provider-utils';
+import { asSchema, DeepPartial, isDeepEqualData, parsePartialJson } from 'ai';
+import { useCallback, useId, useRef, useState } from 'react';
+import useSWR from 'swr';
+
+// use function to allow for mocking in tests:
+const getOriginalFetch = () => fetch;
+
+export type Experimental_UseObjectOptions<
+  SCHEMA extends FlexibleSchema,
+  RESULT,
+> = {
+  /**
+   * The API endpoint. It should stream JSON that matches the schema as chunked text.
+   */
+  api: string;
+
+  /**
+   * A schema that defines the shape of the complete object.
+   */
+  schema: SCHEMA;
+
+  /**
+   * An unique identifier. If not provided, a random one will be
+   * generated. When provided, the `useObject` hook with the same `id` will
+   * have shared states across components.
+   */
+  id?: string;
+
+  /**
+   * An optional value for the initial object.
+   */
+  initialValue?: DeepPartial<RESULT>;
+
+  /**
+Custom fetch implementation. You can use it as a middleware to intercept requests,
+or to provide a custom fetch implementation for e.g. testing.
+    */
+  fetch?: FetchFunction;
+
+  /**
+Callback that is called when the stream has finished.
+     */
+  onFinish?: (event: {
+    /**
+The generated object (typed according to the schema).
+Can be undefined if the final object does not match the schema.
+   */
+    object: RESULT | undefined;
+
+    /**
+Optional error object. This is e.g. a TypeValidationError when the final object does not match the schema.
+ */
+    error: Error | undefined;
+  }) => Promise<void> | void;
+
+  /**
+   * Callback function to be called when an error is encountered.
+   */
+  onError?: (error: Error) => void;
+
+  /**
+   * Additional HTTP headers to be included in the request.
+   * Can be a static object, a function that returns headers, or an async function
+   * for dynamic auth tokens.
+   */
+  headers?: Resolvable<Record<string, string> | Headers>;
+
+  /**
+   * The credentials mode to be used for the fetch request.
+   * Possible values are: 'omit', 'same-origin', 'include'.
+   * Defaults to 'same-origin'.
+   */
+  credentials?: RequestCredentials;
+};
+
+export type Experimental_UseObjectHelpers<RESULT, INPUT> = {
+  /**
+   * Calls the API with the provided input as JSON body.
+   */
+  submit: (input: INPUT) => void;
+
+  /**
+   * The current value for the generated object. Updated as the API streams JSON chunks.
+   */
+  object: DeepPartial<RESULT> | undefined;
+
+  /**
+   * The error object of the API request if any.
+   */
+  error: Error | undefined;
+
+  /**
+   * Flag that indicates whether an API request is in progress.
+   */
+  isLoading: boolean;
+
+  /**
+   * Abort the current request immediately, keep the current partial object if any.
+   */
+  stop: () => void;
+
+  /**
+   * Clear the object state.
+   */
+  clear: () => void;
+};
+
+function useObject<
+  SCHEMA extends FlexibleSchema,
+  RESULT = InferSchema<SCHEMA>,
+  INPUT = any,
+>({
+  api,
+  id,
+  schema, // required, in the future we will use it for validation
+  initialValue,
+  fetch,
+  onError,
+  onFinish,
+  headers,
+  credentials,
+}: Experimental_UseObjectOptions<
+  SCHEMA,
+  RESULT
+>): Experimental_UseObjectHelpers<RESULT, INPUT> {
+  // Generate an unique id if not provided.
+  const hookId = useId();
+  const completionId = id ?? hookId;
+
+  // Store the completion state in SWR, using the completionId as the key to share states.
+  const { data, mutate } = useSWR<DeepPartial<RESULT>>(
+    [api, completionId],
+    null,
+    { fallbackData: initialValue },
+  );
+
+  const [error, setError] = useState<undefined | Error>(undefined);
+  const [isLoading, setIsLoading] = useState(false);
+
+  // Abort controller to cancel the current API call.
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  const stop = useCallback(() => {
+    try {
+      abortControllerRef.current?.abort();
+    } catch (ignored) {
+    } finally {
+      setIsLoading(false);
+      abortControllerRef.current = null;
+    }
+  }, []);
+
+  const submit = async (input: INPUT) => {
+    try {
+      clearObject();
+
+      setIsLoading(true);
+
+      const abortController = new AbortController();
+      abortControllerRef.current = abortController;
+
+      // Resolve headers at request time (supports async functions for dynamic auth tokens)
+      const resolvedHeaders = await resolve(headers);
+
+      const actualFetch = fetch ?? getOriginalFetch();
+      const response = await actualFetch(api, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          ...normalizeHeaders(resolvedHeaders),
+        },
+        credentials,
+        signal: abortController.signal,
+        body: JSON.stringify(input),
+      });
+
+      if (!response.ok) {
+        throw new Error(
+          (await response.text()) ?? 'Failed to fetch the response.',
+        );
+      }
+
+      if (response.body == null) {
+        throw new Error('The response body is empty.');
+      }
+
+      let accumulatedText = '';
+      let latestObject: DeepPartial<RESULT> | undefined = undefined;
+
+      await response.body.pipeThrough(new TextDecoderStream()).pipeTo(
+        new WritableStream<string>({
+          async write(chunk) {
+            accumulatedText += chunk;
+
+            const { value } = await parsePartialJson(accumulatedText);
+            const currentObject = value as DeepPartial<RESULT>;
+
+            if (!isDeepEqualData(latestObject, currentObject)) {
+              latestObject = currentObject;
+
+              mutate(currentObject);
+            }
+          },
+
+          async close() {
+            setIsLoading(false);
+            abortControllerRef.current = null;
+
+            if (onFinish != null) {
+              const validationResult = await safeValidateTypes({
+                value: latestObject,
+                schema: asSchema(schema),
+              });
+
+              onFinish(
+                validationResult.success
+                  ? { object: validationResult.value, error: undefined }
+                  : { object: undefined, error: validationResult.error },
+              );
+            }
+          },
+        }),
+      );
+    } catch (error) {
+      if (isAbortError(error)) {
+        return;
+      }
+
+      if (onError && error instanceof Error) {
+        onError(error);
+      }
+
+      setIsLoading(false);
+      setError(error instanceof Error ? error : new Error(String(error)));
+    }
+  };
+
+  const clear = () => {
+    stop();
+    clearObject();
+  };
+
+  const clearObject = () => {
+    setError(undefined);
+    setIsLoading(false);
+    mutate(undefined);
+  };
+
+  return {
+    submit,
+    object: data,
+    error,
+    isLoading,
+    stop,
+    clear,
+  };
+}
+
+export const experimental_useObject = useObject;