feat: support image outputs during code execution (#674)
Co-authored-by: Bailey Simrell <baileysimrell@gmail.com>
This commit is contained in:
Jeremy
GitHub
parent
50fbc0dab2
commit
50bd4604b6
7 changed files with 313 additions and 117 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
246
components/run-code-button.tsx
Normal file
246
components/run-code-button.tsx
Normal file
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue