dmitry.galkin/chatbot-template
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: modularize block system (#718)

Browse source
This commit is contained in:
Jeremy 2025-01-27 14:19:47 +05:30 • committed by GitHub
parent 5e8cddc886
commit 38527ff92e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 684 additions and 604 deletions
Download patch file Download diff file Expand all files Collapse all files

8
blocks/actions.ts Normal file
View file

@ -0,0 +1,8 @@
'use server';
import { getSuggestionsByDocumentId } from '@/lib/db/queries';
export async function getSuggestions({ documentId }: { documentId: string }) {
const suggestions = await getSuggestionsByDocumentId({ documentId });
return suggestions;
}

154
blocks/code.tsx Normal file
View file

@ -0,0 +1,154 @@
import { Block } from '@/components/create-block';
import { CodeEditor } from '@/components/code-editor';
import {
CopyIcon,
LogsIcon,
MessageIcon,
PlayIcon,
RedoIcon,
UndoIcon,
} from '@/components/icons';
import { toast } from 'sonner';
import { generateUUID } from '@/lib/utils';
import { Console, ConsoleOutput } from '@/components/console';
interface Metadata {
outputs: Array<ConsoleOutput>;
}
export const codeBlock = new Block<'code', Metadata>({
kind: 'code',
description:
'Useful for code generation; Code execution is only available for python code.',
initialize: () => ({
outputs: [],
}),
onStreamPart: ({ streamPart, setBlock }) => {
if (streamPart.type === 'code-delta') {
setBlock((draftBlock) => ({
...draftBlock,
content: streamPart.content as string,
isVisible:
draftBlock.status === 'streaming' &&
draftBlock.content.length > 300 &&
draftBlock.content.length < 310
? true
: draftBlock.isVisible,
status: 'streaming',
}));
}
},
content: ({ metadata, setMetadata, ...props }) => {
return (
<>
<CodeEditor {...props} />
{metadata?.outputs && (
<Console
consoleOutputs={metadata.outputs}
setConsoleOutputs={() => {
setMetadata({
...metadata,
outputs: [],
});
}}
/>
)}
</>
);
},
actions: [
{
icon: <PlayIcon size={18} />,
label: 'Run',
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/',
});
currentPyodideInstance.setStdout({
batched: (output: string) => {
outputs.push({
id: runId,
contents: [
{
type: output.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: output,
},
],
status: 'completed',
});
},
});
await currentPyodideInstance.loadPackagesFromImports(content, {
messageCallback: (message: string) => {
outputs.push({
id: runId,
contents: [{ type: 'text', value: message }],
status: 'loading_packages',
});
},
});
await currentPyodideInstance.runPythonAsync(content);
setMetadata((metadata: any) => ({
...metadata,
outputs,
}));
},
},
{
icon: <UndoIcon size={18} />,
description: 'View Previous version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('prev');
},
},
{
icon: <RedoIcon size={18} />,
description: 'View Next version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('next');
},
},
{
icon: <CopyIcon size={18} />,
description: 'Copy code to clipboard',
onClick: ({ content }) => {
navigator.clipboard.writeText(content);
toast.success('Copied to clipboard!');
},
},
],
toolbar: [
{
icon: <MessageIcon />,
description: 'Add comments',
onClick: ({ appendMessage }) => {
appendMessage({
role: 'user',
content: 'Add comments to the code snippet for understanding',
});
},
},
{
icon: <LogsIcon />,
description: 'Add logs',
onClick: ({ appendMessage }) => {
appendMessage({
role: 'user',
content: 'Add logs to the code snippet for debugging',
});
},
},
],
});

59
blocks/image.tsx Normal file
View file

@ -0,0 +1,59 @@
import { Block } from '@/components/create-block';
import { CopyIcon, RedoIcon, UndoIcon } from '@/components/icons';
import { ImageEditor } from '@/components/image-editor';
export const imageBlock = new Block({
kind: 'image',
description: 'Useful for image generation',
onStreamPart: ({ streamPart, setBlock }) => {
if (streamPart.type === 'image-delta') {
setBlock((draftBlock) => ({
...draftBlock,
content: streamPart.content as string,
isVisible: true,
status: 'streaming',
}));
}
},
content: ImageEditor,
actions: [
{
icon: <UndoIcon size={18} />,
description: 'View Previous version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('prev');
},
},
{
icon: <RedoIcon size={18} />,
description: 'View Next version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('next');
},
},
{
icon: <CopyIcon size={18} />,
description: 'Copy image to clipboard',
onClick: ({ content }) => {
const img = new Image();
img.src = `data:image/png;base64,${content}`;
img.onload = () => {
const canvas = document.createElement('canvas');
canvas.width = img.width;
canvas.height = img.height;
const ctx = canvas.getContext('2d');
ctx?.drawImage(img, 0, 0);
canvas.toBlob((blob) => {
if (blob) {
navigator.clipboard.write([
new ClipboardItem({ 'image/png': blob }),
]);
}
}, 'image/png');
};
},
},
],
toolbar: [],
});

174
blocks/text.tsx Normal file
View file

@ -0,0 +1,174 @@
import { Block } from '@/components/create-block';
import { DiffView } from '@/components/diffview';
import { DocumentSkeleton } from '@/components/document-skeleton';
import { Editor } from '@/components/editor';
import {
ClockRewind,
CopyIcon,
MessageIcon,
PenIcon,
RedoIcon,
UndoIcon,
} from '@/components/icons';
import { Suggestion } from '@/lib/db/schema';
import { toast } from 'sonner';
import { getSuggestions } from './actions';
interface TextBlockMetadata {
suggestions: Array<Suggestion>;
}
export const textBlock = new Block<'text', TextBlockMetadata>({
kind: 'text',
description: 'Useful for text content, like drafting essays and emails.',
initialize: async ({ documentId, setMetadata }) => {
const suggestions = await getSuggestions({ documentId });
setMetadata({
suggestions,
});
},
onStreamPart: ({ streamPart, setMetadata, setBlock }) => {
if (streamPart.type === 'suggestion') {
setMetadata((metadata) => {
return {
suggestions: [
...metadata.suggestions,
streamPart.content as Suggestion,
],
};
});
}
if (streamPart.type === 'text-delta') {
setBlock((draftBlock) => {
return {
...draftBlock,
content: draftBlock.content + (streamPart.content as string),
isVisible:
draftBlock.status === 'streaming' &&
draftBlock.content.length > 400 &&
draftBlock.content.length < 450
? true
: draftBlock.isVisible,
status: 'streaming',
};
});
}
},
content: ({
mode,
status,
content,
isCurrentVersion,
currentVersionIndex,
onSaveContent,
getDocumentContentById,
isLoading,
metadata,
}) => {
if (isLoading) {
return <DocumentSkeleton blockKind="text" />;
}
if (mode === 'diff') {
const oldContent = getDocumentContentById(currentVersionIndex - 1);
const newContent = getDocumentContentById(currentVersionIndex);
return <DiffView oldContent={oldContent} newContent={newContent} />;
}
return (
<>
<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}
</>
);
},
actions: [
{
icon: <ClockRewind size={18} />,
description: 'View changes',
onClick: ({ handleVersionChange }) => {
handleVersionChange('toggle');
},
isDisabled: ({ currentVersionIndex, setMetadata }) => {
if (currentVersionIndex === 0) {
return true;
}
return false;
},
},
{
icon: <UndoIcon size={18} />,
description: 'View Previous version',
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} />,
description: 'Copy to clipboard',
onClick: ({ content }) => {
navigator.clipboard.writeText(content);
toast.success('Copied to clipboard!');
},
},
],
toolbar: [
{
icon: <PenIcon />,
description: 'Add final polish',
onClick: ({ appendMessage }) => {
appendMessage({
role: 'user',
content:
'Please add final polish and check for grammar, add section titles for better structure, and ensure everything reads smoothly.',
});
},
},
{
icon: <MessageIcon />,
description: 'Request suggestions',
onClick: ({ appendMessage }) => {
appendMessage({
role: 'user',
content:
'Please add suggestions you have that could improve the writing.',
});
},
},
],
});

121
components/block-actions.tsx
View file

@ -1,20 +1,18 @@
import { cn } from '@/lib/utils';
import { ClockRewind, CopyIcon, RedoIcon, UndoIcon } from './icons';
import { Button } from './ui/button';
import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
import { toast } from 'sonner';
import { ConsoleOutput, UIBlock } from './block';
import { blockDefinitions, UIBlock } from './block';
import { Dispatch, memo, SetStateAction } from 'react';
import { RunCodeButton } from './run-code-button';
import { useMultimodalCopyToClipboard } from '@/hooks/use-multimodal-copy-to-clipboard';
import { BlockActionContext } from './create-block';
import { cn } from '@/lib/utils';
interface BlockActionsProps {
block: UIBlock;
handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
currentVersionIndex: number;
isCurrentVersion: boolean;
mode: 'read-only' | 'edit' | 'diff';
setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
mode: 'edit' | 'diff';
metadata: any;
setMetadata: Dispatch<SetStateAction<any>>;
}
function PureBlockActions({
@ -23,95 +21,50 @@ function PureBlockActions({
currentVersionIndex,
isCurrentVersion,
mode,
setConsoleOutputs,
metadata,
setMetadata,
}: BlockActionsProps) {
const { copyTextToClipboard, copyImageToClipboard } =
useMultimodalCopyToClipboard();
const blockDefinition = blockDefinitions.find(
(definition) => definition.kind === block.kind,
);
if (!blockDefinition) {
throw new Error('Block definition not found!');
}
const actionContext: BlockActionContext = {
content: block.content,
handleVersionChange,
currentVersionIndex,
isCurrentVersion,
mode,
metadata,
setMetadata,
};
return (
<div className="flex flex-row gap-1">
{block.kind === 'code' && (
<RunCodeButton block={block} setConsoleOutputs={setConsoleOutputs} />
)}
{block.kind === 'text' && (
<Tooltip>
{blockDefinition.actions.map((action) => (
<Tooltip key={action.description}>
<TooltipTrigger asChild>
<Button
variant="outline"
className={cn(
'p-2 h-fit !pointer-events-auto dark:hover:bg-zinc-700',
{
'bg-muted': mode === 'diff',
},
)}
onClick={() => {
handleVersionChange('toggle');
}}
className={cn('h-fit dark:hover:bg-zinc-700', {
'p-2': !action.label,
'py-1.5 px-2': action.label,
})}
onClick={() => action.onClick(actionContext)}
disabled={
block.status === 'streaming' || currentVersionIndex === 0
action.isDisabled ? action.isDisabled(actionContext) : false
}
>
<ClockRewind size={18} />
{action.icon}
{action.label}
</Button>
</TooltipTrigger>
<TooltipContent>View changes</TooltipContent>
<TooltipContent>{action.description}</TooltipContent>
</Tooltip>
)}
<Tooltip>
<TooltipTrigger asChild>
<Button
variant="outline"
className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
onClick={() => {
handleVersionChange('prev');
}}
disabled={currentVersionIndex === 0 || block.status === 'streaming'}
>
<UndoIcon size={18} />
</Button>
</TooltipTrigger>
<TooltipContent>View Previous version</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
variant="outline"
className="p-2 h-fit dark:hover:bg-zinc-700 !pointer-events-auto"
onClick={() => {
handleVersionChange('next');
}}
disabled={isCurrentVersion || block.status === 'streaming'}
>
<RedoIcon size={18} />
</Button>
</TooltipTrigger>
<TooltipContent>View Next version</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
variant="outline"
className="p-2 h-fit dark:hover:bg-zinc-700"
onClick={() => {
if (block.kind === 'image') {
copyImageToClipboard(block.content);
} else {
copyTextToClipboard(block.content);
}
toast.success('Copied to clipboard!');
}}
disabled={block.status === 'streaming'}
>
<CopyIcon size={18} />
</Button>
</TooltipTrigger>
<TooltipContent>Copy to clipboard</TooltipContent>
</Tooltip>
))}
</div>
);
}

145
components/block.tsx
View file

@ -16,27 +16,23 @@ import {
} from 'react';
import useSWR, { useSWRConfig } from 'swr';
import { useDebounceCallback, useWindowSize } from 'usehooks-ts';
import type { Document, Suggestion, Vote } from '@/lib/db/schema';
import type { Document, Vote } from '@/lib/db/schema';
import { cn, fetcher } from '@/lib/utils';
import { DiffView } from './diffview';
import { DocumentSkeleton } from './document-skeleton';
import { Editor } from './editor';
import { MultimodalInput } from './multimodal-input';
import { Toolbar } from './toolbar';
import { VersionFooter } from './version-footer';
import { BlockActions } from './block-actions';
import { BlockCloseButton } from './block-close-button';
import { BlockMessages } from './block-messages';
import { CodeEditor } from './code-editor';
import { Console } from './console';
import { useSidebar } from './ui/sidebar';
import { useBlock } from '@/hooks/use-block';
import equal from 'fast-deep-equal';
import { ImageEditor } from './image-editor';
import { textBlock } from '@/blocks/text';
import { imageBlock } from '@/blocks/image';
import { codeBlock } from '@/blocks/code';
export type BlockKind = 'text' | 'code' | 'image';
export const blockDefinitions = [textBlock, codeBlock, imageBlock] as const;
export type BlockKind = (typeof blockDefinitions)[number]['kind'];
export interface UIBlock {
title: string;
@ -53,17 +49,6 @@ export interface UIBlock {
};
}
export interface ConsoleOutputContent {
type: 'text' | 'image';
value: string;
}
export interface ConsoleOutput {
id: string;
status: 'in_progress' | 'loading_packages' | 'completed' | 'failed';
contents: Array<ConsoleOutputContent>;
}
function PureBlock({
chatId,
input,
@ -105,7 +90,7 @@ function PureBlock({
) => Promise<string | null | undefined>;
isReadonly: boolean;
}) {
const { block, setBlock } = useBlock();
const { block, setBlock, metadata, setMetadata } = useBlock();
const {
data: documents,
@ -118,22 +103,9 @@ function PureBlock({
fetcher,
);
const { data: suggestions } = useSWR<Array<Suggestion>>(
documents && block && block.status !== 'streaming'
? `/api/suggestions?documentId=${block.documentId}`
: null,
fetcher,
{
dedupingInterval: 5000,
},
);
const [mode, setMode] = useState<'edit' | 'diff'>('edit');
const [document, setDocument] = useState<Document | null>(null);
const [currentVersionIndex, setCurrentVersionIndex] = useState(-1);
const [consoleOutputs, setConsoleOutputs] = useState<Array<ConsoleOutput>>(
[],
);
const { open: isSidebarOpen } = useSidebar();
@ -268,6 +240,23 @@ function PureBlock({
const { width: windowWidth, height: windowHeight } = useWindowSize();
const isMobile = windowWidth ? windowWidth < 768 : false;
const blockDefinition = blockDefinitions.find(
(definition) => definition.kind === block.kind,
);
if (!blockDefinition) {
throw new Error('Block definition not found!');
}
useEffect(() => {
if (block && block.documentId !== 'init') {
blockDefinition.initialize({
documentId: block.documentId,
setMetadata,
});
}
}, [block, blockDefinition, setMetadata]);
return (
<AnimatePresence>
{block.isVisible && (
@ -457,7 +446,8 @@ function PureBlock({
handleVersionChange={handleVersionChange}
isCurrentVersion={isCurrentVersion}
mode={mode}
setConsoleOutputs={setConsoleOutputs}
metadata={metadata}
setMetadata={setMetadata}
/>
</div>
@ -465,7 +455,7 @@ function PureBlock({
className={cn(
'dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full pb-40 items-center',
{
'py-2 px-2': block.kind === 'code',
'': block.kind === 'code',
'py-8 md:p-20 px-4': block.kind === 'text',
},
)}
@ -476,61 +466,25 @@ function PureBlock({
'mx-auto max-w-[600px]': block.kind === 'text',
})}
>
{isDocumentsFetching && !block.content ? (
<DocumentSkeleton blockKind={block.kind} />
) : block.kind === 'code' ? (
<CodeEditor
content={
isCurrentVersion
? block.content
: getDocumentContentById(currentVersionIndex)
}
isCurrentVersion={isCurrentVersion}
currentVersionIndex={currentVersionIndex}
suggestions={suggestions ?? []}
status={block.status}
saveContent={saveContent}
/>
) : block.kind === 'text' ? (
mode === 'edit' ? (
<Editor
content={
isCurrentVersion
? block.content
: getDocumentContentById(currentVersionIndex)
}
isCurrentVersion={isCurrentVersion}
currentVersionIndex={currentVersionIndex}
status={block.status}
saveContent={saveContent}
suggestions={isCurrentVersion ? (suggestions ?? []) : []}
/>
) : (
<DiffView
oldContent={getDocumentContentById(
currentVersionIndex - 1,
)}
newContent={getDocumentContentById(currentVersionIndex)}
/>
)
) : block.kind === 'image' ? (
<ImageEditor
title={block.title}
content={
isCurrentVersion
? block.content
: getDocumentContentById(currentVersionIndex)
}
isCurrentVersion={isCurrentVersion}
currentVersionIndex={currentVersionIndex}
status={block.status}
isInline={false}
/>
) : null}
{suggestions && suggestions.length > 0 ? (
<div className="md:hidden h-dvh w-12 shrink-0" />
) : null}
<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 && (
@ -557,13 +511,6 @@ function PureBlock({
/>
)}
</AnimatePresence>
<AnimatePresence>
<Console
consoleOutputs={consoleOutputs}
setConsoleOutputs={setConsoleOutputs}
/>
</AnimatePresence>
</motion.div>
</motion.div>
)}

8
components/code-editor.tsx
View file

@ -10,14 +10,14 @@ import { Suggestion } from '@/lib/db/schema';
type EditorProps = {
content: string;
saveContent: (updatedContent: string, debounce: boolean) => void;
onSaveContent: (updatedContent: string, debounce: boolean) => void;
status: 'streaming' | 'idle';
isCurrentVersion: boolean;
currentVersionIndex: number;
suggestions: Array<Suggestion>;
};
function PureCodeEditor({ content, saveContent, status }: EditorProps) {
function PureCodeEditor({ content, onSaveContent, status }: EditorProps) {
const containerRef = useRef<HTMLDivElement>(null);
const editorRef = useRef<EditorView | null>(null);
@ -54,7 +54,7 @@ function PureCodeEditor({ content, saveContent, status }: EditorProps) {
if (transaction) {
const newContent = update.state.doc.toString();
saveContent(newContent, true);
onSaveContent(newContent, true);
}
}
});
@ -69,7 +69,7 @@ function PureCodeEditor({ content, saveContent, status }: EditorProps) {
editorRef.current.setState(newState);
}
}, [saveContent]);
}, [onSaveContent]);
useEffect(() => {
if (editorRef.current && content) {

12
components/console.tsx
View file

@ -8,10 +8,20 @@ import {
useRef,
useState,
} from 'react';
import { ConsoleOutput } from './block';
import { cn } from '@/lib/utils';
import { useBlockSelector } from '@/hooks/use-block';
export interface ConsoleOutputContent {
type: 'text' | 'image';
value: string;
}
export interface ConsoleOutput {
id: string;
status: 'in_progress' | 'loading_packages' | 'completed' | 'failed';
contents: Array<ConsoleOutputContent>;
}
interface ConsoleProps {
consoleOutputs: Array<ConsoleOutput>;
setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;

100
components/create-block.tsx Normal file
View file

@ -0,0 +1,100 @@
import { Suggestion } from '@/lib/db/schema';
import { UseChatHelpers } from 'ai/react';
import { ComponentType, Dispatch, ReactNode, SetStateAction } from 'react';
import { DataStreamDelta } from './data-stream-handler';
import { UIBlock } from './block';
export type BlockActionContext = {
content: string;
handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
currentVersionIndex: number;
isCurrentVersion: boolean;
mode: 'edit' | 'diff';
metadata: any;
setMetadata: Dispatch<SetStateAction<any>>;
};
type BlockAction = {
icon: ReactNode;
label?: string;
description: string;
onClick: (context: BlockActionContext) => void;
isDisabled?: (context: BlockActionContext) => boolean;
};
export type BlockToolbarContext = {
appendMessage: UseChatHelpers['append'];
};
export type BlockToolbarItem = {
description: string;
icon: ReactNode;
onClick: (context: BlockToolbarContext) => void;
};
type BlockContent = {
title: string;
content: string;
mode: 'edit' | 'diff';
isCurrentVersion: boolean;
currentVersionIndex: number;
status: 'streaming' | 'idle';
suggestions: Array<Suggestion>;
onSaveContent: (updatedContent: string, debounce: boolean) => void;
isInline: boolean;
getDocumentContentById: (index: number) => string;
isLoading: boolean;
metadata: any;
setMetadata: Dispatch<SetStateAction<any>>;
};
interface InitializeParameters<M = any> {
documentId: string;
setMetadata: Dispatch<SetStateAction<M>>;
}
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;
initialize?: (parameters: InitializeParameters<M>) => void;
onStreamPart?: (args: {
setMetadata: Dispatch<SetStateAction<M>>;
setBlock: Dispatch<SetStateAction<UIBlock>>;
streamPart: DataStreamDelta;
}) => void;
};
export class Block<T extends string, M = any> {
readonly kind: T;
readonly description: string;
readonly content: ComponentType<BlockContent>;
readonly actions: BlockAction[];
readonly toolbar: BlockToolbarItem[];
readonly metadata: M;
readonly initialize: (parameters: InitializeParameters) => void;
readonly onStreamPart?: (args: {
setMetadata: Dispatch<SetStateAction<M>>;
setBlock: Dispatch<SetStateAction<UIBlock>>;
streamPart: DataStreamDelta;
}) => void;
constructor(config: BlockConfig<T, M>) {
this.kind = config.kind;
this.description = config.description;
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;
}
}

80
components/data-stream-handler.tsx
View file

@ -1,14 +1,13 @@
'use client';
import { useChat } from 'ai/react';
import { useEffect, useRef, useState } from 'react';
import { BlockKind } from './block';
import { useEffect, useRef } from 'react';
import { blockDefinitions, BlockKind } from './block';
import { Suggestion } from '@/lib/db/schema';
import { initialBlockData, useBlock } from '@/hooks/use-block';
import { useUserMessageId } from '@/hooks/use-user-message-id';
import { useSWRConfig } from 'swr';
type DataStreamDelta = {
export type DataStreamDelta = {
type:
| 'text-delta'
| 'code-delta'
@ -26,22 +25,9 @@ type DataStreamDelta = {
export function DataStreamHandler({ id }: { id: string }) {
const { data: dataStream } = useChat({ id });
const { setUserMessageIdFromServer } = useUserMessageId();
const { setBlock } = useBlock();
const { block, setBlock, setMetadata } = useBlock();
const lastProcessedIndex = useRef(-1);
const { mutate } = useSWRConfig();
const [optimisticSuggestions, setOptimisticSuggestions] = useState<
Array<Suggestion>
>([]);
useEffect(() => {
if (optimisticSuggestions && optimisticSuggestions.length > 0) {
const [optimisticSuggestion] = optimisticSuggestions;
const url = `/api/suggestions?documentId=${optimisticSuggestion.documentId}`;
mutate(url, optimisticSuggestions, false);
}
}, [optimisticSuggestions, mutate]);
useEffect(() => {
if (!dataStream?.length) return;
@ -54,6 +40,18 @@ export function DataStreamHandler({ id }: { id: string }) {
return;
}
const blockDefinition = blockDefinitions.find(
(blockDefinition) => blockDefinition.kind === block.kind,
);
if (blockDefinition?.onStreamPart) {
blockDefinition.onStreamPart({
streamPart: delta,
setBlock,
setMetadata,
});
}
setBlock((draftBlock) => {
if (!draftBlock) {
return { ...initialBlockData, status: 'streaming' };
@ -81,50 +79,6 @@ export function DataStreamHandler({ id }: { id: string }) {
status: 'streaming',
};
case 'text-delta':
return {
...draftBlock,
content: draftBlock.content + (delta.content as string),
isVisible:
draftBlock.status === 'streaming' &&
draftBlock.content.length > 400 &&
draftBlock.content.length < 450
? true
: draftBlock.isVisible,
status: 'streaming',
};
case 'code-delta':
return {
...draftBlock,
content: delta.content as string,
isVisible:
draftBlock.status === 'streaming' &&
draftBlock.content.length > 300 &&
draftBlock.content.length < 310
? true
: draftBlock.isVisible,
status: 'streaming',
};
case 'image-delta':
return {
...draftBlock,
content: delta.content as string,
isVisible: true,
status: 'streaming',
};
case 'suggestion':
setTimeout(() => {
setOptimisticSuggestions((currentSuggestions) => [
...currentSuggestions,
delta.content as Suggestion,
]);
}, 0);
return draftBlock;
case 'clear':
return {
...draftBlock,
@ -143,7 +97,7 @@ export function DataStreamHandler({ id }: { id: string }) {
}
});
});
}, [dataStream, setBlock, setUserMessageIdFromServer]);
}, [dataStream, setBlock, setUserMessageIdFromServer, setMetadata, block]);
return null;
}

4
components/document-preview.tsx
View file

@ -249,11 +249,11 @@ const DocumentContent = ({ document }: { document: Document }) => {
return (
<div className={containerClassName}>
{document.kind === 'text' ? (
<Editor {...commonProps} />
<Editor {...commonProps} onSaveContent={() => {}} />
) : document.kind === 'code' ? (
<div className="flex flex-1 relative w-full">
<div className="absolute inset-0">
<CodeEditor {...commonProps} />
<CodeEditor {...commonProps} onSaveContent={() => {}} />
</div>
</div>
) : document.kind === 'image' ? (

14
components/editor.tsx
View file

@ -25,7 +25,7 @@ import {
type EditorProps = {
content: string;
saveContent: (updatedContent: string, debounce: boolean) => void;
onSaveContent: (updatedContent: string, debounce: boolean) => void;
status: 'streaming' | 'idle';
isCurrentVersion: boolean;
currentVersionIndex: number;
@ -34,7 +34,7 @@ type EditorProps = {
function PureEditor({
content,
saveContent,
onSaveContent,
suggestions,
status,
}: EditorProps) {
@ -80,11 +80,15 @@ function PureEditor({
if (editorRef.current) {
editorRef.current.setProps({
dispatchTransaction: (transaction) => {
handleTransaction({ transaction, editorRef, saveContent });
handleTransaction({
transaction,
editorRef,
onSaveContent,
});
},
});
}
}, [saveContent]);
}, [onSaveContent]);
useEffect(() => {
if (editorRef.current && content) {
@ -153,7 +157,7 @@ function areEqual(prevProps: EditorProps, nextProps: EditorProps) {
prevProps.isCurrentVersion === nextProps.isCurrentVersion &&
!(prevProps.status === 'streaming' && nextProps.status === 'streaming') &&
prevProps.content === nextProps.content &&
prevProps.saveContent === nextProps.saveContent
prevProps.onSaveContent === nextProps.onSaveContent
);
}

2
components/image-editor.tsx
View file

@ -13,8 +13,6 @@ interface ImageEditorProps {
export function ImageEditor({
title,
content,
isCurrentVersion,
currentVersionIndex,
status,
isInline,
}: ImageEditorProps) {

210
components/run-code-button.tsx
View file

@ -1,210 +0,0 @@
import { generateUUID } from '@/lib/utils';
import {
type Dispatch,
type SetStateAction,
startTransition,
useCallback,
useState,
useEffect,
memo,
} from 'react';
import type { ConsoleOutput, ConsoleOutputContent, UIBlock } from './block';
import { Button } from './ui/button';
import { PlayIcon } from './icons';
import { useBlockSelector } from '@/hooks/use-block';
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;
}
export function PureRunCodeButton({
setConsoleOutputs,
}: {
block: UIBlock;
setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
}) {
const isPython = true;
const [pyodide, setPyodide] = useState<any>(null);
const codeContent = useBlockSelector((state) => state.content);
const isBlockStreaming = useBlockSelector(
(state) => state.status === 'streaming',
);
const loadAndRunPython = useCallback(async () => {
const runId = generateUUID();
const stdOutputs: Array<ConsoleOutputContent> = [];
setConsoleOutputs((outputs) => [
...outputs,
{
id: runId,
contents: [],
status: 'in_progress',
},
]);
let currentPyodideInstance = pyodide;
if (isPython) {
try {
if (!currentPyodideInstance) {
// @ts-expect-error - loadPyodide is not defined
const newPyodideInstance = await globalThis.loadPyodide({
indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
});
setPyodide(null);
setPyodide(newPyodideInstance);
currentPyodideInstance = newPyodideInstance;
}
currentPyodideInstance.setStdout({
batched: (output: string) => {
stdOutputs.push({
type: output.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: output,
});
},
});
await currentPyodideInstance.loadPackagesFromImports(codeContent, {
messageCallback: (message: string) => {
setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: message }],
status: 'loading_packages',
},
]);
},
});
const requiredHandlers = detectRequiredHandlers(codeContent);
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(codeContent);
setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: generateUUID(),
contents: stdOutputs.filter((output) => output.value.trim().length),
status: 'completed',
},
]);
} catch (error: any) {
setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: error.message }],
status: 'failed',
},
]);
}
}
}, [pyodide, codeContent, isPython, setConsoleOutputs]);
useEffect(() => {
return () => {
if (pyodide) {
try {
pyodide.runPythonAsync(`
import sys
import gc
has_plt = 'matplotlib.pyplot' in sys.modules
if has_plt:
import matplotlib.pyplot as plt
plt.clf()
plt.close('all')
gc.collect()
`);
} catch (error) {
console.warn('Cleanup failed:', error);
}
}
};
}, [pyodide]);
return (
<Button
variant="outline"
className="py-1.5 px-2 h-fit dark:hover:bg-zinc-700"
onClick={() => {
startTransition(() => {
loadAndRunPython();
});
}}
disabled={isBlockStreaming}
>
<PlayIcon size={18} /> Run
</Button>
);
}
export const RunCodeButton = memo(PureRunCodeButton, (prevProps, nextProps) => {
if (prevProps.block.status !== nextProps.block.status) return false;
return true;
});

145
components/toolbar.tsx
View file

@ -11,6 +11,7 @@ import {
import {
type Dispatch,
memo,
ReactNode,
type SetStateAction,
useEffect,
useRef,
@ -32,22 +33,16 @@ import {
LogsIcon,
MessageIcon,
PenIcon,
SparklesIcon,
StopIcon,
SummarizeIcon,
} from './icons';
import { BlockKind } from './block';
import { blockDefinitions, BlockKind } from './block';
import { BlockToolbarItem } from './create-block';
import { UseChatHelpers } from 'ai/react';
type ToolProps = {
type:
| 'final-polish'
| 'request-suggestions'
| 'adjust-reading-level'
| 'code-review'
| 'add-comments'
| 'add-logs';
description: string;
icon: JSX.Element;
icon: ReactNode;
selectedTool: string | null;
setSelectedTool: Dispatch<SetStateAction<string | null>>;
isToolbarVisible?: boolean;
@ -57,10 +52,14 @@ type ToolProps = {
message: Message | CreateMessage,
chatRequestOptions?: ChatRequestOptions,
) => Promise<string | null | undefined>;
onClick: ({
appendMessage,
}: {
appendMessage: UseChatHelpers['append'];
}) => void;
};
const Tool = ({
type,
description,
icon,
selectedTool,
@ -69,14 +68,15 @@ const Tool = ({
setIsToolbarVisible,
isAnimating,
append,
onClick,
}: ToolProps) => {
const [isHovered, setIsHovered] = useState(false);
useEffect(() => {
if (selectedTool !== type) {
if (selectedTool !== description) {
setIsHovered(false);
}
}, [selectedTool, type]);
}, [selectedTool, description]);
const handleSelect = () => {
if (!isToolbarVisible && setIsToolbarVisible) {
@ -86,44 +86,15 @@ const Tool = ({
if (!selectedTool) {
setIsHovered(true);
setSelectedTool(type);
setSelectedTool(description);
return;
}
if (selectedTool !== type) {
setSelectedTool(type);
if (selectedTool !== description) {
setSelectedTool(description);
} else {
if (type === 'final-polish') {
append({
role: 'user',
content:
'Please add final polish and check for grammar, add section titles for better structure, and ensure everything reads smoothly.',
});
setSelectedTool(null);
} else if (type === 'request-suggestions') {
append({
role: 'user',
content:
'Please add suggestions you have that could improve the writing.',
});
setSelectedTool(null);
} else if (type === 'add-comments') {
append({
role: 'user',
content: 'Please add comments to explain the code.',
});
setSelectedTool(null);
} else if (type === 'add-logs') {
append({
role: 'user',
content: 'Please add logs to help debug the code.',
});
setSelectedTool(null);
}
setSelectedTool(null);
onClick({ appendMessage: append });
}
};
@ -132,13 +103,13 @@ const Tool = ({
<TooltipTrigger asChild>
<motion.div
className={cx('p-3 rounded-full', {
'bg-primary !text-primary-foreground': selectedTool === type,
'bg-primary !text-primary-foreground': selectedTool === description,
})}
onHoverStart={() => {
setIsHovered(true);
}}
onHoverEnd={() => {
if (selectedTool !== type) setIsHovered(false);
if (selectedTool !== description) setIsHovered(false);
}}
onKeyDown={(event) => {
if (event.key === 'Enter') {
@ -158,7 +129,7 @@ const Tool = ({
handleSelect();
}}
>
{selectedTool === type ? <ArrowUpIcon /> : icon}
{selectedTool === description ? <ArrowUpIcon /> : icon}
</motion.div>
</TooltipTrigger>
<TooltipContent
@ -283,52 +254,6 @@ const ReadingLevelSelector = ({
);
};
const toolsByBlockKind: Record<
BlockKind,
Array<{
type:
| 'final-polish'
| 'request-suggestions'
| 'adjust-reading-level'
| 'code-review'
| 'add-comments'
| 'add-logs';
description: string;
icon: JSX.Element;
}>
> = {
text: [
{
type: 'final-polish',
description: 'Add final polish',
icon: <PenIcon />,
},
{
type: 'adjust-reading-level',
description: 'Adjust reading level',
icon: <SummarizeIcon />,
},
{
type: 'request-suggestions',
description: 'Request suggestions',
icon: <MessageIcon />,
},
],
code: [
{
type: 'add-comments',
description: 'Add comments',
icon: <CodeIcon />,
},
{
type: 'add-logs',
description: 'Add logs',
icon: <LogsIcon />,
},
],
image: [],
};
export const Tools = ({
isToolbarVisible,
selectedTool,
@ -336,7 +261,7 @@ export const Tools = ({
append,
isAnimating,
setIsToolbarVisible,
blockKind,
tools,
}: {
isToolbarVisible: boolean;
selectedTool: string | null;
@ -347,9 +272,9 @@ export const Tools = ({
) => Promise<string | null | undefined>;
isAnimating: boolean;
setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
blockKind: BlockKind;
tools: Array<BlockToolbarItem>;
}) => {
const [primaryTool, ...secondaryTools] = toolsByBlockKind[blockKind];
const [primaryTool, ...secondaryTools] = tools;
return (
<motion.div
@ -362,20 +287,19 @@ export const Tools = ({
{isToolbarVisible &&
secondaryTools.map((secondaryTool) => (
<Tool
key={secondaryTool.type}
type={secondaryTool.type}
key={secondaryTool.description}
description={secondaryTool.description}
icon={secondaryTool.icon}
selectedTool={selectedTool}
setSelectedTool={setSelectedTool}
append={append}
isAnimating={isAnimating}
onClick={secondaryTool.onClick}
/>
))}
</AnimatePresence>
<Tool
type={primaryTool.type}
description={primaryTool.description}
icon={primaryTool.icon}
selectedTool={selectedTool}
@ -384,6 +308,7 @@ export const Tools = ({
setIsToolbarVisible={setIsToolbarVisible}
append={append}
isAnimating={isAnimating}
onClick={primaryTool.onClick}
/>
</motion.div>
);
@ -451,7 +376,17 @@ const PureToolbar = ({
}
}, [isLoading, setIsToolbarVisible]);
if (toolsByBlockKind[blockKind].length === 0) {
const blockDefinition = blockDefinitions.find(
(definition) => definition.kind === blockKind,
);
if (!blockDefinition) {
throw new Error('Block definition not found!');
}
const toolsByBlockKind = blockDefinition.toolbar;
if (toolsByBlockKind.length === 0) {
return null;
}
@ -473,7 +408,7 @@ const PureToolbar = ({
: {
opacity: 1,
y: 0,
height: toolsByBlockKind[blockKind].length * 50,
height: toolsByBlockKind.length * 50,
transition: { delay: 0 },
scale: 1,
}
@ -530,7 +465,7 @@ const PureToolbar = ({
selectedTool={selectedTool}
setIsToolbarVisible={setIsToolbarVisible}
setSelectedTool={setSelectedTool}
blockKind={blockKind}
tools={toolsByBlockKind}
/>
)}
</motion.div>

22
hooks/use-block.ts
View file

@ -1,8 +1,8 @@
'use client';
import useSWR from 'swr';
import { UIBlock } from '@/components/block';
import { useCallback, useMemo } from 'react';
import useSWR from 'swr';
export const initialBlockData: UIBlock = {
documentId: 'init',
@ -19,7 +19,6 @@ export const initialBlockData: UIBlock = {
},
};
// Add type for selector function
type Selector<T> = (state: UIBlock) => T;
export function useBlockSelector<Selected>(selector: Selector<Selected>) {
@ -64,5 +63,22 @@ export function useBlock() {
[setLocalBlock],
);
return useMemo(() => ({ block, setBlock }), [block, setBlock]);
const { data: localBlockMetadata, mutate: setLocalBlockMetadata } =
useSWR<any>(
() => (block.documentId ? `block-metadata-${block.documentId}` : null),
null,
{
fallbackData: null,
},
);
return useMemo(
() => ({
block,
setBlock,
metadata: localBlockMetadata,
setMetadata: setLocalBlockMetadata,
}),
[block, setBlock, localBlockMetadata, setLocalBlockMetadata],
);
}

22
hooks/use-multimodal-copy-to-clipboard.ts
View file

@ -1,22 +0,0 @@
import { useCopyToClipboard } from 'usehooks-ts';
async function copyImageToClipboard(base64String: string) {
try {
const blob = await fetch(`data:image/png;base64,${base64String}`).then(
(res) => res.blob(),
);
const item = new ClipboardItem({
'image/png': blob,
});
await navigator.clipboard.write([item]);
} catch (error) {
console.error('Failed to copy image to clipboard:', error);
}
}
export function useMultimodalCopyToClipboard() {
const [_, copyTextToClipboard] = useCopyToClipboard();
return { copyTextToClipboard, copyImageToClipboard };
}

8
lib/editor/config.ts
View file

@ -24,11 +24,11 @@ export function headingRule(level: number) {
export const handleTransaction = ({
transaction,
editorRef,
saveContent,
onSaveContent,
}: {
transaction: Transaction;
editorRef: MutableRefObject<EditorView | null>;
saveContent: (updatedContent: string, debounce: boolean) => void;
onSaveContent: (updatedContent: string, debounce: boolean) => void;
}) => {
if (!editorRef || !editorRef.current) return;
@ -39,9 +39,9 @@ export const handleTransaction = ({
const updatedContent = buildContentFromDocument(newState.doc);
if (transaction.getMeta('no-debounce')) {
saveContent(updatedContent, false);
onSaveContent(updatedContent, false);
} else {
saveContent(updatedContent, true);
onSaveContent(updatedContent, true);
}
}
};