Skip to content

Commit

Permalink
fix: restore code block loading states (vercel#728)
Browse files Browse the repository at this point in the history
jeremyphilemon authored Jan 27, 2025
1 parent 38527ff commit 085f4a8
Showing 6 changed files with 267 additions and 134 deletions.
196 changes: 156 additions & 40 deletions blocks/code.tsx
Original file line number Diff line number Diff line change
@@ -10,7 +10,57 @@ import {
} from '@/components/icons';
import { toast } from 'sonner';
import { generateUUID } from '@/lib/utils';
import { Console, ConsoleOutput } from '@/components/console';
import {
Console,
ConsoleOutput,
ConsoleOutputContent,
} from '@/components/console';

const OUTPUT_HANDLERS = {
matplotlib: `
import io
import base64
from matplotlib import pyplot as plt
# Clear any existing plots
plt.clf()
plt.close('all')
# Switch to agg backend
plt.switch_backend('agg')
def setup_matplotlib_output():
def custom_show():
if plt.gcf().get_size_inches().prod() * plt.gcf().dpi ** 2 > 25_000_000:
print("Warning: Plot size too large, reducing quality")
plt.gcf().set_dpi(100)
png_buf = io.BytesIO()
plt.savefig(png_buf, format='png')
png_buf.seek(0)
png_base64 = base64.b64encode(png_buf.read()).decode('utf-8')
print(f'data:image/png;base64,{png_base64}')
png_buf.close()
plt.clf()
plt.close('all')
plt.show = custom_show
`,
basic: `
# Basic output capture setup
`,
};

function detectRequiredHandlers(code: string): string[] {
const handlers: string[] = ['basic'];

if (code.includes('matplotlib') || code.includes('plt.')) {
handlers.push('matplotlib');
}

return handlers;
}

interface Metadata {
outputs: Array<ConsoleOutput>;
@@ -20,9 +70,11 @@ export const codeBlock = new Block<'code', Metadata>({
kind: 'code',
description:
'Useful for code generation; Code execution is only available for python code.',
initialize: () => ({
outputs: [],
}),
initialize: async ({ setMetadata }) => {
setMetadata({
outputs: [],
});
},
onStreamPart: ({ streamPart, setBlock }) => {
if (streamPart.type === 'code-delta') {
setBlock((draftBlock) => ({
@@ -41,7 +93,9 @@ export const codeBlock = new Block<'code', Metadata>({
content: ({ metadata, setMetadata, ...props }) => {
return (
<>
<CodeEditor {...props} />
<div className="px-1">
<CodeEditor {...props} />
</div>

{metadata?.outputs && (
<Console
@@ -64,46 +118,94 @@ export const codeBlock = new Block<'code', Metadata>({
description: 'Execute code',
onClick: async ({ content, setMetadata }) => {
const runId = generateUUID();
const outputs: any[] = [];

// @ts-expect-error - loadPyodide is not defined
const currentPyodideInstance = await globalThis.loadPyodide({
indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
});
const outputContent: Array<ConsoleOutputContent> = [];

currentPyodideInstance.setStdout({
batched: (output: string) => {
outputs.push({
setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs,
{
id: runId,
contents: [
{
type: output.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: output,
},
],
status: 'completed',
});
},
});
contents: [],
status: 'in_progress',
},
],
}));

await currentPyodideInstance.loadPackagesFromImports(content, {
messageCallback: (message: string) => {
outputs.push({
id: runId,
contents: [{ type: 'text', value: message }],
status: 'loading_packages',
});
},
});
try {
// @ts-expect-error - loadPyodide is not defined
const currentPyodideInstance = await globalThis.loadPyodide({
indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
});

await currentPyodideInstance.runPythonAsync(content);
currentPyodideInstance.setStdout({
batched: (output: string) => {
outputContent.push({
type: output.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: output,
});
},
});

setMetadata((metadata: any) => ({
...metadata,
outputs,
}));
await currentPyodideInstance.loadPackagesFromImports(content, {
messageCallback: (message: string) => {
setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: message }],
status: 'loading_packages',
},
],
}));
},
});

const requiredHandlers = detectRequiredHandlers(content);
for (const handler of requiredHandlers) {
if (OUTPUT_HANDLERS[handler as keyof typeof OUTPUT_HANDLERS]) {
await currentPyodideInstance.runPythonAsync(
OUTPUT_HANDLERS[handler as keyof typeof OUTPUT_HANDLERS],
);

if (handler === 'matplotlib') {
await currentPyodideInstance.runPythonAsync(
'setup_matplotlib_output()',
);
}
}
}

await currentPyodideInstance.runPythonAsync(content);

setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: outputContent,
status: 'completed',
},
],
}));
} catch (error: any) {
setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: error.message }],
status: 'failed',
},
],
}));
}
},
},
{
@@ -112,13 +214,27 @@ export const codeBlock = new Block<'code', Metadata>({
onClick: ({ handleVersionChange }) => {
handleVersionChange('prev');
},
isDisabled: ({ currentVersionIndex }) => {
if (currentVersionIndex === 0) {
return true;
}

return false;
},
},
{
icon: <RedoIcon size={18} />,
description: 'View Next version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('next');
},
isDisabled: ({ isCurrentVersion }) => {
if (isCurrentVersion) {
return true;
}

return false;
},
},
{
icon: <CopyIcon size={18} />,
17 changes: 17 additions & 0 deletions blocks/image.tsx
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { Block } from '@/components/create-block';
import { CopyIcon, RedoIcon, UndoIcon } from '@/components/icons';
import { ImageEditor } from '@/components/image-editor';
import { toast } from 'sonner';

export const imageBlock = new Block({
kind: 'image',
@@ -23,13 +24,27 @@ export const imageBlock = new Block({
onClick: ({ handleVersionChange }) => {
handleVersionChange('prev');
},
isDisabled: ({ currentVersionIndex }) => {
if (currentVersionIndex === 0) {
return true;
}

return false;
},
},
{
icon: <RedoIcon size={18} />,
description: 'View Next version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('next');
},
isDisabled: ({ isCurrentVersion }) => {
if (isCurrentVersion) {
return true;
}

return false;
},
},
{
icon: <CopyIcon size={18} />,
@@ -52,6 +67,8 @@ export const imageBlock = new Block({
}
}, 'image/png');
};

toast.success('Copied image to clipboard!');
},
},
],
26 changes: 15 additions & 11 deletions blocks/text.tsx
Original file line number Diff line number Diff line change
@@ -80,18 +80,22 @@ export const textBlock = new Block<'text', TextBlockMetadata>({

return (
<>
<Editor
content={content}
suggestions={metadata ? metadata.suggestions : []}
isCurrentVersion={isCurrentVersion}
currentVersionIndex={currentVersionIndex}
status={status}
onSaveContent={onSaveContent}
/>
<div className="flex flex-row py-8 md:p-20 px-4">
<Editor
content={content}
suggestions={metadata ? metadata.suggestions : []}
isCurrentVersion={isCurrentVersion}
currentVersionIndex={currentVersionIndex}
status={status}
onSaveContent={onSaveContent}
/>

{metadata && metadata.suggestions && metadata.suggestions.length > 0 ? (
<div className="md:hidden h-dvh w-12 shrink-0" />
) : null}
{metadata &&
metadata.suggestions &&
metadata.suggestions.length > 0 ? (
<div className="md:hidden h-dvh w-12 shrink-0" />
) : null}
</div>
</>
);
},
23 changes: 20 additions & 3 deletions components/block-actions.tsx
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import { Button } from './ui/button';
import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
import { blockDefinitions, UIBlock } from './block';
import { Dispatch, memo, SetStateAction } from 'react';
import { Dispatch, memo, SetStateAction, useState } from 'react';
import { BlockActionContext } from './create-block';
import { cn } from '@/lib/utils';
import { toast } from 'sonner';

interface BlockActionsProps {
block: UIBlock;
@@ -24,6 +25,8 @@ function PureBlockActions({
metadata,
setMetadata,
}: BlockActionsProps) {
const [isLoading, setIsLoading] = useState(false);

const blockDefinition = blockDefinitions.find(
(definition) => definition.kind === block.kind,
);
@@ -53,9 +56,23 @@ function PureBlockActions({
'p-2': !action.label,
'py-1.5 px-2': action.label,
})}
onClick={() => action.onClick(actionContext)}
onClick={async () => {
setIsLoading(true);

try {
await Promise.resolve(action.onClick(actionContext));
} catch (error) {
toast.error('Failed to execute action');
} finally {
setIsLoading(false);
}
}}
disabled={
action.isDisabled ? action.isDisabled(actionContext) : false
isLoading || block.status === 'streaming'
? true
: action.isDisabled
? action.isDisabled(actionContext)
: false
}
>
{action.icon}
95 changes: 41 additions & 54 deletions components/block.tsx
Original file line number Diff line number Diff line change
@@ -17,7 +17,7 @@ import {
import useSWR, { useSWRConfig } from 'swr';
import { useDebounceCallback, useWindowSize } from 'usehooks-ts';
import type { Document, Vote } from '@/lib/db/schema';
import { cn, fetcher } from '@/lib/utils';
import { fetcher } from '@/lib/utils';
import { MultimodalInput } from './multimodal-input';
import { Toolbar } from './toolbar';
import { VersionFooter } from './version-footer';
@@ -26,10 +26,10 @@ import { BlockCloseButton } from './block-close-button';
import { BlockMessages } from './block-messages';
import { useSidebar } from './ui/sidebar';
import { useBlock } from '@/hooks/use-block';
import equal from 'fast-deep-equal';
import { textBlock } from '@/blocks/text';
import { imageBlock } from '@/blocks/image';
import { codeBlock } from '@/blocks/code';
import equal from 'fast-deep-equal';

export const blockDefinitions = [textBlock, codeBlock, imageBlock] as const;
export type BlockKind = (typeof blockDefinitions)[number]['kind'];
@@ -250,10 +250,12 @@ function PureBlock({

useEffect(() => {
if (block && block.documentId !== 'init') {
blockDefinition.initialize({
documentId: block.documentId,
setMetadata,
});
if (blockDefinition.initialize) {
blockDefinition.initialize({
documentId: block.documentId,
setMetadata,
});
}
}
}, [block, blockDefinition, setMetadata]);

@@ -451,55 +453,40 @@ function PureBlock({
/>
</div>

<div
className={cn(
'dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full pb-40 items-center',
{
'': block.kind === 'code',
'py-8 md:p-20 px-4': block.kind === 'text',
},
)}
>
<div
className={cn('flex flex-row', {
'': block.kind === 'code',
'mx-auto max-w-[600px]': block.kind === 'text',
})}
>
<blockDefinition.content
title={block.title}
content={
isCurrentVersion
? block.content
: getDocumentContentById(currentVersionIndex)
}
mode={mode}
status={block.status}
currentVersionIndex={currentVersionIndex}
suggestions={[]}
onSaveContent={saveContent}
isInline={false}
isCurrentVersion={isCurrentVersion}
getDocumentContentById={getDocumentContentById}
isLoading={isDocumentsFetching && !block.content}
metadata={metadata}
setMetadata={setMetadata}
/>
<div className="dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full items-center">
<blockDefinition.content
title={block.title}
content={
isCurrentVersion
? block.content
: getDocumentContentById(currentVersionIndex)
}
mode={mode}
status={block.status}
currentVersionIndex={currentVersionIndex}
suggestions={[]}
onSaveContent={saveContent}
isInline={false}
isCurrentVersion={isCurrentVersion}
getDocumentContentById={getDocumentContentById}
isLoading={isDocumentsFetching && !block.content}
metadata={metadata}
setMetadata={setMetadata}
/>

<AnimatePresence>
{isCurrentVersion && (
<Toolbar
isToolbarVisible={isToolbarVisible}
setIsToolbarVisible={setIsToolbarVisible}
append={append}
isLoading={isLoading}
stop={stop}
setMessages={setMessages}
blockKind={block.kind}
/>
)}
</AnimatePresence>
</div>
<AnimatePresence>
{isCurrentVersion && (
<Toolbar
isToolbarVisible={isToolbarVisible}
setIsToolbarVisible={setIsToolbarVisible}
append={append}
isLoading={isLoading}
stop={stop}
setMessages={setMessages}
blockKind={block.kind}
/>
)}
</AnimatePresence>
</div>

<AnimatePresence>
44 changes: 18 additions & 26 deletions components/create-block.tsx
Original file line number Diff line number Diff line change
@@ -4,22 +4,22 @@ import { ComponentType, Dispatch, ReactNode, SetStateAction } from 'react';
import { DataStreamDelta } from './data-stream-handler';
import { UIBlock } from './block';

export type BlockActionContext = {
export type BlockActionContext<M = any> = {
content: string;
handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
currentVersionIndex: number;
isCurrentVersion: boolean;
mode: 'edit' | 'diff';
metadata: any;
setMetadata: Dispatch<SetStateAction<any>>;
metadata: M;
setMetadata: Dispatch<SetStateAction<M>>;
};

type BlockAction = {
type BlockAction<M = any> = {
icon: ReactNode;
label?: string;
description: string;
onClick: (context: BlockActionContext) => void;
isDisabled?: (context: BlockActionContext) => boolean;
onClick: (context: BlockActionContext<M>) => Promise<void> | void;
isDisabled?: (context: BlockActionContext<M>) => boolean;
};

export type BlockToolbarContext = {
@@ -32,7 +32,7 @@ export type BlockToolbarItem = {
onClick: (context: BlockToolbarContext) => void;
};

type BlockContent = {
interface BlockContent<M = any> {
title: string;
content: string;
mode: 'edit' | 'diff';
@@ -44,9 +44,9 @@ type BlockContent = {
isInline: boolean;
getDocumentContentById: (index: number) => string;
isLoading: boolean;
metadata: any;
setMetadata: Dispatch<SetStateAction<any>>;
};
metadata: M;
setMetadata: Dispatch<SetStateAction<M>>;
}

interface InitializeParameters<M = any> {
documentId: string;
@@ -56,17 +56,11 @@ interface InitializeParameters<M = any> {
type BlockConfig<T extends string, M = any> = {
kind: T;
description: string;
content: ComponentType<
Omit<BlockContent, 'metadata' | 'setMetadata'> & {
metadata: M;
setMetadata: Dispatch<SetStateAction<M>>;
}
>;
actions?: BlockAction[];
toolbar?: BlockToolbarItem[];
metadata?: M;
content: ComponentType<BlockContent<M>>;
actions: Array<BlockAction<M>>;
toolbar: BlockToolbarItem[];
initialize?: (parameters: InitializeParameters<M>) => void;
onStreamPart?: (args: {
onStreamPart: (args: {
setMetadata: Dispatch<SetStateAction<M>>;
setBlock: Dispatch<SetStateAction<UIBlock>>;
streamPart: DataStreamDelta;
@@ -76,12 +70,11 @@ type BlockConfig<T extends string, M = any> = {
export class Block<T extends string, M = any> {
readonly kind: T;
readonly description: string;
readonly content: ComponentType<BlockContent>;
readonly actions: BlockAction[];
readonly content: ComponentType<BlockContent<M>>;
readonly actions: Array<BlockAction<M>>;
readonly toolbar: BlockToolbarItem[];
readonly metadata: M;
readonly initialize: (parameters: InitializeParameters) => void;
readonly onStreamPart?: (args: {
readonly initialize?: (parameters: InitializeParameters) => void;
readonly onStreamPart: (args: {
setMetadata: Dispatch<SetStateAction<M>>;
setBlock: Dispatch<SetStateAction<UIBlock>>;
streamPart: DataStreamDelta;
@@ -93,7 +86,6 @@ export class Block<T extends string, M = any> {
this.content = config.content;
this.actions = config.actions || [];
this.toolbar = config.toolbar || [];
this.metadata = config.metadata as M;
this.initialize = config.initialize || (async () => ({}));
this.onStreamPart = config.onStreamPart;
}

0 comments on commit 085f4a8

Please sign in to comment.