feat: add code and text block types (#609)

2024-12-16 18:14:40 +05:30 · 2024-12-16 18:14:40 +05:30 · 9778631d6f
commit 9778631d6f
parent 3df0fd4c0f
27 changed files with 1754 additions and 290 deletions
--- a/components/block-actions.tsx
+++ b/components/block-actions.tsx
@ -1,11 +1,18 @@
-import { cn } from '@/lib/utils';
-import { CopyIcon, DeltaIcon, RedoIcon, UndoIcon } from './icons';
+import { cn, generateUUID } from '@/lib/utils';
+import { ClockRewind, CopyIcon, PlayIcon, RedoIcon, UndoIcon } from './icons';
 import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { useCopyToClipboard } from 'usehooks-ts';
 import { toast } from 'sonner';
-import { UIBlock } from './block';
-import { memo } from 'react';
+import { ConsoleOutput, UIBlock } from './block';
+import {
+  Dispatch,
+  memo,
+  SetStateAction,
+  startTransition,
+  useCallback,
+  useState,
+} from 'react';

 interface BlockActionsProps {
  block: UIBlock;
@ -13,6 +20,99 @@ interface BlockActionsProps {
  currentVersionIndex: number;
  isCurrentVersion: boolean;
  mode: 'read-only' | 'edit' | 'diff';
+  setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
+}
+
+export function RunCodeButton({
+  block,
+  setConsoleOutputs,
+}: {
+  block: UIBlock;
+  setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
+}) {
+  const [pyodide, setPyodide] = useState<any>(null);
+  const isPython = true;
+  const codeContent = block.content;
+
+  const updateConsoleOutput = useCallback(
+    (runId: string, content: string | null, status: 'completed' | 'failed') => {
+      setConsoleOutputs((consoleOutputs) => {
+        const index = consoleOutputs.findIndex((output) => output.id === runId);
+
+        if (index === -1) return consoleOutputs;
+
+        const updatedOutputs = [...consoleOutputs];
+        updatedOutputs[index] = {
+          id: runId,
+          content,
+          status,
+        };
+
+        return updatedOutputs;
+      });
+    },
+    [setConsoleOutputs],
+  );
+
+  const loadAndRunPython = useCallback(async () => {
+    const runId = generateUUID();
+
+    setConsoleOutputs((consoleOutputs) => [
+      ...consoleOutputs,
+      {
+        id: runId,
+        content: null,
+        status: 'in_progress',
+      },
+    ]);
+
+    let currentPyodideInstance = pyodide;
+
+    if (isPython) {
+      if (!currentPyodideInstance) {
+        // @ts-expect-error - pyodide is not defined
+        const newPyodideInstance = await loadPyodide({
+          indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
+        });
+
+        setPyodide(newPyodideInstance);
+        currentPyodideInstance = newPyodideInstance;
+      }
+
+      try {
+        await currentPyodideInstance.runPythonAsync(`
+            import sys
+            import io
+            sys.stdout = io.StringIO()
+          `);
+
+        await currentPyodideInstance.runPythonAsync(codeContent);
+
+        const output: string = await currentPyodideInstance.runPythonAsync(
+          `sys.stdout.getvalue()`,
+        );
+
+        updateConsoleOutput(runId, output, 'completed');
+      } catch (error: any) {
+        updateConsoleOutput(runId, error.message, 'failed');
+      }
+    }
+  }, [pyodide, codeContent, isPython, setConsoleOutputs, updateConsoleOutput]);
+
+  return (
+    <Button
+      variant="outline"
+      className="py-1.5 px-2 h-fit dark:hover:bg-zinc-700"
+      onClick={() => {
+        startTransition(() => {
+          loadAndRunPython();
+        });
+      }}
+      disabled={block.status === 'streaming'}
+    >
+      <PlayIcon size={18} /> Run
+    </Button>
+  );
 }

 function PureBlockActions({
@ -21,11 +121,73 @@ function PureBlockActions({
  currentVersionIndex,
  isCurrentVersion,
  mode,
+  setConsoleOutputs,
 }: BlockActionsProps) {
  const [_, copyToClipboard] = useCopyToClipboard();

  return (
    <div className="flex flex-row gap-1">
+      {block.kind === 'code' && (
+        <RunCodeButton block={block} setConsoleOutputs={setConsoleOutputs} />
+      )}
+
+      {block.kind === 'text' && (
+        <Tooltip>
+          <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');
+              }}
+              disabled={
+                block.status === 'streaming' || currentVersionIndex === 0
+              }
+            >
+              <ClockRewind size={18} />
+            </Button>
+          </TooltipTrigger>
+          <TooltipContent>View changes</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
@ -42,56 +204,6 @@ function PureBlockActions({
        </TooltipTrigger>
        <TooltipContent>Copy to clipboard</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={cn(
-              'p-2 h-fit !pointer-events-auto dark:hover:bg-zinc-700',
-              {
-                'bg-muted': mode === 'diff',
-              },
-            )}
-            onClick={() => {
-              handleVersionChange('toggle');
-            }}
-            disabled={block.status === 'streaming' || currentVersionIndex === 0}
-          >
-            <DeltaIcon size={18} />
-          </Button>
-        </TooltipTrigger>
-        <TooltipContent>View changes</TooltipContent>
-      </Tooltip>
    </div>
  );
 }
--- a/components/block.tsx
+++ b/components/block.tsx
@ -18,7 +18,7 @@ import useSWR, { useSWRConfig } from 'swr';
 import { useDebounceCallback, useWindowSize } from 'usehooks-ts';

 import type { Document, Suggestion, Vote } from '@/lib/db/schema';
-import { fetcher } from '@/lib/utils';
+import { cn, fetcher } from '@/lib/utils';

 import { DiffView } from './diffview';
 import { DocumentSkeleton } from './document-skeleton';
@ -29,10 +29,15 @@ 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';
+
+export type BlockKind = 'text' | 'code';

 export interface UIBlock {
  title: string;
  documentId: string;
+  kind: BlockKind;
  content: string;
  isVisible: boolean;
  status: 'streaming' | 'idle';
@ -44,6 +49,12 @@ export interface UIBlock {
  };
 }

+export interface ConsoleOutput {
+  id: string;
+  status: 'in_progress' | 'completed' | 'failed';
+  content: string | null;
+}
+
 function PureBlock({
  chatId,
  input,
@ -113,6 +124,9 @@ function PureBlock({
  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>>(
+    [],
+  );

  useEffect(() => {
    if (documents && documents.length > 0) {
@ -158,6 +172,7 @@ function PureBlock({
              body: JSON.stringify({
                title: block.title,
                content: updatedContent,
+                kind: block.kind,
              }),
            });

@ -318,7 +333,7 @@ function PureBlock({
      )}

      <motion.div
-        className="fixed dark:bg-muted bg-background h-dvh flex flex-col shadow-xl overflow-y-scroll"
+        className="fixed dark:bg-muted bg-background h-dvh flex flex-col overflow-y-scroll"
        initial={
          isMobile
            ? {
@ -415,15 +430,29 @@ function PureBlock({
            handleVersionChange={handleVersionChange}
            isCurrentVersion={isCurrentVersion}
            mode={mode}
+            setConsoleOutputs={setConsoleOutputs}
          />
        </div>

-        <div className="prose dark:prose-invert dark:bg-muted bg-background h-full overflow-y-scroll px-4 py-8 md:p-20 !max-w-full pb-40 items-center">
-          <div className="flex flex-row max-w-[600px] mx-auto">
+        <div
+          className={cn(
+            'prose dark:prose-invert dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full pb-40 items-center',
+            {
+              'py-2 px-2': 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',
+            })}
+          >
            {isDocumentsFetching && !block.content ? (
              <DocumentSkeleton />
-            ) : mode === 'edit' ? (
-              <Editor
+            ) : block.kind === 'code' ? (
+              <CodeEditor
                content={
                  isCurrentVersion
                    ? block.content
@ -431,16 +460,31 @@ function PureBlock({
                }
                isCurrentVersion={isCurrentVersion}
                currentVersionIndex={currentVersionIndex}
+                suggestions={suggestions ?? []}
                status={block.status}
                saveContent={saveContent}
-                suggestions={isCurrentVersion ? (suggestions ?? []) : []}
              />
-            ) : (
-              <DiffView
-                oldContent={getDocumentContentById(currentVersionIndex - 1)}
-                newContent={getDocumentContentById(currentVersionIndex)}
-              />
-            )}
+            ) : 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)}
+                />
+              )
+            ) : null}

            {suggestions ? (
              <div className="md:hidden h-dvh w-12 shrink-0" />
@ -455,6 +499,7 @@ function PureBlock({
                  isLoading={isLoading}
                  stop={stop}
                  setMessages={setMessages}
+                  blockKind={block.kind}
                />
              )}
            </AnimatePresence>
@ -471,6 +516,13 @@ function PureBlock({
            />
          )}
        </AnimatePresence>
+
+        <AnimatePresence>
+          <Console
+            consoleOutputs={consoleOutputs}
+            setConsoleOutputs={setConsoleOutputs}
+          />
+        </AnimatePresence>
      </motion.div>
    </motion.div>
  );
--- a/components/chat.tsx
+++ b/components/chat.tsx
@ -58,6 +58,7 @@ export function Chat({
  const [block, setBlock] = useState<UIBlock>({
    documentId: 'init',
    content: '',
+    kind: 'text',
    title: '',
    status: 'idle',
    isVisible: false,
--- a/components/code-block.tsx
+++ b/components/code-block.tsx
@ -0,0 +1,59 @@
+'use client';
+
+import { useCallback, useState } from 'react';
+import { CodeIcon, LoaderIcon, PlayIcon, PythonIcon } from './icons';
+import { Button } from './ui/button';
+import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
+import { cn } from '@/lib/utils';
+
+interface CodeBlockProps {
+  node: any;
+  inline: boolean;
+  className: string;
+  children: any;
+}
+
+export function CodeBlock({
+  node,
+  inline,
+  className,
+  children,
+  ...props
+}: CodeBlockProps) {
+  const [output, setOutput] = useState<string | null>(null);
+  const [pyodide, setPyodide] = useState<any>(null);
+  const match = /language-(\w+)/.exec(className || '');
+  const isPython = match && match[1] === 'python';
+  const codeContent = String(children).replace(/\n$/, '');
+  const [tab, setTab] = useState<'code' | 'run'>('code');
+
+  if (!inline) {
+    return (
+      <div className="not-prose flex flex-col">
+        {tab === 'code' && (
+          <pre
+            {...props}
+            className={`text-sm w-full overflow-x-auto dark:bg-zinc-900 p-4 border border-zinc-200 dark:border-zinc-700 rounded-xl dark:text-zinc-50 text-zinc-900`}
+          >
+            <code className="whitespace-pre-wrap break-words">{children}</code>
+          </pre>
+        )}
+
+        {tab === 'run' && output && (
+          <div className="text-sm w-full overflow-x-auto bg-zinc-800 dark:bg-zinc-900 p-4 border border-zinc-200 dark:border-zinc-700 border-t-0 rounded-b-xl text-zinc-50">
+            <code>{output}</code>
+          </div>
+        )}
+      </div>
+    );
+  } else {
+    return (
+      <code
+        className={`${className} text-sm bg-zinc-100 dark:bg-zinc-800 py-0.5 px-1 rounded-md`}
+        {...props}
+      >
+        {children}
+      </code>
+    );
+  }
+}
--- a/components/code-editor.tsx
+++ b/components/code-editor.tsx
@ -0,0 +1,108 @@
+'use client';
+
+import { EditorView } from '@codemirror/view';
+import { EditorState, Transaction } from '@codemirror/state';
+import { python } from '@codemirror/lang-python';
+import { oneDark } from '@codemirror/theme-one-dark';
+import { basicSetup } from 'codemirror';
+import React, { memo, useEffect, useRef } from 'react';
+import { Suggestion } from '@/lib/db/schema';
+
+type EditorProps = {
+  content: string;
+  saveContent: (updatedContent: string, debounce: boolean) => void;
+  status: 'streaming' | 'idle';
+  isCurrentVersion: boolean;
+  currentVersionIndex: number;
+  suggestions: Array<Suggestion>;
+};
+
+function PureCodeEditor({ content, saveContent, status }: EditorProps) {
+  const containerRef = useRef<HTMLDivElement>(null);
+  const editorRef = useRef<EditorView | null>(null);
+
+  useEffect(() => {
+    if (containerRef.current && !editorRef.current) {
+      const startState = EditorState.create({
+        doc: content,
+        extensions: [basicSetup, python(), oneDark],
+      });
+
+      editorRef.current = new EditorView({
+        state: startState,
+        parent: containerRef.current,
+      });
+    }
+
+    return () => {
+      if (editorRef.current) {
+        editorRef.current.destroy();
+        editorRef.current = null;
+      }
+    };
+  }, []);
+
+  useEffect(() => {
+    if (editorRef.current) {
+      const updateListener = EditorView.updateListener.of((update) => {
+        if (update.docChanged) {
+          const transaction = update.transactions.find(
+            (tr) => !tr.annotation(Transaction.remote),
+          );
+
+          if (transaction) {
+            const newContent = update.state.doc.toString();
+            saveContent(newContent, true);
+          }
+        }
+      });
+
+      const newState = EditorState.create({
+        doc: editorRef.current.state.doc,
+        extensions: [basicSetup, python(), oneDark, updateListener],
+      });
+
+      editorRef.current.setState(newState);
+    }
+  }, [saveContent]);
+
+  useEffect(() => {
+    if (editorRef.current && content) {
+      const currentContent = editorRef.current.state.doc.toString();
+
+      if (status === 'streaming' || currentContent !== content) {
+        const transaction = editorRef.current.state.update({
+          changes: {
+            from: 0,
+            to: currentContent.length,
+            insert: content,
+          },
+          annotations: [Transaction.remote.of(true)],
+        });
+
+        editorRef.current.dispatch(transaction);
+      }
+    }
+  }, [content, status]);
+
+  return (
+    <div
+      className="relative not-prose w-full pb-[calc(80dvh)] text-sm"
+      ref={containerRef}
+    />
+  );
+}
+
+function areEqual(prevProps: EditorProps, nextProps: EditorProps) {
+  if (prevProps.suggestions !== nextProps.suggestions) return false;
+  if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex)
+    return false;
+  if (prevProps.isCurrentVersion !== nextProps.isCurrentVersion) return false;
+  if (prevProps.status === 'streaming' && nextProps.status === 'streaming')
+    return false;
+  if (prevProps.content !== nextProps.content) return false;
+
+  return true;
+}
+
+export const CodeEditor = memo(PureCodeEditor, areEqual);
--- a/components/console.tsx
+++ b/components/console.tsx
@ -0,0 +1,126 @@
+import { TerminalWindowIcon, LoaderIcon, CrossSmallIcon } from './icons';
+import { Button } from './ui/button';
+import {
+  Dispatch,
+  SetStateAction,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from 'react';
+import { ConsoleOutput } from './block';
+import { cn } from '@/lib/utils';
+
+interface ConsoleProps {
+  consoleOutputs: Array<ConsoleOutput>;
+  setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
+}
+
+export function Console({ consoleOutputs, setConsoleOutputs }: ConsoleProps) {
+  const [height, setHeight] = useState<number>(300);
+  const [isResizing, setIsResizing] = useState(false);
+  const consoleEndRef = useRef<HTMLDivElement>(null);
+
+  const minHeight = 100;
+  const maxHeight = 800;
+
+  const startResizing = useCallback(() => {
+    setIsResizing(true);
+  }, []);
+
+  const stopResizing = useCallback(() => {
+    setIsResizing(false);
+  }, []);
+
+  const resize = useCallback(
+    (e: MouseEvent) => {
+      if (isResizing) {
+        const newHeight = window.innerHeight - e.clientY;
+        if (newHeight >= minHeight && newHeight <= maxHeight) {
+          setHeight(newHeight);
+        }
+      }
+    },
+    [isResizing],
+  );
+
+  useEffect(() => {
+    window.addEventListener('mousemove', resize);
+    window.addEventListener('mouseup', stopResizing);
+    return () => {
+      window.removeEventListener('mousemove', resize);
+      window.removeEventListener('mouseup', stopResizing);
+    };
+  }, [resize, stopResizing]);
+
+  useEffect(() => {
+    consoleEndRef.current?.scrollIntoView({ behavior: 'smooth' });
+  }, [consoleOutputs]);
+
+  return consoleOutputs.length > 0 ? (
+    <>
+      <div
+        className="h-2 w-full fixed cursor-ns-resize z-50"
+        onMouseDown={startResizing}
+        style={{ bottom: height - 4 }}
+      />
+
+      <div
+        className={cn(
+          'fixed flex flex-col bottom-0 dark:bg-zinc-900 bg-zinc-50 w-full border-t z-40 overflow-y-scroll dark:border-zinc-700 border-zinc-200',
+          {
+            'select-none': isResizing,
+          },
+        )}
+        style={{ height }}
+      >
+        <div className="flex flex-row justify-between items-center w-full h-fit border-b dark:border-zinc-700 border-zinc-200 px-2 py-1 sticky top-0 z-50 bg-muted">
+          <div className="text-sm pl-2 dark:text-zinc-50 text-zinc-800 flex flex-row gap-3 items-center">
+            <div className="text-muted-foreground">
+              <TerminalWindowIcon />
+            </div>
+            <div>Console</div>
+          </div>
+          <Button
+            variant="ghost"
+            className="size-fit p-1 hover:dark:bg-zinc-700 hover:bg-zinc-200"
+            size="icon"
+            onClick={() => setConsoleOutputs([])}
+          >
+            <CrossSmallIcon />
+          </Button>
+        </div>
+
+        <div>
+          {consoleOutputs.map((consoleOutput, index) => (
+            <div
+              key={consoleOutput.id}
+              className="px-4 py-2 flex flex-row text-sm border-b dark:border-zinc-700 border-zinc-200 dark:bg-zinc-900 bg-zinc-50 font-mono"
+            >
+              <div
+                className={cn('w-12 shrink-0', {
+                  'text-muted-foreground':
+                    consoleOutput.status === 'in_progress',
+                  'text-emerald-500': consoleOutput.status === 'completed',
+                  'text-red-400': consoleOutput.status === 'failed',
+                })}
+              >
+                [{index + 1}]
+              </div>
+              {consoleOutput.status === 'in_progress' ? (
+                <div className="animate-spin size-fit self-center">
+                  <LoaderIcon />
+                </div>
+              ) : (
+                <div className="dark:text-zinc-50 text-zinc-900 whitespace-pre-line">
+                  {consoleOutput.content}
+                </div>
+              )}
+            </div>
+          ))}
+          <div ref={consoleEndRef} />
+        </div>
+      </div>
+    </>
+  ) : null;
+}
--- a/components/document.tsx
+++ b/components/document.tsx
@ -1,6 +1,6 @@
 import { memo, type SetStateAction } from 'react';

-import type { UIBlock } from './block';
+import type { BlockKind, UIBlock } from './block';
 import { FileIcon, LoaderIcon, MessageIcon, PencilEditIcon } from './icons';
 import { toast } from 'sonner';

@ -24,7 +24,7 @@ const getActionText = (

 interface DocumentToolResultProps {
  type: 'create' | 'update' | 'request-suggestions';
-  result: { id: string; title: string };
+  result: { id: string; title: string; kind: BlockKind };
  block: UIBlock;
  setBlock: (value: SetStateAction<UIBlock>) => void;
  isReadonly: boolean;
@ -59,6 +59,7 @@ function PureDocumentToolResult({

        setBlock({
          documentId: result.id,
+          kind: result.kind,
          content: '',
          title: result.title,
          isVisible: true,
--- a/components/icons.tsx
+++ b/components/icons.tsx
@ -627,6 +627,23 @@ export const CrossIcon = ({ size = 16 }: { size?: number }) => (
  </svg>
 );

+export const CrossSmallIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: 'currentcolor' }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M9.96966 11.0303L10.5 11.5607L11.5607 10.5L11.0303 9.96966L9.06065 7.99999L11.0303 6.03032L11.5607 5.49999L10.5 4.43933L9.96966 4.96966L7.99999 6.93933L6.03032 4.96966L5.49999 4.43933L4.43933 5.49999L4.96966 6.03032L6.93933 7.99999L4.96966 9.96966L4.43933 10.5L5.49999 11.5607L6.03032 11.0303L7.99999 9.06065L9.96966 11.0303Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
 export const UndoIcon = ({ size = 16 }: { size?: number }) => (
  <svg
    height={size}
@ -931,3 +948,138 @@ export const ShareIcon = ({ size = 16 }: { size?: number }) => {
    </svg>
  );
 };
+
+export const CodeIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M4.21969 12.5303L4.75002 13.0607L5.81068 12L5.28035 11.4697L1.81068 7.99999L5.28035 4.53032L5.81068 3.99999L4.75002 2.93933L4.21969 3.46966L0.39647 7.29289C0.00594562 7.68341 0.00594562 8.31658 0.39647 8.7071L4.21969 12.5303ZM11.7804 12.5303L11.25 13.0607L10.1894 12L10.7197 11.4697L14.1894 7.99999L10.7197 4.53032L10.1894 3.99999L11.25 2.93933L11.7804 3.46966L15.6036 7.29289C15.9941 7.68341 15.9941 8.31658 15.6036 8.7071L11.7804 12.5303Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const PlayIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M13.4549 7.22745L13.3229 7.16146L2.5 1.74999L2.4583 1.72914L1.80902 1.4045L1.3618 1.18089C1.19558 1.09778 1 1.21865 1 1.4045L1 1.9045L1 2.63041L1 2.67704L1 13.3229L1 13.3696L1 14.0955L1 14.5955C1 14.7813 1.19558 14.9022 1.3618 14.8191L1.80902 14.5955L2.4583 14.2708L2.5 14.25L13.3229 8.83852L13.4549 8.77253L14.2546 8.37267L14.5528 8.2236C14.737 8.13147 14.737 7.86851 14.5528 7.77638L14.2546 7.62731L13.4549 7.22745ZM11.6459 7.99999L2.5 3.42704L2.5 12.5729L11.6459 7.99999Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const PythonIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        d="M7.90474 0.00013087C7.24499 0.00316291 6.61494 0.0588153 6.06057 0.15584C4.42745 0.441207 4.13094 1.0385 4.13094 2.14002V3.59479H7.9902V4.07971H4.13094H2.68259C1.56099 4.07971 0.578874 4.7465 0.271682 6.01496C-0.0826597 7.4689 -0.0983767 8.37619 0.271682 9.89434C0.546012 11.0244 1.20115 11.8296 2.32276 11.8296H3.64966V10.0856C3.64966 8.82574 4.75179 7.71441 6.06057 7.71441H9.91533C10.9884 7.71441 11.845 6.84056 11.845 5.77472V2.14002C11.845 1.10556 10.9626 0.328487 9.91533 0.15584C9.25237 0.046687 8.56448 -0.00290121 7.90474 0.00013087ZM5.81768 1.17017C6.21631 1.17017 6.54185 1.49742 6.54185 1.89978C6.54185 2.30072 6.21631 2.62494 5.81768 2.62494C5.41761 2.62494 5.09351 2.30072 5.09351 1.89978C5.09351 1.49742 5.41761 1.17017 5.81768 1.17017Z"
+        fill="currentColor"
+      ></path>
+      <path
+        d="M12.3262 4.07971V5.77472C12.3262 7.08883 11.1997 8.19488 9.91525 8.19488H6.06049C5.0046 8.19488 4.13086 9.0887 4.13086 10.1346V13.7693C4.13086 14.8037 5.04033 15.4122 6.06049 15.709C7.28211 16.0642 8.45359 16.1285 9.91525 15.709C10.8868 15.4307 11.8449 14.8708 11.8449 13.7693V12.3145H7.99012V11.8296H11.8449H13.7745C14.8961 11.8296 15.3141 11.0558 15.7041 9.89434C16.1071 8.69865 16.0899 7.5488 15.7041 6.01495C15.4269 4.91058 14.8975 4.07971 13.7745 4.07971H12.3262ZM10.1581 13.2843C10.5582 13.2843 10.8823 13.6086 10.8823 14.0095C10.8823 14.4119 10.5582 14.7391 10.1581 14.7391C9.7595 14.7391 9.43397 14.4119 9.43397 14.0095C9.43397 13.6086 9.7595 13.2843 10.1581 13.2843Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const TerminalWindowIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M1.5 2.5H14.5V12.5C14.5 13.0523 14.0523 13.5 13.5 13.5H2.5C1.94772 13.5 1.5 13.0523 1.5 12.5V2.5ZM0 1H1.5H14.5H16V2.5V12.5C16 13.8807 14.8807 15 13.5 15H2.5C1.11929 15 0 13.8807 0 12.5V2.5V1ZM4 11.1339L4.44194 10.6919L6.51516 8.61872C6.85687 8.27701 6.85687 7.72299 6.51517 7.38128L4.44194 5.30806L4 4.86612L3.11612 5.75L3.55806 6.19194L5.36612 8L3.55806 9.80806L3.11612 10.25L4 11.1339ZM8 9.75494H8.6225H11.75H12.3725V10.9999H11.75H8.6225H8V9.75494Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const TerminalIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M1.53035 12.7804L1.00002 13.3108L-0.0606384 12.2501L0.469692 11.7198L4.18936 8.00011L0.469692 4.28044L-0.0606384 3.75011L1.00002 2.68945L1.53035 3.21978L5.60358 7.29301C5.9941 7.68353 5.9941 8.3167 5.60357 8.70722L1.53035 12.7804ZM8.75002 12.5001H8.00002V14.0001H8.75002H15.25H16V12.5001H15.25H8.75002Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const ClockRewind = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M7.96452 2.5C11.0257 2.5 13.5 4.96643 13.5 8C13.5 11.0336 11.0257 13.5 7.96452 13.5C6.12055 13.5 4.48831 12.6051 3.48161 11.2273L3.03915 10.6217L1.828 11.5066L2.27046 12.1122C3.54872 13.8617 5.62368 15 7.96452 15C11.8461 15 15 11.87 15 8C15 4.13001 11.8461 1 7.96452 1C5.06835 1 2.57851 2.74164 1.5 5.23347V3.75V3H0V3.75V7.25C0 7.66421 0.335786 8 0.75 8H3.75H4.5V6.5H3.75H2.63724C3.29365 4.19393 5.42843 2.5 7.96452 2.5ZM8.75 5.25V4.5H7.25V5.25V7.8662C7.25 8.20056 7.4171 8.51279 7.6953 8.69825L9.08397 9.62404L9.70801 10.0401L10.5401 8.79199L9.91603 8.37596L8.75 7.59861V5.25Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
+
+export const LogsIcon = ({ size = 16 }: { size?: number }) => {
+  return (
+    <svg
+      height={size}
+      strokeLinejoin="round"
+      viewBox="0 0 16 16"
+      width={size}
+      style={{ color: 'currentcolor' }}
+    >
+      <path
+        fillRule="evenodd"
+        clipRule="evenodd"
+        d="M9 2H9.75H14.25H15V3.5H14.25H9.75H9V2ZM9 12.5H9.75H14.25H15V14H14.25H9.75H9V12.5ZM9.75 7.25H9V8.75H9.75H14.25H15V7.25H14.25H9.75ZM1 12.5H1.75H2.25H3V14H2.25H1.75H1V12.5ZM1.75 2H1V3.5H1.75H2.25H3V2H2.25H1.75ZM1 7.25H1.75H2.25H3V8.75H2.25H1.75H1V7.25ZM5.75 12.5H5V14H5.75H6.25H7V12.5H6.25H5.75ZM5 2H5.75H6.25H7V3.5H6.25H5.75H5V2ZM5.75 7.25H5V8.75H5.75H6.25H7V7.25H6.25H5.75Z"
+        fill="currentColor"
+      ></path>
+    </svg>
+  );
+};
--- a/components/markdown.tsx
+++ b/components/markdown.tsx
@ -2,29 +2,13 @@ import Link from 'next/link';
 import React, { memo } from 'react';
 import ReactMarkdown, { type Components } from 'react-markdown';
 import remarkGfm from 'remark-gfm';
+import { CodeBlock } from './code-block';

 const NonMemoizedMarkdown = ({ children }: { children: string }) => {
  const components: Partial<Components> = {
    // @ts-expect-error
-    code: ({ node, inline, className, children, ...props }) => {
-      const match = /language-(\w+)/.exec(className || '');
-      return !inline && match ? (
-        // @ts-expect-error
-        <pre
-          {...props}
-          className={`${className} text-sm w-[80dvw] md:max-w-[500px] overflow-x-scroll bg-zinc-100 p-3 rounded-lg mt-2 dark:bg-zinc-800`}
-        >
-          <code className={match[1]}>{children}</code>
-        </pre>
-      ) : (
-        <code
-          className={`${className} text-sm bg-zinc-100 dark:bg-zinc-800 py-0.5 px-1 rounded-md`}
-          {...props}
-        >
-          {children}
-        </code>
-      );
-    },
+    code: CodeBlock,
+    pre: ({ children }) => <>{children}</>,
    ol: ({ node, children, ...props }) => {
      return (
        <ol className="list-decimal list-outside ml-4" {...props}>
--- a/components/message.tsx
+++ b/components/message.tsx
@ -64,7 +64,7 @@ const PurePreviewMessage = ({
        )}
      >
        {message.role === 'assistant' && (
-          <div className="size-8 flex items-center rounded-full justify-center ring-1 shrink-0 ring-border">
+          <div className="size-8 flex items-center rounded-full justify-center ring-1 shrink-0 ring-border bg-background">
            <SparklesIcon size={14} />
          </div>
        )}
--- a/components/suggested-actions.tsx
+++ b/components/suggested-actions.tsx
@ -21,9 +21,9 @@ function PureSuggestedActions({ chatId, append }: SuggestedActionsProps) {
      action: 'What is the weather in San Francisco?',
    },
    {
-      title: 'Help me draft an essay',
-      label: 'about Silicon Valley',
-      action: 'Help me draft a short essay about Silicon Valley',
+      title: 'Help me write code to',
+      label: 'demonstrate the djikstra algorithm!',
+      action: 'Help me write code to demonstrate the djikstra algorithm!',
    },
  ];

--- a/components/suggestion.tsx
+++ b/components/suggestion.tsx
@ -8,13 +8,17 @@ import type { UISuggestion } from '@/lib/editor/suggestions';

 import { CrossIcon, MessageIcon } from './icons';
 import { Button } from './ui/button';
+import { cn } from '@/lib/utils';
+import { BlockKind } from './block';

 export const Suggestion = ({
  suggestion,
  onApply,
+  blockKind,
 }: {
  suggestion: UISuggestion;
  onApply: () => void;
+  blockKind: BlockKind;
 }) => {
  const [isExpanded, setIsExpanded] = useState(false);
  const { width: windowWidth } = useWindowSize();
@ -23,7 +27,10 @@ export const Suggestion = ({
    <AnimatePresence>
      {!isExpanded ? (
        <motion.div
-          className="absolute cursor-pointer text-muted-foreground -right-8 p-1"
+          className={cn('cursor-pointer text-muted-foreground p-1', {
+            'absolute -right-8': blockKind === 'text',
+            'sticky top-0 right-4': blockKind === 'code',
+          })}
          onClick={() => {
            setIsExpanded(true);
          }}
@ -34,7 +41,7 @@ export const Suggestion = ({
      ) : (
        <motion.div
          key={suggestion.id}
-          className="absolute bg-background p-3 flex flex-col gap-3 rounded-2xl border text-sm w-56 shadow-xl z-50 -right-12 md:-right-16"
+          className="absolute bg-background p-3 flex flex-col gap-3 rounded-2xl border text-sm w-56 shadow-xl z-50 -right-12 md:-right-16 font-sans"
          transition={{ type: 'spring', stiffness: 500, damping: 30 }}
          initial={{ opacity: 0, y: -10 }}
          animate={{ opacity: 1, y: -20 }}
--- a/components/toolbar.tsx
+++ b/components/toolbar.tsx
@ -28,15 +28,25 @@ import { sanitizeUIMessages } from '@/lib/utils';

 import {
  ArrowUpIcon,
+  CodeIcon,
+  FileIcon,
+  LogsIcon,
  MessageIcon,
  PenIcon,
  StopIcon,
  SummarizeIcon,
+  TerminalIcon,
 } from './icons';
-import equal from 'fast-deep-equal';
+import { BlockKind } from './block';

 type ToolProps = {
-  type: 'final-polish' | 'request-suggestions' | 'adjust-reading-level';
+  type:
+    | 'final-polish'
+    | 'request-suggestions'
+    | 'adjust-reading-level'
+    | 'code-review'
+    | 'add-comments'
+    | 'add-logs';
  description: string;
  icon: JSX.Element;
  selectedTool: string | null;
@ -99,6 +109,20 @@ const Tool = ({
            '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);
      }
    }
@ -260,6 +284,51 @@ 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 />,
+    },
+  ],
+};
+
 export const Tools = ({
  isToolbarVisible,
  selectedTool,
@ -267,6 +336,7 @@ export const Tools = ({
  append,
  isAnimating,
  setIsToolbarVisible,
+  blockKind,
 }: {
  isToolbarVisible: boolean;
  selectedTool: string | null;
@ -277,7 +347,10 @@ export const Tools = ({
  ) => Promise<string | null | undefined>;
  isAnimating: boolean;
  setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
+  blockKind: 'text' | 'code';
 }) => {
+  const [primaryTool, ...secondaryTools] = toolsByBlockKind[blockKind];
+
  return (
    <motion.div
      className="flex flex-col"
@ -286,35 +359,25 @@ export const Tools = ({
      exit={{ opacity: 0, scale: 0.95 }}
    >
      <AnimatePresence>
-        {isToolbarVisible && (
-          <>
+        {isToolbarVisible &&
+          secondaryTools.map((secondaryTool) => (
            <Tool
-              type="adjust-reading-level"
-              description="Adjust reading level"
-              icon={<SummarizeIcon />}
+              key={secondaryTool.type}
+              type={secondaryTool.type}
+              description={secondaryTool.description}
+              icon={secondaryTool.icon}
              selectedTool={selectedTool}
              setSelectedTool={setSelectedTool}
              append={append}
              isAnimating={isAnimating}
            />
-
-            <Tool
-              type="request-suggestions"
-              description="Request suggestions"
-              icon={<MessageIcon />}
-              selectedTool={selectedTool}
-              setSelectedTool={setSelectedTool}
-              append={append}
-              isAnimating={isAnimating}
-            />
-          </>
-        )}
+          ))}
      </AnimatePresence>

      <Tool
-        type="final-polish"
-        description="Add final polish"
-        icon={<PenIcon />}
+        type={primaryTool.type}
+        description={primaryTool.description}
+        icon={primaryTool.icon}
        selectedTool={selectedTool}
        setSelectedTool={setSelectedTool}
        isToolbarVisible={isToolbarVisible}
@ -333,6 +396,7 @@ const PureToolbar = ({
  isLoading,
  stop,
  setMessages,
+  blockKind,
 }: {
  isToolbarVisible: boolean;
  setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
@ -343,6 +407,7 @@ const PureToolbar = ({
  ) => Promise<string | null | undefined>;
  stop: () => void;
  setMessages: Dispatch<SetStateAction<Message[]>>;
+  blockKind: 'text' | 'code';
 }) => {
  const toolbarRef = useRef<HTMLDivElement>(null);
  const timeoutRef = useRef<ReturnType<typeof setTimeout>>();
@ -404,7 +469,7 @@ const PureToolbar = ({
              : {
                  opacity: 1,
                  y: 0,
-                  height: 3 * 45,
+                  height: toolsByBlockKind[blockKind].length * 47,
                  transition: { delay: 0 },
                  scale: 1,
                }
@ -461,6 +526,7 @@ const PureToolbar = ({
            selectedTool={selectedTool}
            setIsToolbarVisible={setIsToolbarVisible}
            setSelectedTool={setSelectedTool}
+            blockKind={blockKind}
          />
        )}
      </motion.div>
@ -471,6 +537,7 @@ const PureToolbar = ({
 export const Toolbar = memo(PureToolbar, (prevProps, nextProps) => {
  if (prevProps.isLoading !== nextProps.isLoading) return false;
  if (prevProps.isToolbarVisible !== nextProps.isToolbarVisible) return false;
+  if (prevProps.blockKind !== nextProps.blockKind) return false;

  return true;
 });
--- a/components/use-block-stream.tsx
+++ b/components/use-block-stream.tsx
@ -4,18 +4,20 @@ import { useSWRConfig } from 'swr';

 import type { Suggestion } from '@/lib/db/schema';

-import type { UIBlock } from './block';
+import type { BlockKind, UIBlock } from './block';
 import { useUserMessageId } from '@/hooks/use-user-message-id';

 type StreamingDelta = {
  type:
    | 'text-delta'
+    | 'code-delta'
    | 'title'
    | 'id'
    | 'suggestion'
    | 'clear'
    | 'finish'
-    | 'user-message-id';
+    | 'user-message-id'
+    | 'kind';

  content: string | Suggestion;
 };
@ -67,6 +69,12 @@ export function useBlockStream({
            title: delta.content as string,
          };

+        case 'kind':
+          return {
+            ...draftBlock,
+            kind: delta.content as BlockKind,
+          };
+
        case 'text-delta':
          return {
            ...draftBlock,
@ -80,6 +88,19 @@ export function useBlockStream({
            status: 'streaming',
          };

+        case 'code-delta':
+          return {
+            ...draftBlock,
+            content: delta.content as string,
+            isVisible:
+              draftBlock.status === 'streaming' &&
+              draftBlock.content.length > 20 &&
+              draftBlock.content.length < 30
+                ? true
+                : draftBlock.isVisible,
+            status: 'streaming',
+          };
+
        case 'suggestion':
          setTimeout(() => {
            setOptimisticSuggestions((currentSuggestions) => [
@ -107,5 +128,5 @@ export function useBlockStream({
          return draftBlock;
      }
    });
-  }, [streamingData, setBlock]);
+  }, [streamingData, setBlock, setUserMessageIdFromServer]);
 }