@marimo-team/islands 0.15.2 → 0.15.4

package/src/plugins/impl/chat/chat-ui.tsx

@@ -1,10 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { useChat } from "@ai-sdk/react";
+import { type UIMessage, useChat } from "@ai-sdk/react";
 import { ChatBubbleIcon } from "@radix-ui/react-icons";
 import { PopoverAnchor } from "@radix-ui/react-popover";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
-import type { Message } from "ai/react";
+import { DefaultChatTransport, type FileUIPart } from "ai";
 import { startCase } from "lodash-es";
 import {
   BotMessageSquareIcon,
@@ -18,6 +18,7 @@ import {
   X,
 } from "lucide-react";
 import React, { useEffect, useRef, useState } from "react";
+import { convertToFileUIPart } from "@/components/chat/chat-utils";
 import {
   type AdditionalCompletions,
   PromptInput,
@@ -49,12 +50,7 @@ import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
 import { ErrorBanner } from "../common/error-banner";
 import type { PluginFunctions } from "./ChatPlugin";
-import type {
-  ChatAttachment,
-  ChatConfig,
-  ChatMessage,
-  ChatRole,
-} from "./types";
+import type { ChatConfig, ChatMessage } from "./types";
 
 interface Props extends PluginFunctions {
   prompts: string[];
@@ -67,8 +63,9 @@ interface Props extends PluginFunctions {
 }
 
 export const Chatbot: React.FC<Props> = (props) => {
+  const [input, setInput] = useState("");
   const [config, setConfig] = useState<ChatConfig>(props.config);
-  const [files, setFiles] = useState<FileList | undefined>(undefined);
+  const [files, setFiles] = useState<File[] | undefined>(undefined);
   const fileInputRef = useRef<HTMLInputElement>(null);
   const formRef = useRef<HTMLFormElement>(null);
   const codeMirrorInputRef = useRef<ReactCodeMirrorRef>(null);
@@ -76,75 +73,86 @@ export const Chatbot: React.FC<Props> = (props) => {
 
   const { data: initialMessages } = useAsyncData(async () => {
     const chatMessages = await props.get_chat_history({});
-    const messages: Message[] = chatMessages.messages.map((message, idx) => ({
+    const messages: UIMessage[] = chatMessages.messages.map((message, idx) => ({
       id: idx.toString(),
       role: message.role,
-      content: message.content,
-      experimental_attachments: message.attachments,
+      parts: message.parts ?? [],
     }));
     return messages;
   }, []);
 
   const {
     messages,
+    sendMessage,
     setMessages,
-    input,
-    setInput,
-    handleSubmit,
     status,
     stop,
     error,
-    reload,
+    regenerate,
   } = useChat({
-    keepLastMessageOnError: true,
-    streamProtocol: "text",
-    fetch: async (_url, request) => {
-      const body = JSON.parse(request?.body as string) as {
-        messages: Message[];
-      };
-      try {
-        const response = await props.send_prompt({
-          messages: body.messages.map((m) => ({
-            role: m.role as ChatRole,
-            content: m.content,
-            attachments: m.experimental_attachments,
-          })),
-          config: {
-            max_tokens: config.max_tokens,
-            temperature: config.temperature,
-            top_p: config.top_p,
-            top_k: config.top_k,
-            frequency_penalty: config.frequency_penalty,
-            presence_penalty: config.presence_penalty,
-          },
-        });
-        return new Response(response);
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      } catch (error: any) {
-        // HACK: strip the error message to clean up the response
-        const strippedError = error.message
-          .split("failed with exception ")
-          .pop();
-        return new Response(strippedError, { status: 400 });
-      }
-    },
-    initialMessages: initialMessages,
-    onFinish: (message, { usage, finishReason }) => {
+    transport: new DefaultChatTransport({
+      fetch: async (
+        request: RequestInfo | URL,
+        init: RequestInit | undefined,
+      ) => {
+        if (init === undefined) {
+          return fetch(request);
+        }
+
+        const body = JSON.parse(init.body as unknown as string) as {
+          messages: UIMessage[];
+        };
+        try {
+          const messages = body.messages.map((m) => ({
+            role: m.role,
+            content: m.parts
+              ?.map((p) => ("text" in p ? p.text : ""))
+              .join("\n"),
+            parts: m.parts,
+          }));
+          const response = await props.send_prompt({
+            messages: messages,
+            config: {
+              max_tokens: config.max_tokens,
+              temperature: config.temperature,
+              top_p: config.top_p,
+              top_k: config.top_k,
+              frequency_penalty: config.frequency_penalty,
+              presence_penalty: config.presence_penalty,
+            },
+          });
+          // Update local state with AI response
+          setMessages((prev) => [
+            ...prev,
+            {
+              id: Date.now().toString(),
+              role: "assistant",
+              parts: [{ type: "text", text: response }],
+            },
+          ]);
+          return new Response(response);
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        } catch (error: any) {
+          // HACK: strip the error message to clean up the response
+          const strippedError = error.message
+            .split("failed with exception ")
+            .pop();
+          return new Response(strippedError, { status: 400 });
+        }
+      },
+    }),
+    messages: initialMessages,
+    onFinish: (message) => {
       setFiles(undefined);
 
       if (fileInputRef.current) {
         fileInputRef.current.value = "";
       }
       Logger.debug("Finished streaming message:", message);
-      Logger.debug("Token usage:", usage);
-      Logger.debug("Finish reason:", finishReason);
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);
     },
-    onResponse: (response) => {
-      Logger.debug("Received HTTP response from server:", response);
-    },
   });
 
   const isLoading = status === "submitted" || status === "streaming";
@@ -157,12 +165,12 @@ export const Chatbot: React.FC<Props> = (props) => {
     }
   };
 
-  const renderAttachment = (attachment: ChatAttachment) => {
-    if (attachment.contentType?.startsWith("image")) {
+  const renderAttachment = (attachment: FileUIPart) => {
+    if (attachment.mediaType?.startsWith("image")) {
       return (
         <img
           src={attachment.url}
-          alt={attachment.name || "Attachment"}
+          alt={attachment.filename || "Attachment"}
           className="object-contain rounded-sm"
           width={100}
           height={100}
@@ -177,18 +185,20 @@ export const Chatbot: React.FC<Props> = (props) => {
         rel="noopener noreferrer"
         className="text-background hover:underline"
       >
-        {attachment.name || "Attachment"}
+        {attachment.filename || "Attachment"}
       </a>
     );
   };
 
-  const renderMessage = (message: Message) => {
+  const renderMessage = (message: UIMessage) => {
+    const textParts = message.parts?.filter((p) => p.type === "text");
+    const textContent = textParts?.map((p) => p.text).join("\n");
     const content =
       message.role === "assistant"
-        ? renderHTML({ html: message.content })
-        : message.content;
+        ? renderHTML({ html: textContent })
+        : textContent;
 
-    const attachments = message.experimental_attachments;
+    const attachments = message.parts?.filter((p) => p.type === "file");
 
     return (
       <>
@@ -204,7 +214,7 @@ export const Chatbot: React.FC<Props> = (props) => {
                     size: "icon",
                   })}
                   href={attachment.url}
-                  download={attachment.name}
+                  download={attachment.filename}
                 >
                   <DownloadIcon className="size-3" />
                 </a>
@@ -274,50 +284,53 @@ export const Chatbot: React.FC<Props> = (props) => {
             </p>
           </div>
         )}
-        {messages.map((message) => (
-          <div
-            key={message.id}
-            className={cn(
-              "flex flex-col group gap-2",
-              message.role === "user" ? "items-end" : "items-start",
-            )}
-          >
+        {messages.map((message) => {
+          const textContent = message.parts
+            ?.filter((p) => p.type === "text")
+            .map((p) => p.text)
+            .join("\n");
+
+          return (
             <div
-              className={`max-w-[80%] p-3 rounded-lg ${
-                message.role === "user"
-                  ? "bg-(--sky-11) text-(--slate-1)"
-                  : "bg-(--slate-4) text-(--slate-12)"
-              }`}
+              key={message.id}
+              className={cn(
+                "flex flex-col group gap-2",
+                message.role === "user" ? "items-end" : "items-start",
+              )}
             >
-              <p
-                className={cn(message.role === "user" && "whitespace-pre-wrap")}
+              <div
+                className={`max-w-[80%] p-3 rounded-lg ${
+                  message.role === "user"
+                    ? "bg-(--sky-11) text-(--slate-1) whitespace-pre-wrap"
+                    : "bg-(--slate-4) text-(--slate-12)"
+                }`}
               >
                 {renderMessage(message)}
-              </p>
-            </div>
-            <div className="flex justify-end text-xs gap-2 invisible group-hover:visible">
-              <button
-                type="button"
-                onClick={async () => {
-                  await copyToClipboard(message.content);
-                  toast({
-                    title: "Copied to clipboard",
-                  });
-                }}
-                className="text-xs text-(--slate-9) hover:text-(--slate-11)"
-              >
-                <ClipboardIcon className="h-3 w-3" />
-              </button>
-              <button
-                type="button"
-                onClick={() => handleDelete(message.id)}
-                className="text-xs text-(--slate-9) hover:text-(--slate-11)"
-              >
-                <Trash2Icon className="h-3 w-3 text-(--red-9)" />
-              </button>
+              </div>
+              <div className="flex justify-end text-xs gap-2 invisible group-hover:visible">
+                <button
+                  type="button"
+                  onClick={async () => {
+                    await copyToClipboard(textContent);
+                    toast({
+                      title: "Copied to clipboard",
+                    });
+                  }}
+                  className="text-xs text-(--slate-9) hover:text-(--slate-11)"
+                >
+                  <ClipboardIcon className="h-3 w-3" />
+                </button>
+                <button
+                  type="button"
+                  onClick={() => handleDelete(message.id)}
+                  className="text-xs text-(--slate-9) hover:text-(--slate-11)"
+                >
+                  <Trash2Icon className="h-3 w-3 text-(--red-9)" />
+                </button>
+              </div>
             </div>
-          </div>
-        ))}
+          );
+        })}
 
         {isLoading && (
           <div className="flex items-center justify-center space-x-2 mb-4">
@@ -336,17 +349,23 @@ export const Chatbot: React.FC<Props> = (props) => {
         {error && (
           <div className="flex items-center justify-center space-x-2 mb-4">
             <ErrorBanner error={error} />
-            <Button variant="outline" size="sm" onClick={() => reload()}>
+            <Button variant="outline" size="sm" onClick={() => regenerate()}>
               Retry
             </Button>
           </div>
         )}
       </div>
-
       <form
-        onSubmit={(evt) => {
-          handleSubmit(evt, {
-            experimental_attachments: files,
+        onSubmit={async (evt) => {
+          evt.preventDefault();
+
+          const fileParts = files
+            ? await convertToFileUIPart(files)
+            : undefined;
+
+          sendMessage({
+            role: "user",
+            parts: [{ type: "text", text: input }, ...(fileParts ?? [])],
           });
         }}
         ref={formRef}
@@ -388,7 +407,7 @@ export const Chatbot: React.FC<Props> = (props) => {
         {files && files.length === 1 && (
           <span
             title={files[0].name}
-            className="text-sm text-(--slate-11) truncate shrink-0 w-24"
+            className="text-sm text-(--slate-11) truncate shrink-0 w-fit max-w-24"
           >
             {files[0].name}
           </span>
@@ -439,7 +458,7 @@ export const Chatbot: React.FC<Props> = (props) => {
               }
               onChange={(event) => {
                 if (event.target.files) {
-                  setFiles(event.target.files);
+                  setFiles([...event.target.files]);
                 }
               }}
             />

package/src/plugins/impl/chat/types.ts

@@ -1,16 +1,13 @@
 /* Copyright 2024 Marimo. All rights reserved. */
+
+import type { UIMessage } from "ai";
+
 export type ChatRole = "system" | "user" | "assistant";
 
 export interface ChatMessage {
   role: ChatRole;
-  content: string;
-  attachments?: ChatAttachment[];
-}
-
-export interface ChatAttachment {
-  name?: string;
-  contentType?: string;
-  url: string;
+  content: string; // TODO: Deprecate content
+  parts: UIMessage["parts"];
 }
 
 export interface SendMessageRequest {

package/src/plugins/impl/code/__tests__/language.test.ts

@@ -0,0 +1,15 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { langs } from "@uiw/codemirror-extensions-langs";
+import { describe, expect, it } from "vitest";
+import { LANGUAGE_MAP } from "../any-language-editor";
+
+describe("Codemirror Languages", () => {
+  const codemirrorLanguages = Object.keys(langs);
+
+  it("LANGUAGE_MAP should have all the languages in CodeMirror", () => {
+    for (const language of Object.values(LANGUAGE_MAP)) {
+      expect(codemirrorLanguages).toContain(language);
+    }
+  });
+});

package/src/plugins/impl/code/any-language-editor.tsx

@@ -1,10 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import {
-  type LanguageName,
-  langs,
-  loadLanguage,
-} from "@uiw/codemirror-extensions-langs";
+import { langs, loadLanguage } from "@uiw/codemirror-extensions-langs";
 import ReactCodeMirror, {
   type Extension,
   type ReactCodeMirrorProps,
@@ -15,6 +11,12 @@ import type { ResolvedTheme } from "@/theme/useTheme";
 import { Logger } from "@/utils/Logger";
 import { ErrorBanner } from "../common/error-banner";
 
+export const LANGUAGE_MAP: Record<string, string> = {
+  python: "py",
+  javascript: "js",
+  typescript: "ts",
+};
+
 /**
  * A code editor that supports any language.
  *
@@ -28,6 +30,9 @@ const AnyLanguageCodeMirror: React.FC<
     showCopyButton?: boolean;
   }
 > = ({ language, showCopyButton, extensions = [], ...props }) => {
+  // Maybe normalize the language to the extension
+  language = LANGUAGE_MAP[language || ""] || language;
+
   const isNotSupported = language && !(language in langs);
   if (isNotSupported) {
     Logger.warn(`Language ${language} not found in CodeMirror.`);
@@ -37,9 +42,7 @@ const AnyLanguageCodeMirror: React.FC<
     if (!language) {
       return extensions;
     }
-    return [loadLanguage(language as LanguageName), ...extensions].filter(
-      Boolean,
-    );
+    return [loadLanguage(language), ...extensions].filter(Boolean);
   }, [language, extensions]);
 
   return (

package/src/plugins/impl/vega/vega.css

@@ -29,3 +29,124 @@
     right: 2px;
   }
 }
+
+/**
+ * Due to a bug when defining our own VegaPlugin where these styles are not applied,
+ * we provide the default styles for vega-actions.
+ * https://github.com/vega/vega-embed/blob/8f4ec8e136951a7e34b358e4e9b6037582a6ccdd/vega-embed.scss
+ */
+
+.vega-embed.has-actions {
+  padding-right: 38px;
+}
+.vega-embed details:not([open]) > :not(summary) {
+  display: none !important;
+}
+.vega-embed summary {
+  list-style: none;
+  position: absolute;
+  top: 0;
+  right: 0;
+  padding: 6px;
+  z-index: 1000;
+  background: white;
+  box-shadow: 1px 1px 3px rgba(0, 0, 0, 0.1);
+  color: #1b1e23;
+  border: 1px solid #aaa;
+  border-radius: 999px;
+  opacity: 0.2;
+  transition: opacity 0.4s ease-in;
+  cursor: pointer;
+  line-height: 0px;
+}
+.vega-embed summary::-webkit-details-marker {
+  display: none;
+}
+.vega-embed summary:active {
+  box-shadow: #aaa 0px 0px 0px 1px inset;
+}
+.vega-embed summary svg {
+  width: 14px;
+  height: 14px;
+}
+.vega-embed details[open] summary {
+  opacity: 0.7;
+}
+.vega-embed:hover summary,
+.vega-embed:focus-within summary {
+  opacity: 1 !important;
+  transition: opacity 0.2s ease;
+}
+.vega-embed .vega-actions {
+  position: absolute;
+  z-index: 1001;
+  top: 35px;
+  /* right: -9px; Remove right for our own styles */
+  display: flex;
+  flex-direction: column;
+  padding-bottom: 8px;
+  padding-top: 8px;
+  border-radius: 4px;
+  box-shadow: 0 2px 8px 0 rgba(0, 0, 0, 0.2);
+  border: 1px solid #d9d9d9;
+  background: white;
+  animation-duration: 0.15s;
+  animation-name: scale-in;
+  animation-timing-function: cubic-bezier(0.2, 0, 0.13, 1.5);
+  text-align: left;
+}
+.vega-embed .vega-actions a {
+  padding: 8px 16px;
+  font-family: sans-serif;
+  font-size: 14px;
+  font-weight: 600;
+  white-space: nowrap;
+  color: #434a56;
+  text-decoration: none;
+}
+.vega-embed .vega-actions a:hover,
+.vega-embed .vega-actions a:focus {
+  background-color: #f7f7f9;
+  color: black;
+}
+.vega-embed .vega-actions::before,
+.vega-embed .vega-actions::after {
+  content: "";
+  display: inline-block;
+  position: absolute;
+}
+.vega-embed .vega-actions::before {
+  left: auto;
+  /* right: 14px; */
+  top: -16px;
+  border: 8px solid #000 0;
+  border-bottom-color: #d9d9d9;
+}
+.vega-embed .vega-actions::after {
+  left: auto;
+  /* right: 15px; */
+  top: -14px;
+  border: 7px solid #000 0;
+  border-bottom-color: #fff;
+}
+.vega-embed .chart-wrapper.fit-x {
+  width: 100%;
+}
+.vega-embed .chart-wrapper.fit-y {
+  height: 100%;
+}
+.vega-embed-wrapper {
+  max-width: 100%;
+  overflow: auto;
+  padding-right: 14px;
+}
+@keyframes scale-in {
+  from {
+    opacity: 0;
+    transform: scale(0.6);
+  }
+  to {
+    opacity: 1;
+    transform: scale(1);
+  }
+}

package/src/plugins/layout/MimeRenderPlugin.tsx

@@ -19,12 +19,9 @@ export class MimeRendererPlugin implements IStatelessPlugin<Data> {
 
   validator = z.object({
     mime: z.string().transform((val) => val as OutputMessage["mimetype"]),
-    data: z.union([
-      z.string(),
-      z.null(),
-      z.record(z.unknown()),
-      z.array(z.any()),
-    ]),
+    data: z
+      .union([z.string(), z.null(), z.record(z.unknown()), z.array(z.any())])
+      .transform((val) => val as OutputMessage["data"]),
   });
 
   render({ data }: IStatelessPluginProps<Data>): JSX.Element {

package/src/stories/cell.stories.tsx

@@ -1,4 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
+
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { createStore, Provider } from "jotai";
 import { createRef } from "react";
@@ -9,6 +10,8 @@ import {
 } from "@/core/cells/types";
 import { defaultUserConfig } from "@/core/config/config-schema";
 import { connectionAtom } from "@/core/network/connection";
+import { requestClientAtom } from "@/core/network/requests";
+import { resolveRequestClient } from "@/core/network/resolve.ts";
 import type { CellConfig } from "@/core/network/types";
 import { WebSocketState } from "@/core/websocket/types";
 import { MultiColumn } from "@/utils/id-tree";
@@ -78,6 +81,7 @@ const Cell: React.FC<{
   const store = createStore();
   store.set(notebookAtom, notebook);
   store.set(connectionAtom, { state: WebSocketState.OPEN });
+  store.set(requestClientAtom, resolveRequestClient());
   return (
     <Provider store={store}>
       <TooltipProvider>
@@ -219,6 +223,7 @@ export const Disabled: Story = {
           config: {
             disabled: true,
             hide_code: false,
+            column: null,
           },
           output: {
             channel: "output",
@@ -299,6 +304,7 @@ export const DisabledAndStaleStatus: Story = {
         overrides={{
           runElapsedTimeMs: 20 as Milliseconds,
           config: {
+            column: null,
             disabled: true,
             hide_code: false,
           },