286 lines · 7.4 KB · TypeScript
RawDownload
import { toast } from "sonner";import { CodeEditor } from "@/components/chat/code-editor";import {  Console,  type ConsoleOutput,  type ConsoleOutputContent,} from "@/components/chat/console";import { Artifact } from "@/components/chat/create-artifact";import {  CopyIcon,  LogsIcon,  MessageIcon,  PlayIcon,  RedoIcon,  UndoIcon,} from "@/components/chat/icons";import { generateUUID } from "@/lib/utils"; 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;} type Metadata = {  outputs: ConsoleOutput[];}; export const codeArtifact = new Artifact<"code", Metadata>({  kind: "code",  description:    "Useful for code generation; Code execution is only available for python code.",  initialize: ({ setMetadata }) => {    setMetadata({      outputs: [],    });  },  onStreamPart: ({ streamPart, setArtifact }) => {    if (streamPart.type === "data-codeDelta") {      setArtifact((draftArtifact) => ({        ...draftArtifact,        content: streamPart.data,        isVisible:          draftArtifact.status === "streaming" &&          draftArtifact.content.length > 300 &&          draftArtifact.content.length < 310            ? true            : draftArtifact.isVisible,        status: "streaming",      }));    }  },  content: ({ metadata, setMetadata, ...props }) => {    return (      <>        <div className="relative min-h-[200px]">          <CodeEditor {...props} />        </div>         {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 outputContent: ConsoleOutputContent[] = [];         setMetadata((metadata) => ({          ...metadata,          outputs: [            ...metadata.outputs,            {              id: runId,              contents: [],              status: "in_progress",            },          ],        }));         try {          // @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) => {              outputContent.push({                type: output.startsWith("data:image/png;base64")                  ? "image"                  : "text",                value: output,              });            },          });           await currentPyodideInstance.loadPackagesFromImports(content, {            messageCallback: (message: string) => {              setMetadata((metadata) => ({                ...metadata,                outputs: [                  ...metadata.outputs.filter((output) => output.id !== runId),                  {                    id: runId,                    contents: [{ type: "text", value: message }],                    status: "loading_packages",                  },                ],              }));            },          });           const requiredHandlers = detectRequiredHandlers(content);          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(content);           setMetadata((metadata) => ({            ...metadata,            outputs: [              ...metadata.outputs.filter((output) => output.id !== runId),              {                id: runId,                contents: outputContent,                status: "completed",              },            ],          }));        } catch (error: unknown) {          setMetadata((metadata) => ({            ...metadata,            outputs: [              ...metadata.outputs.filter((output) => output.id !== runId),              {                id: runId,                contents: [                  {                    type: "text",                    value:                      error instanceof Error ? error.message : String(error),                  },                ],                status: "failed",              },            ],          }));        }      },    },    {      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 code to clipboard",      onClick: ({ content }) => {        navigator.clipboard.writeText(content);        toast.success("Copied to clipboard!");      },    },  ],  toolbar: [    {      icon: <MessageIcon />,      description: "Add comments",      onClick: ({ sendMessage }) => {        sendMessage({          role: "user",          parts: [            {              type: "text",              text: "Add comments to the code snippet for understanding",            },          ],        });      },    },    {      icon: <LogsIcon />,      description: "Add logs",      onClick: ({ sendMessage }) => {        sendMessage({          role: "user",          parts: [            {              type: "text",              text: "Add logs to the code snippet for debugging",            },          ],        });      },    },  ],});