Skip to content

Commit

Permalink
docs: add memoization recipe (#4116)
Browse files Browse the repository at this point in the history
  • Loading branch information
nicoalbanese authored Dec 18, 2024
1 parent 5d0e8aa commit eb0632d
Showing 1 changed file with 133 additions and 0 deletions.
133 changes: 133 additions & 0 deletions content/cookbook/01-next/25-markdown-chatbot-with-memoization.mdx
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>

0 comments on commit eb0632d

Please sign in to comment.