Managing Generative UI State

State is an essential part of any application. State is particularly important in AI applications as it is passed to large language models (LLMs) on each request to ensure they have the necessary context to produce a great generation. Traditional chatbots are text-based and have a structure that mirrors that of any chat application.

For example, in a chatbot, state is an array of messages where each message has:

• id: a unique identifier
• role: who sent the message (user/assistant/system/tool)
• content: the content of the message

This state can be rendered in the UI and sent to the model without any modifications.

With Generative UI, the model can now return a React component, rather than a plain text message. The client can render that component without issue, but that state can't be sent back to the model because React components aren't serialisable. So, what can you do?

The solution is to split the state in two, where one (AI State) becomes a proxy for the other (UI State).

One way to understand this concept is through a Lego analogy. Imagine a 10,000 piece Lego model that, once built, cannot be easily transported because it is fragile. By taking the model apart, it can be easily transported, and then rebuilt following the steps outlined in the instructions pamphlet. In this way, the instructions pamphlet is a proxy to the physical structure. Similarly, AI State provides a serialisable (JSON) representation of your UI that can be passed back and forth to the model.

What is AI and UI State?

The RSC API simplifies how you manage AI State and UI State, providing a robust way to keep them in sync between your database, server and client.

AI State

AI State refers to the state of your application in a serialisable format that will be used on the server and can be shared with the language model.

For a chat app, the AI State is the conversation history (messages) between the user and the assistant. Components generated by the model would be represented in a JSON format as a tool alongside any necessary props. AI State can also be used to store other values and meta information such as createdAt for each message and chatId for each conversation. The LLM reads this history so it can generate the next message. This state serves as the source of truth for the current application state.

UI State

UI State refers to the state of your application that is rendered on the client. It is a fully client-side state (similar to useState) that can store anything from Javascript values to React elements. UI state is a list of actual UI elements that are rendered on the client.

Using AI / UI State

Creating the AI Context

AI SDK RSC simplifies managing AI and UI state across your application by providing several hooks. These hooks are powered by React context under the hood.

Notably, this means you do not have to pass the message history to the server explicitly for each request. You also can access and update your application state in any child component of the context provider. As you begin building multistep generative interfaces, this will be particularly helpful.

To use ai/rsc to manage AI and UI State in your application, you can create a React context using createAI:

// Define the AI state and UI state types
export type ServerMessage = {
  role: 'user' | 'assistant';
  content: string;
};

export type ClientMessage = {
  id: string;
  role: 'user' | 'assistant';
  display: ReactNode;
};

export const sendMessage = async (input: string): Promise<ClientMessage> => {
  "use server"
  ...
}

import { createAI } from 'ai/rsc';
import { ClientMessage, ServerMessage, sendMessage } from './actions';

export type AIState = ServerMessage[];
export type UIState = ClientMessage[];

// Create the AI provider with the initial states and allowed actions
export const AI = createAI<AIState, UIState>({
  initialAIState: [],
  initialUIState: [],
  actions: {
    sendMessage,
  },
});

In this example, you define types for AI State and UI State, respectively.

Next, wrap your application with your newly created context. With that, you can get and set AI and UI State across your entire application.

import { type ReactNode } from 'react';
import { AI } from './ai';

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.

'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.

'use client';

import { useAIState } 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 within any Server Action provided to the createAI context using the getAIState function. It returns the current AI state as a read-only value:

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 from within your Server Action with the getMutableAIState function. This function is similar to getAIState, but it returns the state with methods to read and update it:

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;
}

Calling Server Actions from the 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:

'use client';

import { useActions, useUIState } from 'ai/rsc';
import { AI } from './ai';

export default function Page() {
  const { sendMessage } = useActions<typeof AI>();
  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.

To learn more, check out this example on managing AI and UI state using ai/rsc.

Next, you will learn how you can save and restore state with ai/rsc.