Restore Messages From Database

Learn how to restore messages from an external database using the AI TOOLKIT and React Server Components

3 min readrsctool useView source

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 ai

Client

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>

);

}

app/layout.tsx
tsx
  1. 1import { ServerMessage } from './actions';
  2. 2import { AI } from './ai';
  3. 3export default function RootLayout({
  4. 4 children,
  5. 5}: Readonly<{
  6. 6 children: React.ReactNode;
  7. 7}>) {
  8. 8 // Fetch stored messages from your database
  9. 9 const savedMessages: ServerMessage[] = getSavedMessages();
  10. 10 return (
  11. 11 <html lang="en">
  12. 12 <body>
  13. 13 <AI initialAIState={savedMessages} initialUIState={[]}>
  14. 14 {children}
  15. 15 </AI>
  16. 16 </body>
  17. 17 </html>
  18. 18 );
  19. 19}
app/page.tsx
tsx
  1. 1'use client';
  2. 2import { useState, useEffect } from 'react';
  3. 3import { ClientMessage } from './actions';
  4. 4import { useActions, useUIState } from '@ai-toolkit/rsc';
  5. 5import { generateId } from 'ai-toolkit';
  6. 6export default function Home() {
  7. 7 const [conversation, setConversation] = useUIState();
  8. 8 const [input, setInput] = useState<string>('');
  9. 9 const { continueConversation } = useActions();
  10. 10 return (
  11. 11 <div>
  12. 12 <div className="conversation-history">
  13. 13 {conversation.map((message: ClientMessage) => (
  14. 14 <div key={message.id} className={`message ${message.role}`}>
  15. 15 {message.role}: {message.display}
  16. 16 </div>
  17. 17 ))}
  18. 18 </div>
  19. 19 <div className="input-area">
  20. 20 <input
  21. 21 type="text"
  22. 22 value={input}
  23. 23 onChange={e => setInput(e.target.value)}
  24. 24 placeholder="Type your message..."
  25. 25 />
  26. 26 <button
  27. 27 onClick={async () => {
  28. 28 // Add user message to UI
  29. 29 setConversation((currentConversation: ClientMessage[]) => [
  30. 30 ...currentConversation,
  31. 31 { id: generateId(), role: 'user', display: input },
  32. 32 ]);
  33. 33 // Get AI response
  34. 34 const message = await continueConversation(input);
  35. 35 // Add AI response to UI
  36. 36 setConversation((currentConversation: ClientMessage[]) => [
  37. 37 ...currentConversation,
  38. 38 message,
  39. 39 ]);
  40. 40 setInput('');
  41. 41 }}
  42. 42 >
  43. 43 Send
  44. 44 </button>
  45. 45 </div>
  46. 46 </div>
  47. 47 );
  48. 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();

}

app/ai.ts
tsx
  1. 1import { createAI } from '@ai-toolkit/rsc';
  2. 2import { ServerMessage, ClientMessage, continueConversation } from './actions';
  3. 3import { Stock } from '@ai-studio/components/stock';
  4. 4import { generateId } from 'ai-toolkit';
  5. 5export const AI = createAI<ServerMessage[], ClientMessage[]>({
  6. 6 actions: {
  7. 7 continueConversation,
  8. 8 },
  9. 9 onGetUIState: async () => {
  10. 10 'use server';
  11. 11 // Get the current AI state (stored messages)
  12. 12 const history: ServerMessage[] = getAIState();
  13. 13 // Transform server messages into client messages
  14. 14 return history.map(({ role, content }) => ({
  15. 15 id: generateId(),
  16. 16 role,
  17. 17 display:
  18. 18 role === 'function' ? <Stock {...JSON.parse(content)} /> : content,
  19. 19 }));
  20. 20 },
  21. 21});
app/actions.tsx
tsx
  1. 1'use server';
  2. 2import { getAIState } from '@ai-toolkit/rsc';
  3. 3export interface ServerMessage {
  4. 4 role: 'user' | 'assistant' | 'function';
  5. 5 content: string;
  6. 6}
  7. 7export interface ClientMessage {
  8. 8 id: string;
  9. 9 role: 'user' | 'assistant' | 'function';
  10. 10 display: ReactNode;
  11. 11}
  12. 12// Function to get saved messages from database
  13. 13export async function getSavedMessages(): Promise<ServerMessage[]> {
  14. 14 'use server';
  15. 15 // Implement your database fetching logic here
  16. 16 return await fetchMessagesFromDatabase();
  17. 17}