Restore Messages From Database
Learn how to restore messages from an external database using the AI TOOLKIT and React Server Components
When building AI applications, you might want to restore previous conversations from a database to allow users to continue their conversations or review past interactions. The AI TOOLKIT provides mechanisms to restore conversation state through initialAIState and onGetUIState.
Run it locally
$ npm install aiClient
import { ServerMessage } from './actions';
import { AI } from './ai';
export default function RootLayout({
children,
}: Readonly<{
children: React.ReactNode;
}>) {
// Fetch stored messages from your database
const savedMessages: ServerMessage[] = getSavedMessages();
return (
<html lang="en">
<body>
<AI initialAIState={savedMessages} initialUIState={[]}>
{children}
</AI>
</body>
</html>
);
}
'use client';
import { useState, useEffect } from 'react';
import { ClientMessage } from './actions';
import { useActions, useUIState } from '@ai-toolkit/rsc';
import { generateId } from 'ai-toolkit';
export default function Home() {
const [conversation, setConversation] = useUIState();
const [input, setInput] = useState<string>('');
const { continueConversation } = useActions();
return (
<div>
<div className="conversation-history">
{conversation.map((message: ClientMessage) => (
<div key={message.id} className={message ${message.role}}>
{message.role}: {message.display}
</div>
))}
</div>
<div className="input-area">
<input
type="text"
value={input}
onChange={e => setInput(e.target.value)}
placeholder="Type your message..."
/>
<button
onClick={async () => {
// Add user message to UI
setConversation((currentConversation: ClientMessage[]) => [
...currentConversation,
{ id: generateId(), role: 'user', display: input },
]);
// Get AI response
const message = await continueConversation(input);
// Add AI response to UI
setConversation((currentConversation: ClientMessage[]) => [
...currentConversation,
message,
]);
setInput('');
}}
>
Send
</button>
</div>
</div>
);
}
- 1import { ServerMessage } from './actions';
- 2import { AI } from './ai';
- 3export default function RootLayout({
- 4 children,
- 5}: Readonly<{
- 6 children: React.ReactNode;
- 7}>) {
- 8 // Fetch stored messages from your database
- 9 const savedMessages: ServerMessage[] = getSavedMessages();
- 10 return (
- 11 <html lang="en">
- 12 <body>
- 13 <AI initialAIState={savedMessages} initialUIState={[]}>
- 14 {children}
- 15 </AI>
- 16 </body>
- 17 </html>
- 18 );
- 19}
- 1'use client';
- 2import { useState, useEffect } from 'react';
- 3import { ClientMessage } from './actions';
- 4import { useActions, useUIState } from '@ai-toolkit/rsc';
- 5import { generateId } from 'ai-toolkit';
- 6export default function Home() {
- 7 const [conversation, setConversation] = useUIState();
- 8 const [input, setInput] = useState<string>('');
- 9 const { continueConversation } = useActions();
- 10 return (
- 11 <div>
- 12 <div className="conversation-history">
- 13 {conversation.map((message: ClientMessage) => (
- 14 <div key={message.id} className={`message ${message.role}`}>
- 15 {message.role}: {message.display}
- 16 </div>
- 17 ))}
- 18 </div>
- 19 <div className="input-area">
- 20 <input
- 21 type="text"
- 22 value={input}
- 23 onChange={e => setInput(e.target.value)}
- 24 placeholder="Type your message..."
- 25 />
- 26 <button
- 27 onClick={async () => {
- 28 // Add user message to UI
- 29 setConversation((currentConversation: ClientMessage[]) => [
- 30 ...currentConversation,
- 31 { id: generateId(), role: 'user', display: input },
- 32 ]);
- 33 // Get AI response
- 34 const message = await continueConversation(input);
- 35 // Add AI response to UI
- 36 setConversation((currentConversation: ClientMessage[]) => [
- 37 ...currentConversation,
- 38 message,
- 39 ]);
- 40 setInput('');
- 41 }}
- 42 >
- 43 Send
- 44 </button>
- 45 </div>
- 46 </div>
- 47 );
- 48}
Server
The server-side implementation handles the restoration of messages and their transformation into the appropriate format for display.
import { createAI } from '@ai-toolkit/rsc';
import { ServerMessage, ClientMessage, continueConversation } from './actions';
import { Stock } from '@ai-studio/components/stock';
import { generateId } from 'ai-toolkit';
export const AI = createAI<ServerMessage[], ClientMessage[]>({
actions: {
continueConversation,
},
onGetUIState: async () => {
'use server';
// Get the current AI state (stored messages)
const history: ServerMessage[] = getAIState();
// Transform server messages into client messages
return history.map(({ role, content }) => ({
id: generateId(),
role,
display:
role === 'function' ? <Stock {...JSON.parse(content)} /> : content,
}));
},
});
'use server';
import { getAIState } from '@ai-toolkit/rsc';
export interface ServerMessage {
role: 'user' | 'assistant' | 'function';
content: string;
}
export interface ClientMessage {
id: string;
role: 'user' | 'assistant' | 'function';
display: ReactNode;
}
// Function to get saved messages from database
export async function getSavedMessages(): Promise<ServerMessage[]> {
'use server';
// Implement your database fetching logic here
return await fetchMessagesFromDatabase();
}
- 1import { createAI } from '@ai-toolkit/rsc';
- 2import { ServerMessage, ClientMessage, continueConversation } from './actions';
- 3import { Stock } from '@ai-studio/components/stock';
- 4import { generateId } from 'ai-toolkit';
- 5export const AI = createAI<ServerMessage[], ClientMessage[]>({
- 6 actions: {
- 7 continueConversation,
- 8 },
- 9 onGetUIState: async () => {
- 10 'use server';
- 11 // Get the current AI state (stored messages)
- 12 const history: ServerMessage[] = getAIState();
- 13 // Transform server messages into client messages
- 14 return history.map(({ role, content }) => ({
- 15 id: generateId(),
- 16 role,
- 17 display:
- 18 role === 'function' ? <Stock {...JSON.parse(content)} /> : content,
- 19 }));
- 20 },
- 21});
- 1'use server';
- 2import { getAIState } from '@ai-toolkit/rsc';
- 3export interface ServerMessage {
- 4 role: 'user' | 'assistant' | 'function';
- 5 content: string;
- 6}
- 7export interface ClientMessage {
- 8 id: string;
- 9 role: 'user' | 'assistant' | 'function';
- 10 display: ReactNode;
- 11}
- 12// Function to get saved messages from database
- 13export async function getSavedMessages(): Promise<ServerMessage[]> {
- 14 'use server';
- 15 // Implement your database fetching logic here
- 16 return await fetchMessagesFromDatabase();
- 17}