AI TOOLKIT
Get started
Recipes/React Server Components

Save Messages To Database

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

3 min readrsctool useView source

Sometimes conversations with language models can get interesting and you might want to save the state of so you can revisit it or continue the conversation later. createAI has an experimental callback function called onSetAIState that gets called whenever the AI state changes. You can use this to save the AI state to a file or a database.

Run it locally

$ npm install ai

Client

import { ServerMessage } from './actions';

import { AI } from './ai';

export default function RootLayout({

children,

}: Readonly<{

children: React.ReactNode;

}>) {

// get chat history from database

const history: ServerMessage[] = getChat();

return (

<html lang="en">

<body>

<AI initialAIState={history} initialUIState={[]}>

{children}

</AI>

</body>

</html>

);

}

'use client';

import { useState } from 'react';

import { ClientMessage } from './actions';

import { useActions, useUIState } from '@ai-toolkit/rsc';

import { generateId } from 'ai-toolkit';

// Allow streaming responses up to 30 seconds

export const maxDuration = 30;

export default function Home() {

const [input, setInput] = useState<string>('');

const [conversation, setConversation] = useUIState();

const { continueConversation } = useActions();

return (

<div>

<div>

{conversation.map((message: ClientMessage) => (

<div key={message.id}>

{message.role}: {message.display}

</div>

))}

</div>

<div>

<input

type="text"

value={input}

onChange={event => {

setInput(event.target.value);

}}

/>

<button

onClick={async () => {

setConversation((currentConversation: ClientMessage[]) => [

...currentConversation,

{ id: generateId(), role: 'user', display: input },

]);

const message = await continueConversation(input);

setConversation((currentConversation: ClientMessage[]) => [

...currentConversation,

message,

]);

}}

>

Send Message

</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 // get chat history from database
  9. 9 const history: ServerMessage[] = getChat();
  10. 10 return (
  11. 11 <html lang="en">
  12. 12 <body>
  13. 13 <AI initialAIState={history} 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 } from 'react';
  3. 3import { ClientMessage } from './actions';
  4. 4import { useActions, useUIState } from '@ai-toolkit/rsc';
  5. 5import { generateId } from 'ai-toolkit';
  6. 6// Allow streaming responses up to 30 seconds
  7. 7export const maxDuration = 30;
  8. 8export default function Home() {
  9. 9 const [input, setInput] = useState<string>('');
  10. 10 const [conversation, setConversation] = useUIState();
  11. 11 const { continueConversation } = useActions();
  12. 12 return (
  13. 13 <div>
  14. 14 <div>
  15. 15 {conversation.map((message: ClientMessage) => (
  16. 16 <div key={message.id}>
  17. 17 {message.role}: {message.display}
  18. 18 </div>
  19. 19 ))}
  20. 20 </div>
  21. 21 <div>
  22. 22 <input
  23. 23 type="text"
  24. 24 value={input}
  25. 25 onChange={event => {
  26. 26 setInput(event.target.value);
  27. 27 }}
  28. 28 />
  29. 29 <button
  30. 30 onClick={async () => {
  31. 31 setConversation((currentConversation: ClientMessage[]) => [
  32. 32 ...currentConversation,
  33. 33 { id: generateId(), role: 'user', display: input },
  34. 34 ]);
  35. 35 const message = await continueConversation(input);
  36. 36 setConversation((currentConversation: ClientMessage[]) => [
  37. 37 ...currentConversation,
  38. 38 message,
  39. 39 ]);
  40. 40 }}
  41. 41 >
  42. 42 Send Message
  43. 43 </button>
  44. 44 </div>
  45. 45 </div>
  46. 46 );
  47. 47}

Server

We will use the callback function to listen to state changes and save the conversation once we receive a done event.

'use server';

import { getAIState, getMutableAIState, streamUI } from '@ai-toolkit/rsc';

import { openai } from '@ai-toolkit/openai';

import { ReactNode } from 'react';

import { z } from 'zod';

import { generateId } from 'ai-toolkit';

import { Stock } from '@ai-studio/components/stock';

export interface ServerMessage {

role: 'user' | 'assistant' | 'function';

content: string;

}

export interface ClientMessage {

id: string;

role: 'user' | 'assistant' | 'function';

display: ReactNode;

}

export async function continueConversation(

input: string,

): Promise<ClientMessage> {

'use server';

const history = getMutableAIState();

const result = await streamUI({

model: openai('gpt-3.5-turbo'),

messages: [...history.get(), { role: 'user', content: input }],

text: ({ content, done }) => {

if (done) {

history.done([

...history.get(),

{ role: 'user', content: input },

{ role: 'assistant', content },

]);

}

return <div>{content}</div>;

},

tools: {

showStockInformation: {

description:

'Get stock information for symbol for the last numOfMonths months',

inputSchema: z.object({

symbol: z

.string()

.describe('The stock symbol to get information for'),

numOfMonths: z

.number()

.describe('The number of months to get historical information for'),

}),

generate: async ({ symbol, numOfMonths }) => {

history.done([

...history.get(),

{

role: 'function',

name: 'showStockInformation',

content: JSON.stringify({ symbol, numOfMonths }),

},

]);

return <Stock symbol={symbol} numOfMonths={numOfMonths} />;

},

},

},

});

return {

id: generateId(),

role: 'assistant',

display: result.value,

};

}

import { createAI } from '@ai-toolkit/rsc';

import { ServerMessage, ClientMessage, continueConversation } from './actions';

export const AI = createAI<ServerMessage[], ClientMessage[]>({

actions: {

continueConversation,

},

onSetAIState: async ({ state, done }) => {

'use server';

if (done) {

saveChat(state);

}

},

onGetUIState: async () => {

'use server';

const history: ServerMessage[] = getAIState();

return history.map(({ role, content }) => ({

id: generateId(),

role,

display:

role === 'function' ? <Stock {...JSON.parse(content)} /> : content,

}));

},

});

app/actions.tsx
tsx
  1. 1'use server';
  2. 2import { getAIState, getMutableAIState, streamUI } from '@ai-toolkit/rsc';
  3. 3import { openai } from '@ai-toolkit/openai';
  4. 4import { ReactNode } from 'react';
  5. 5import { z } from 'zod';
  6. 6import { generateId } from 'ai-toolkit';
  7. 7import { Stock } from '@ai-studio/components/stock';
  8. 8export interface ServerMessage {
  9. 9 role: 'user' | 'assistant' | 'function';
  10. 10 content: string;
  11. 11}
  12. 12export interface ClientMessage {
  13. 13 id: string;
  14. 14 role: 'user' | 'assistant' | 'function';
  15. 15 display: ReactNode;
  16. 16}
  17. 17export async function continueConversation(
  18. 18 input: string,
  19. 19): Promise<ClientMessage> {
  20. 20 'use server';
  21. 21 const history = getMutableAIState();
  22. 22 const result = await streamUI({
  23. 23 model: openai('gpt-3.5-turbo'),
  24. 24 messages: [...history.get(), { role: 'user', content: input }],
  25. 25 text: ({ content, done }) => {
  26. 26 if (done) {
  27. 27 history.done([
  28. 28 ...history.get(),
  29. 29 { role: 'user', content: input },
  30. 30 { role: 'assistant', content },
  31. 31 ]);
  32. 32 }
  33. 33 return <div>{content}</div>;
  34. 34 },
  35. 35 tools: {
  36. 36 showStockInformation: {
  37. 37 description:
  38. 38 'Get stock information for symbol for the last numOfMonths months',
  39. 39 inputSchema: z.object({
  40. 40 symbol: z
  41. 41 .string()
  42. 42 .describe('The stock symbol to get information for'),
  43. 43 numOfMonths: z
  44. 44 .number()
  45. 45 .describe('The number of months to get historical information for'),
  46. 46 }),
  47. 47 generate: async ({ symbol, numOfMonths }) => {
  48. 48 history.done([
  49. 49 ...history.get(),
  50. 50 {
  51. 51 role: 'function',
  52. 52 name: 'showStockInformation',
  53. 53 content: JSON.stringify({ symbol, numOfMonths }),
  54. 54 },
  55. 55 ]);
  56. 56 return <Stock symbol={symbol} numOfMonths={numOfMonths} />;
  57. 57 },
  58. 58 },
  59. 59 },
  60. 60 });
  61. 61 return {
  62. 62 id: generateId(),
  63. 63 role: 'assistant',
  64. 64 display: result.value,
  65. 65 };
  66. 66}
app/ai.ts
ts
  1. 1import { createAI } from '@ai-toolkit/rsc';
  2. 2import { ServerMessage, ClientMessage, continueConversation } from './actions';
  3. 3export const AI = createAI<ServerMessage[], ClientMessage[]>({
  4. 4 actions: {
  5. 5 continueConversation,
  6. 6 },
  7. 7 onSetAIState: async ({ state, done }) => {
  8. 8 'use server';
  9. 9 if (done) {
  10. 10 saveChat(state);
  11. 11 }
  12. 12 },
  13. 13 onGetUIState: async () => {
  14. 14 'use server';
  15. 15 const history: ServerMessage[] = getAIState();
  16. 16 return history.map(({ role, content }) => ({
  17. 17 id: generateId(),
  18. 18 role,
  19. 19 display:
  20. 20 role === 'function' ? <Stock {...JSON.parse(content)} /> : content,
  21. 21 }));
  22. 22 },
  23. 23});
Browse all recipes