-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
docs: add memoization recipe (#4116)
- Loading branch information
1 parent
5d0e8aa
commit eb0632d
Showing
1 changed file
with
133 additions
and
0 deletions.
There are no files selected for viewing
133 changes: 133 additions & 0 deletions
133
content/cookbook/01-next/25-markdown-chatbot-with-memoization.mdx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,133 @@ | ||
--- | ||
title: Markdown Chatbot with Memoization | ||
description: Learn how to build a performant chatbot that renders Markdown responses and uses memoization for optimal performance with Next.js and the AI SDK. | ||
tags: ['next', 'streaming', 'chatbot', 'markdown'] | ||
--- | ||
|
||
# Markdown Chatbot with Memoization | ||
|
||
When building a chatbot with Next.js and the AI SDK, you'll likely want to render the model's responses in Markdown format using a library like `react-markdown`. However, this can have negative performance implications as the Markdown is re-rendered on each new token received from the streaming response. | ||
|
||
As conversations get longer and more complex, this performance impact becomes exponentially worse since the entire conversation history is re-rendered with each new token. | ||
|
||
This recipe uses memoization - a performance optimization technique where the results of expensive function calls are cached and reused to avoid unnecessary re-computation. In this case, parsed Markdown blocks are memoized to prevent them from being re-parsed and re-rendered on each token update, which means that once a block is fully parsed, it's cached and reused rather than being regenerated. This approach significantly improves rendering performance for long conversations by eliminating redundant parsing and rendering operations. | ||
|
||
## Server | ||
|
||
On the server, you use a simple route handler that streams the response from the language model. | ||
|
||
```tsx filename='app/api/chat/route.ts' | ||
import { openai } from '@ai-sdk/openai'; | ||
import { streamText } from 'ai'; | ||
|
||
export const maxDuration = 60; | ||
|
||
export async function POST(req: Request) { | ||
const { messages } = await req.json(); | ||
|
||
const result = streamText({ | ||
system: | ||
'You are a helpful assistant. Respond to the user in Markdown format.', | ||
model: openai('gpt-4o'), | ||
messages, | ||
}); | ||
|
||
return result.toDataStreamResponse(); | ||
} | ||
``` | ||
|
||
## Memoized Markdown Component | ||
|
||
Next, create a memoized markdown component that will take in raw Markdown text into blocks and only updates when the content actually changes. This component splits Markdown content into blocks using the `marked` library to identify discrete Markdown elements, then uses React's memoization features to optimize re-rendering by only updating blocks that have actually changed. | ||
|
||
```tsx filename='components/memoized-markdown.tsx' | ||
import { marked } from 'marked'; | ||
import { memo, useMemo } from 'react'; | ||
import ReactMarkdown from 'react-markdown'; | ||
|
||
function parseMarkdownIntoBlocks(markdown: string): string[] { | ||
const tokens = marked.lexer(markdown); | ||
return tokens.map(token => token.raw); | ||
} | ||
|
||
const MemoizedMarkdownBlock = memo( | ||
({ content }: { content: string }) => { | ||
return <ReactMarkdown>{content}</ReactMarkdown>; | ||
}, | ||
(prevProps, nextProps) => { | ||
if (prevProps.content !== nextProps.content) return false; | ||
return true; | ||
}, | ||
); | ||
|
||
MemoizedMarkdownBlock.displayName = 'MemoizedMarkdownBlock'; | ||
|
||
export const MemoizedMarkdown = memo( | ||
({ content, id }: { content: string; id: string }) => { | ||
const blocks = useMemo(() => parseMarkdownIntoBlocks(content), [content]); | ||
|
||
return blocks.map((block, index) => ( | ||
<MemoizedMarkdownBlock content={block} key={`${id}-block_${index}`} /> | ||
)); | ||
}, | ||
); | ||
|
||
MemoizedMarkdown.displayName = 'MemoizedMarkdown'; | ||
``` | ||
|
||
## Client | ||
|
||
Finally, on the client, use the `useChat` hook to manage the chat state and render the chat interface. You can use the `MemoizedMarkdown` component to render the message contents in Markdown format without compromising on performance. Additionally, you can render the form in it's own component so as to not trigger unnecessary re-renders of the chat messages. You can also use the `experimental_throttle` option that will throttle data updates to a specified interval, helping to manage rendering performance. | ||
|
||
```typescript filename='app/page.tsx' | ||
'use client'; | ||
|
||
import { useChat } from 'ai/react'; | ||
import { MemoizedMarkdown } from '@/components/memoized-markdown'; | ||
|
||
export default function Page() { | ||
const { messages } = useChat({ | ||
id: 'chat', | ||
// Throttle the messages and data updates to 50ms: | ||
experimental_throttle: 50, | ||
}); | ||
|
||
return ( | ||
<div className="flex flex-col w-full max-w-xl py-24 mx-auto stretch"> | ||
<div className="space-y-8 mb-4"> | ||
{messages.map(message => ( | ||
<div key={message.id}> | ||
<div className="font-bold mb-2"> | ||
{message.role === 'user' ? 'You' : 'Assistant'} | ||
</div> | ||
<div className="prose space-y-2"> | ||
<MemoizedMarkdown id={message.id} content={message.content} /> | ||
</div> | ||
</div> | ||
))} | ||
</div> | ||
<MessageInput /> | ||
</div> | ||
); | ||
} | ||
|
||
const MessageInput = () => { | ||
const { input, handleSubmit, handleInputChange } = useChat({ id: 'chat' }); | ||
return ( | ||
<form onSubmit={handleSubmit}> | ||
<input | ||
className="fixed bottom-0 w-full max-w-xl p-2 mb-8 dark:bg-zinc-900 border border-zinc-300 dark:border-zinc-800 rounded shadow-xl" | ||
placeholder="Say something..." | ||
value={input} | ||
onChange={handleInputChange} | ||
/> | ||
</form> | ||
); | ||
}; | ||
``` | ||
|
||
<Note> | ||
The chat state is shared between both components by using the same `id` value. | ||
This allows you to split the form and chat messages into separate components | ||
while maintaining synchronized state. | ||
</Note> |