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

Co-authored-by: Bailey Simrell <baileysimrell@gmail.com>

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(),

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,

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({

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>

components/run-code-button.tsx

@@ -0,0 +1,246 @@
+import { generateUUID } from '@/lib/utils';
+import {
+  type Dispatch,
+  type SetStateAction,
+  startTransition,
+  useCallback,
+  useState,
+  useEffect,
+} from 'react';
+import type { ConsoleOutput, UIBlock } from './block';
+import { Button } from './ui/button';
+import { PlayIcon } from './icons';
+
+function detectPythonImports(code: string, pyodide: any): Set<string> {
+  const imports = new Set<string>();
+
+  const importPatterns = [
+    /import\s+(\w+)(?:\s+as\s+\w+)?/g,
+    /from\s+(\w+(?:\.\w+)*)\s+import/g,
+    /import\s+(\w+(?:\.\w+)*)/g,
+  ];
+
+  for (const pattern of importPatterns) {
+    let match: RegExpExecArray | null;
+
+    while (true) {
+      match = pattern.exec(code);
+      if (match === null) break;
+
+      const rootPackage = match[1].split('.')[0];
+      if (rootPackage) {
+        imports.add(rootPackage);
+      }
+    }
+  }
+
+  // Get standard libraries dynamically when initializing Pyodide
+  let standardLibs = new Set<string>();
+  if (pyodide) {
+    const stdLibModules = pyodide.runPython(`
+      import sys
+      list(sys.stdlib_module_names)
+    `);
+    standardLibs = new Set(stdLibModules);
+  }
+
+  return new Set(Array.from(imports).filter((pkg) => !standardLibs.has(pkg)));
+}
+
+const OUTPUT_HANDLERS = {
+  matplotlib: `
+    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 RunCodeButton({
+  block,
+  setConsoleOutputs,
+}: {
+  block: UIBlock;
+  setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
+}) {
+  const [pyodide, setPyodide] = useState<any>(null);
+  const isPython = true;
+  const codeContent = block.content;
+
+  const loadAndRunPython = useCallback(async () => {
+    const runId = generateUUID();
+
+    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 loadPyodide({
+            indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
+          });
+
+          setPyodide(newPyodideInstance);
+          currentPyodideInstance = newPyodideInstance;
+        }
+
+        await currentPyodideInstance.runPythonAsync(`
+          import sys, io, gc
+          import base64
+
+          sys.stdout = io.StringIO()
+        `);
+
+        // Detect and load required packages
+        const requiredPackages = detectPythonImports(
+          codeContent,
+          currentPyodideInstance,
+        );
+
+        if (requiredPackages.size > 0) {
+          setConsoleOutputs((outputs) => [
+            ...outputs.filter((output) => output.id !== runId),
+            {
+              id: runId,
+              contents: [],
+              status: 'loading_packages',
+            },
+          ]);
+
+          await currentPyodideInstance.loadPackage(
+            Array.from(requiredPackages),
+          );
+        }
+
+        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);
+
+        const runOutput = await currentPyodideInstance.runPythonAsync(
+          `sys.stdout.getvalue()`,
+        );
+
+        const runOutputByLines: string[] = runOutput.split('\n');
+
+        setConsoleOutputs((outputs) => [
+          ...outputs.filter((output) => output.id !== runId),
+          {
+            id: generateUUID(),
+            contents: runOutputByLines
+              .filter((line) => line.trim().length)
+              .map((line) => ({
+                type: line.startsWith('data:image/png;base64')
+                  ? 'image'
+                  : 'text',
+                value: line,
+              })),
+            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
+            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={block.status === 'streaming'}
+    >
+      <PlayIcon size={18} /> Run
+    </Button>
+  );
+}

lib/ai/prompts.ts

@@ -1,3 +1,5 @@
+import { BlockKind } from '@/components/block';
+
 export const blocksPrompt = `
 Blocks is a special user interface mode that helps users with writing, editing, and other content creation tasks. When block is open, it is on the right side of the screen, while the conversation is on the left side. When creating or updating documents, changes are reflected in real-time on the blocks and visible to the user.
 
@@ -62,8 +64,20 @@ print(f"Factorial of 5 is: {factorial(5)}")
 \`\`\`
 `;
 
-export const updateDocumentPrompt = (currentContent: string | null) => `\
-Update the following contents of the document based on the given prompt.
+export const updateDocumentPrompt = (
+  currentContent: string | null,
+  type: BlockKind,
+) =>
+  type === 'text'
+    ? `\
+Improve the following contents of the document based on the given prompt.
 
 ${currentContent}
-`;
+`
+    : type === 'code'
+      ? `\
+Improve the following code snippet based on the given prompt.
+
+${currentContent}
+`
+      : '';

next.config.ts

@@ -1,7 +1,6 @@
 import type { NextConfig } from 'next';
 
 const nextConfig: NextConfig = {
-  /* config options here */
   experimental: {
     ppr: true,
   },