AI SDK RSCAI and UI States

What is AI and UI State?

There is a recurring pattern of having a state for the language model on the server, and a state for the UI on the client. However, it can get tricky to manage these two states separately and keep them in sync.

For example, in a chat application, the AI state is usually the conversation history (messages) between the user and the assistant. The LLM reads this history so it can generate the next message. This value is also the source of truth for the current application state. On the client side, the UI state is the list of actual UI elements that are displayed to the user.

In traditional text-based chat applications, the displayed UI elements are just plain text. They are identical to the conversation history passed to the AI. However with Generative UI, what gets presented to the user can be more complex, dynamic and hybrid than just plain text. This is where the AI and UI states can diverge.

The RSC API gives you the flexibility to manage these two states separately, while providing a way to keep them in sync between your database, server and client.

AI State

It contains the context the language model needs to read. For a chat app, the AI state is generally the conversation history between the user and the assistant. In practice, it can also be used to store other values and meta information such as createdAt for each message and chatId for each conversation. The AI state, can be accessed/modified from both the server and the client.

UI State

It contains the generated UI and other information for the client-side of the application, that is displayed to the user. It is a fully client-side state (similar to useState) that can store anything from JS values to React elements. It can't be accessed from the server-side.

Creating the AI Context

To start, you need create an React context using createAI, and wrap your application with it. Then, the AI and UI states will be available to all its children:

app/actions.tsx
// Define the AI state and UI state types
export type AIState = Array<{
role: 'user' | 'assistant';
content: string;
}>;
export type UIState = Array<{
id: string;
role: 'user' | 'assistant';
display: ReactNode;
}>;
async function sendMessage(message: string) {
'use server';
// Handle the message, covered in the following sections.
}
// Create the AI provider with the initial states and allowed actions
export const AI = createAI({
initialAIState: [] as AIState,
initialUIState: [] as UIState,
actions: {
sendMessage,
},
});

Using the AI Context

The AI context can be used in any Server Component, to wrap the children components that need access to the AI. For example, you can wrap your root layout component with it

app/layout.tsx
import { type ReactNode } from 'react';
import { AI } from './actions';
export default function RootLayout({
children,
}: Readonly<{ children: ReactNode }>) {
return (
<AI>
<html lang="en">
<body>{children}</body>
</html>
</AI>
);
}

Reading UI State in Client

The UI state can be accessed in Client Components using the useUIState hook provided by the RSC API. The hook returns the current UI state and a function to update the UI state like React's useState.

app/page.tsx
'use client';
import { useUIState } from 'ai/rsc';
export default function Page() {
const [messages, setMessages] = useUIState();
return (
<ul>
{messages.map(message => (
<li key={message.id}>{message.display}</li>
))}
</ul>
);
}

Reading AI State in Client

The AI state can be accessed in Client Components using the useAIState hook provided by the RSC API. The hook returns the current AI state.

app/page.tsx
'use client';
import { useAIAtate } from 'ai/rsc';
export default function Page() {
const [messages, setMessages] = useAIState();
return (
<ul>
{messages.map(message => (
<li key={message.id}>{message.content}</li>
))}
</ul>
);
}

Reading AI State on Server

The AI State can be accessed by the Server Action we provided to createAI, using the getAIState function. It returns the current AI state as a read-only value:

app/actions.ts
import { getAIState } from 'ai/rsc';
export async function sendMessage(message: string) {
'use server';
const history = getAIState();
const response = await generateText({
model: openai('gpt-3.5-turbo'),
messages: [...history, { role: 'user', content: message }],
});
return response;
}

Updating AI State on Server

The AI State can also be updated by the Server Action with the getMutableAIState function. Similar to getAIState, but it returns the state with methods to read and update it:

app/actions.ts
import { getMutableAIState } from 'ai/rsc';
export async function sendMessage(message: string) {
'use server';
const history = getMutableAIState();
// Update the AI state with the new user message.
history.update([...history.get(), { role: 'user', content: message }]);
const response = await generateText({
model: openai('gpt-3.5-turbo'),
messages: history.get(),
});
// Update the AI state again with the response from the model.
history.done([...history.get(), { role: 'assistant', content: response }]);
return response;
}

When the sendMessage action is triggered multiple times, the AI state will always be updated with the new messages and responses, keeping the conversation history in sync.

Calling Server Actions from Client

To call the sendMessage action from the client, you can use the useActions hook. The hook returns all the available Actions that were provided to createAI:

app/page.tsx
'use client';
import { useActions, useUIState } from 'ai/rsc';
export default function Page() {
const { sendMessage } = useActions();
const [messages, setMessages] = useUIState();
const handleSubmit = async event => {
event.preventDefault();
setMessages([
...messages,
{ id: Date.now(), role: 'user', display: event.target.message.value },
]);
const response = await sendMessage(event.target.message.value);
setMessages([
...messages,
{ id: Date.now(), role: 'assistant', display: response },
]);
};
return (
<>
<ul>
{messages.map(message => (
<li key={message.id}>{message.display}</li>
))}
</ul>
<form onSubmit={handleSubmit}>
<input type="text" name="message" />
<button type="submit">Send</button>
</form>
</>
);
}

When the user submits a message, the sendMessage action is called with the message content. The response from the action is then added to the UI state, updating the displayed messages.