feat: support image outputs during code execution (vercel#674)

Co-authored-by: Bailey Simrell <[email protected]>
ziyaaktas · Jan 7, 2025 · 50bd460 · 50bd460
1 parent 50fbc0d
commit 50bd460
Show file tree

Hide file tree

Showing 7 changed files with 313 additions and 117 deletions.
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
@@ -254,7 +254,7 @@ export async function POST(request: Request) {
               if (document.kind === 'text') {
                 const { fullStream } = streamText({
                   model: customModel(model.apiIdentifier),
-                  system: updateDocumentPrompt(currentContent),
+                  system: updateDocumentPrompt(currentContent, 'text'),
                   prompt: description,
                   experimental_providerMetadata: {
                     openai: {
@@ -284,7 +284,7 @@ export async function POST(request: Request) {
               } else if (document.kind === 'code') {
                 const { fullStream } = streamObject({
                   model: customModel(model.apiIdentifier),
-                  system: updateDocumentPrompt(currentContent),
+                  system: updateDocumentPrompt(currentContent, 'code'),
                   prompt: description,
                   schema: z.object({
                     code: z.string(),

diff --git a/components/block-actions.tsx b/components/block-actions.tsx
@@ -1,18 +1,12 @@
-import { cn, generateUUID } from '@/lib/utils';
-import { ClockRewind, CopyIcon, PlayIcon, RedoIcon, UndoIcon } from './icons';
+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 { useCopyToClipboard } from 'usehooks-ts';
 import { toast } from 'sonner';
 import { ConsoleOutput, UIBlock } from './block';
-import {
-  Dispatch,
-  memo,
-  SetStateAction,
-  startTransition,
-  useCallback,
-  useState,
-} from 'react';
+import { Dispatch, memo, SetStateAction } from 'react';
+import { RunCodeButton } from './run-code-button';
 
 interface BlockActionsProps {
   block: UIBlock;
@@ -23,98 +17,6 @@ interface BlockActionsProps {
   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({
   block,
   handleVersionChange,

diff --git a/components/block.tsx b/components/block.tsx
@@ -52,10 +52,15 @@ export interface UIBlock {
   };
 }
 
+export interface ConsoleOutputContent {
+  type: 'text' | 'image';
+  value: string;
+}
+
 export interface ConsoleOutput {
   id: string;
-  status: 'in_progress' | 'completed' | 'failed';
-  content: string | null;
+  status: 'in_progress' | 'loading_packages' | 'completed' | 'failed';
+  contents: Array<ConsoleOutputContent>;
 }
 
 function PureBlock({

diff --git a/components/console.tsx b/components/console.tsx
@@ -101,21 +101,51 @@ export function Console({ consoleOutputs, setConsoleOutputs }: ConsoleProps) {
             >
               <div
                 className={cn('w-12 shrink-0', {
-                  'text-muted-foreground':
-                    consoleOutput.status === 'in_progress',
+                  'text-muted-foreground': [
+                    'in_progress',
+                    'loading_packages',
+                  ].includes(consoleOutput.status),
                   '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 />
+              {['in_progress', 'loading_packages'].includes(
+                consoleOutput.status,
+              ) ? (
+                <div className="flex flex-row gap-2">
+                  <div className="animate-spin size-fit self-center">
+                    <LoaderIcon />
+                  </div>
+                  <div className="text-muted-foreground">
+                    {consoleOutput.status === 'in_progress'
+                      ? 'Initializing...'
+                      : consoleOutput.status === 'loading_packages'
+                        ? 'Loading Packages...'
+                        : null}
+                  </div>
                 </div>
               ) : (
-                <div className="dark:text-zinc-50 text-zinc-900 whitespace-pre-line">
-                  {consoleOutput.content}
+                <div className="dark:text-zinc-50 text-zinc-900 w-full flex flex-col gap-2">
+                  {consoleOutput.contents.map((content, index) =>
+                    content.type === 'image' ? (
+                      <picture key={`${consoleOutput.id}-${index}`}>
+                        <img
+                          src={content.value}
+                          alt="output"
+                          className="rounded-md max-w-[600px] w-full"
+                        />
+                      </picture>
+                    ) : (
+                      <div
+                        key={`${consoleOutput.id}-${index}`}
+                        className="whitespace-pre-line"
+                      >
+                        {content.value}
+                      </div>
+                    ),
+                  )}
                 </div>
               )}
             </div>