improve api and add docs

README.md

@@ -277,7 +277,7 @@ export async function POST() {
 
 ### `useChat(options: UseChatOptions): ChatHelpers`
 
-An SWR-powered React hook for streaming text completion or chat messages and handling chat and prompt input state.
+An SWR-powered React hook for streaming chat messages and handling chat and prompt input state.
 
 The `useChat` hook is designed to provide an intuitive interface for building ChatGPT-like UI's in React with streaming text responses. It leverages the [SWR](https://swr.vercel.app) library for efficient data fetching and state synchronization.
 
@@ -306,6 +306,10 @@ type UseChatOptions = {
   id?: string
   initialMessages?: Message[]
   initialInput?: string
+  onResponse?: (response: Response) => void
+  onFinish?: (message: Message) => void
+  headers?: Record<string, string> | Headers
+  body?: any
 }
 ```
 
@@ -383,3 +387,87 @@ export default function Chat() {
 ```
 
 In this example, chat is an object of type `UseChatHelpers`, which contains various utilities to interact with and control the chat. You can use these utilities to render chat messages, handle input changes, submit messages, and manage the chat state in your UI.
+
+### `useCompletion(options: UseCompletionOptions): CompletionHelpers`
+
+An SWR-powered React hook for streaming text completion and handling completion and prompt input state.
+
+The `useCompletion` hook is designed to provide an intuitive interface for building a UI that streams text completions in React. It leverages the [SWR](https://swr.vercel.app) library for efficient data fetching and state synchronization.
+
+#### Types
+
+#### `UseCompletionOptions`
+
+The `UseCompletionOptions` type defines the configuration options for the `useCompletion` hook.
+
+```tsx
+type UseCompletionOptions = {
+  api?: string
+  id?: string
+  initialInput?: string
+  initialCompletion?: string
+  onResponse?: (response: Response) => void
+  onFinish?: (prompt: string, completion: string) => void
+  headers?: Record<string, string> | Headers
+  body?: any
+}
+```
+
+##### `CompletionHelpers`
+
+The `CompletionHelpers` type is the return type of the `useCompletion` hook. It provides various utilities to interact with and manipulate the completion.
+
+```tsx
+type UseChatHelpers = {
+  completion: string
+  complete: (prompt: string) => void
+  error: any
+  set: (completion: string) => void
+  stop: () => void
+  input: string
+  setInput: react.Dispatch<react.SetStateAction<string>>
+  handleInputChange: (e: any) => void
+  handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void
+  isLoading: boolean
+}
+```
+
+#### Example
+
+Below is a basic example of the `useCompletion` hook in a component:
+
+```tsx
+// app/completion.tsx
+'use client'
+
+import { useCompletion } from '@vercel/ai-utils'
+
+export default function Completion() {
+  const { completion, input, stop, isLoading, handleInputChange, handleSubmit } =
+    useCompletion({
+      api: '/api/some-custom-endpoint'
+    })
+
+  return (
+    <div className="mx-auto w-full max-w-md py-24 flex flex-col stretch">
+      <form onSubmit={handleSubmit}>
+        <input
+          className="fixed w-full max-w-md bottom-0 border border-gray-300 rounded mb-8 shadow-xl p-2"
+          value={input}
+          placeholder="Say something..."
+          onChange={handleInputChange}
+        />
+        <p>Completion result: {completion}</p>
+        <button type="button" onClick={stop}>
+          Stop
+        </button>
+        <button disabled={isLoading} type="submit">
+          Send
+        </button>
+      </form>
+    </div>
+  )
+}
+```
+
+In this example, completion is an object of type `CompletionHelpers`, which contains various utilities to interact with and control the completion. You can use these utilities to render the completion, handle input changes, submit prompts, and manage the completion state in your UI.

apps/docs/pages/api-reference.mdx

@@ -6,7 +6,7 @@ title: API Reference
 
 ## `useChat(options: UseChatOptions): ChatHelpers`
 
-An SWR-powered React hook for streaming text completion or chat messages and handling chat and prompt input state.
+An SWR-powered React hook for streaming chat messages and handling chat and prompt input state.
 
 ```tsx filename="app/chat.tsx"
 'use client'
@@ -46,6 +46,58 @@ export default function Chat() {
 - **`id?: string`** - An unique identifier for the chat. If not provided, a random one will be generated. When provided, the `useChat` hook with the same `id` will have shared states across components thanks to SWR.
 - **`initialInput?: string = ''`** - An optional string of initial prompt input
 - **`initialMessages?: Messages[] = []`** - An optional array of initial chat messages
+- **`onResponse?: (res: Response) => void`** - An optional callback that will be called with the response from the API endpoint. Useful for throwing customized errors or logging
+- **`onFinish?: (message: Message) => void`** - An optional callback that will be called when the chat stream ends
+- **`headers?: Record<string, string> | Headers`** - An optional object of headers to be passed to the API endpoint
+- **`body?: any`** - An optional, extra body object to be passed to the API endpoint
+
+## `useCompletion(options: UseCompletionOptions): CompletionHelpers`
+
+An SWR-powered React hook for streaming text completion and handling prompt input state.
+
+```tsx filename="app/completion.tsx"
+'use client'
+
+import { useCompletion } from '@vercel/ai-utils'
+
+export default function Completion() {
+  const { completion, input, stop, isLoading, handleInputChange, handleSubmit } =
+    useCompletion({
+      api: '/api/some-custom-endpoint'
+    })
+
+  return (
+    <div className="mx-auto w-full max-w-md py-24 flex flex-col stretch">
+      <form onSubmit={handleSubmit}>
+        <input
+          className="fixed w-full max-w-md bottom-0 border border-gray-300 rounded mb-8 shadow-xl p-2"
+          value={input}
+          placeholder="Say something..."
+          onChange={handleInputChange}
+        />
+        <p>Completion result: {completion}</p>
+        <button type="button" onClick={stop}>
+          Stop
+        </button>
+        <button disabled={isLoading} type="submit">
+          Send
+        </button>
+      </form>
+    </div>
+  )
+}
+```
+
+### `UseCompletionOptions`
+
+- **`api?: string = '/api/completion'`** - The API endpoint that accepts a `{ prompt: string }` object and returns a stream of tokens of the AI completion response. Defaults to `/api/completion`
+- **`id?: string`** - An unique identifier for the completion. If not provided, a random one will be generated. When provided, the `useCompletion` hook with the same `id` will have shared states across components thanks to SWR
+- **`initialInput?: string = ''`** - An optional string of initial prompt input
+- **`initialCompletion?: string = ''`** - An optional string of initial completion result
+- **`onResponse?: (res: Response) => void`** - An optional callback that will be called with the response from the API endpoint. Useful for throwing customized errors or logging
+- **`onFinish?: (prompt: string, completion: string) => void`** - An optional callback that will be called when the completion stream ends
+- **`headers?: Record<string, string> | Headers`** - An optional object of headers to be passed to the API endpoint
+- **`body?: any`** - An optional, extra body object to be passed to the API endpoint
 
 ## `OpenAIStream(res: Response, cb: AIStreamCallbacks): ReadableStream`
 

packages/core/src/use-chat.ts

@@ -37,12 +37,14 @@ export type UseChatOptions = {
    * a stream of tokens of the AI chat response. Defaults to `/api/chat`.
    */
   api?: string
+
   /**
    * An unique identifier for the chat. If not provided, a random one will be
    * generated. When provided, the `useChat` hook with the same `id` will
    * have shared states across components.
    */
   id?: string
+
   /**
    * Initial messages of the chat. Useful to load an existing chat history.
    */
@@ -52,6 +54,26 @@ export type UseChatOptions = {
    * Initial input of the chat.
    */
   initialInput?: string
+
+  /**
+   * Callback function to be called when the API response is received.
+   */
+  onResponse?: (response: Response) => void
+
+  /**
+   * Callback function to be called when the chat is finished streaming.
+   */
+  onFinish?: (message: Message) => void
+
+  /**
+   * HTTP headers to be sent with the API request.
+   */
+  headers?: Record<string, string> | Headers
+
+  /**
+   * Extra body to be sent with the API request.
+   */
+  body?: any
 }
 
 export type UseChatHelpers = {
@@ -85,18 +107,16 @@ export type UseChatHelpers = {
   isLoading: boolean
 }
 
-export function useChat(
-  {
-    api = '/api/chat',
-    id,
-    initialMessages = [],
-    initialInput = ''
-  }: UseChatOptions = {
-    api: '/api/chat',
-    initialMessages: [],
-    initialInput: ''
-  }
-): UseChatHelpers {
+export function useChat({
+  api = '/api/chat',
+  id,
+  initialMessages = [],
+  initialInput = '',
+  onResponse,
+  onFinish,
+  headers,
+  body
+}: UseChatOptions): UseChatHelpers {
   // Generate an unique id for the chat if not provided.
   const hookId = useId()
   const chatId = id || hookId
@@ -117,6 +137,17 @@ export function useChat(
   const [abortController, setAbortController] =
     useState<AbortController | null>(null)
 
+  const extraMetadataRef = useRef<any>({
+    headers,
+    body
+  })
+  useEffect(() => {
+    extraMetadataRef.current = {
+      headers,
+      body
+    }
+  }, [headers, body])
+
   // Actual mutation hook to send messages to the API endpoint and update the
   // chat state.
   const { error, trigger, isMutating } = useSWRMutation<
@@ -139,15 +170,25 @@ export function useChat(
         const res = await fetch(api, {
           method: 'POST',
           body: JSON.stringify({
-            messages: messagesSnapshot
+            messages: messagesSnapshot,
+            ...extraMetadataRef.current.body
           }),
+          headers: extraMetadataRef.current.headers || {},
           signal: abortController.signal
         }).catch(err => {
           // Restore the previous messages if the request fails.
           mutate(previousMessages, false)
           throw err
         })
 
+        if (onResponse) {
+          try {
+            await onResponse(res)
+          } catch (err) {
+            throw err
+          }
+        }
+
         if (!res.ok) {
           // Restore the previous messages if the request fails.
           mutate(previousMessages, false)
@@ -184,6 +225,15 @@ export function useChat(
           )
         }
 
+        if (onFinish) {
+          onFinish({
+            id: replyId,
+            createdAt,
+            content: result,
+            role: 'assistant'
+          })
+        }
+
         setAbortController(null)
         return null
       } catch (err) {

packages/core/src/use-completion.ts

@@ -22,22 +22,31 @@ export type UseCompletionOptions = {
    * have shared states across components.
    */
   id?: string
+
+  /**
+   * Initial prompt input of the completion.
+   */
+  initialInput?: string
+
   /**
    * Initial completion result. Useful to load an existing history.
    */
   initialCompletion?: string
 
   /**
-   * Initial prompt input of the completion.
+   * Callback function to be called when the API response is received.
    */
-  initialInput?: string
+  onResponse?: (response: Response) => void
+  /**
+   * Callback function to be called when the completion is finished streaming.
+   */
+  onFinish?: (prompt: string, completion: string) => void
 
   /**
    * HTTP headers to be sent with the API request.
    */
   headers?: Record<string, string> | Headers
-  onResponse?: (response: Response) => void
-  onCompletionEnd?: (prompt: string, completion: string) => void
+
   /**
    * Extra body to be sent with the API request.
    */
@@ -52,7 +61,7 @@ export function useCompletion({
   headers,
   body,
   onResponse,
-  onCompletionEnd
+  onFinish
 }: UseCompletionOptions) {
   // Generate an unique id for the completion if not provided.
   const hookId = useId()
@@ -137,8 +146,8 @@ export function useCompletion({
           mutate(result, false)
         }
 
-        if (onCompletionEnd) {
-          onCompletionEnd(prompt, result)
+        if (onFinish) {
+          onFinish(prompt, result)
         }
 
         setAbortController(null)