@marimo-team/frontend 0.15.4 → 0.16.0-dev96986

package/src/components/chat/chat-panel.tsx

@@ -18,6 +18,7 @@ import {
   SquareIcon,
 } from "lucide-react";
 import { memo, useEffect, useMemo, useRef, useState } from "react";
+import { useLocale } from "react-aria";
 import useEvent from "react-use-event-hook";
 import { Button } from "@/components/ui/button";
 import {
@@ -48,7 +49,6 @@ import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
 import { useRequestClient } from "@/core/network/requests";
 import { useRuntimeManager } from "@/core/runtime/config";
-import type { ChatMessage } from "@/plugins/impl/chat/types";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { timeAgo } from "@/utils/dates";
@@ -59,7 +59,6 @@ import { PromptInput } from "../editor/ai/add-cell-with-ai";
 import {
   addContextCompletion,
   CONTEXT_TRIGGER,
-  getAICompletionBodyWithAttachments,
 } from "../editor/ai/completion-utils";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
@@ -68,8 +67,11 @@ import { Tooltip, TooltipProvider } from "../ui/tooltip";
 import { toast } from "../ui/use-toast";
 import { AttachmentRenderer, FileAttachmentPill } from "./chat-components";
 import {
+  buildCompletionRequestBody,
   convertToFileUIPart,
   generateChatTitle,
+  handleToolCall,
+  hasPendingToolCalls,
   isLastMessageReasoning,
 } from "./chat-utils";
 import { MarkdownRenderer } from "./markdown-renderer";
@@ -102,6 +104,7 @@ const ChatHeader: React.FC<ChatHeaderProps> = ({
 }) => {
   const { handleClick } = useOpenSettingsToTab();
   const chatState = useAtomValue(chatStateAtom);
+  const { locale } = useLocale();
   const chats = useMemo(() => {
     return [...chatState.chats.values()].sort(
       (a, b) => b.updatedAt - a.updatedAt,
@@ -158,7 +161,7 @@ const ChatHeader: React.FC<ChatHeaderProps> = ({
                   >
                     <div className="font-medium">{chat.title}</div>
                     <div className="text-sm text-muted-foreground">
-                      {timeAgo(chat.updatedAt)}
+                      {timeAgo(chat.updatedAt, locale)}
                     </div>
                   </button>
                 ))}
@@ -233,7 +236,9 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                   index={i}
                   toolName={part.type}
                   result={part.output}
+                  className="my-2"
                   state={part.state}
+                  input={part.input}
                 />
               );
             }
@@ -265,6 +270,8 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                     toolName={part.type}
                     result={part.output}
                     state={part.state}
+                    input={part.input}
+                    className="my-2"
                   />
                 );
 
@@ -282,10 +289,11 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
                 Logger.error("Unhandled part type:", part.type);
                 try {
                   return (
-                    <MarkdownRenderer
-                      key={i}
-                      content={JSON.stringify(part, null, 2)}
-                    />
+                    <div className="text-xs text-muted-foreground" key={i}>
+                      <MarkdownRenderer
+                        content={JSON.stringify(part, null, 2)}
+                      />
+                    </div>
                   );
                 } catch (error) {
                   Logger.error("Error rendering part:", part.type, error);
@@ -429,7 +437,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
               </>
             )}
 
-            <Tooltip content="Submit">
+            <Tooltip content={isLoading ? "Stop" : "Submit"}>
               <Button
                 variant="text"
                 size="sm"
@@ -565,80 +573,20 @@ const ChatPanelBody = () => {
     id: chatId,
   } = useChat({
     id: activeChatId,
-    // Only automatically submit if we have tool calls but no text response yet
-    sendAutomaticallyWhen: ({ messages }) => {
-      if (messages.length === 0) {
-        return false;
-      }
-
-      const lastMessage = messages[messages.length - 1];
-      const parts = lastMessage.parts;
-
-      if (parts.length === 0) {
-        return false;
-      }
-
-      // Only auto-send if the last message is an assistant message
-      // Because assistant messages are the ones that can have tool calls
-      if (lastMessage.role !== "assistant") {
-        return false;
-      }
-
-      const toolParts = parts.filter((part) =>
-        part.type.startsWith("tool-"),
-      ) as ToolUIPart[];
-
-      const hasCompletedToolCalls = toolParts.some(
-        (part) => part.state === "output-available",
-      );
-
-      // Check if the last part has any text content
-      const lastPart = parts[parts.length - 1];
-      const hasTextContent =
-        lastPart.type === "text" && lastPart.text?.trim().length > 0;
-
-      // Only auto-send if we have completed tool calls and there is no reply yet
-      return hasCompletedToolCalls && !hasTextContent;
-    },
+    sendAutomaticallyWhen: ({ messages }) => hasPendingToolCalls(messages),
     messages: activeChat?.messages || [], // initial messages
     transport: new DefaultChatTransport({
       api: runtimeManager.getAiURL("chat").toString(),
       headers: runtimeManager.headers(),
       prepareSendMessagesRequest: async (options) => {
-        // Map from parts to a single string
-        function toContent(parts: UIMessage["parts"]): string {
-          return parts
-            .map((part) => (part.type === "text" ? part.text : ""))
-            .join("\n");
-        }
-
-        const input = toContent(options.messages.flatMap((m) => m.parts));
-        const completionBody = await getAICompletionBodyWithAttachments({
-          input,
-        });
-
-        // Map from UIMessage to our ChatMessage type
-        // If it's the last message, add the attachments from the completion body
-        function mapMessage(m: UIMessage, isLastMessage: boolean): ChatMessage {
-          const parts = m.parts;
-          if (isLastMessage) {
-            parts.push(...completionBody.attachments);
-          }
-          return {
-            role: m.role,
-            content: toContent(m.parts),
-            parts: parts,
-          };
-        }
+        const completionBody = await buildCompletionRequestBody(
+          options.messages,
+        );
 
         return {
           body: {
             ...options,
-            ...completionBody.body,
-            messages: options.messages.map((m, idx) => ({
-              ...m,
-              ...mapMessage(m, idx === options.messages.length - 1),
-            })),
+            ...completionBody,
           },
         };
       },
@@ -653,24 +601,15 @@ const ChatPanelBody = () => {
       });
     },
     onToolCall: async ({ toolCall }) => {
-      try {
-        const response = await invokeAiTool({
+      await handleToolCall({
+        invokeAiTool,
+        addToolResult,
+        toolCall: {
           toolName: toolCall.toolName,
-          arguments: toolCall.input as Record<string, never>,
-        });
-        addToolResult({
-          tool: toolCall.toolName,
           toolCallId: toolCall.toolCallId,
-          output: response.result || response.error,
-        });
-      } catch (error) {
-        Logger.error("Tool call failed:", error);
-        addToolResult({
-          tool: toolCall.toolName,
-          toolCallId: toolCall.toolCallId,
-          output: `Error: ${error instanceof Error ? error.message : String(error)}`,
-        });
-      }
+          input: toolCall.input as Record<string, never>,
+        },
+      });
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);

package/src/components/chat/chat-utils.ts

@@ -1,7 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { FileUIPart, UIMessage } from "ai";
+import type { components } from "@marimo-team/marimo-api";
+import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import type {
+  InvokeAiToolRequest,
+  InvokeAiToolResponse,
+} from "@/core/network/types";
+import type { ChatMessage } from "@/plugins/impl/chat/types";
 import { blobToString } from "@/utils/fileToBase64";
+import { Logger } from "@/utils/Logger";
+import { getAICompletionBodyWithAttachments } from "../editor/ai/completion-utils";
 
 export function generateChatTitle(message: string): string {
   return message.length > 50 ? `${message.slice(0, 50)}...` : message;
@@ -48,3 +56,121 @@ export function isLastMessageReasoning(messages: UIMessage[]): boolean {
   const lastPart = parts[parts.length - 1];
   return lastPart.type === "reasoning";
 }
+
+function stringifyTextParts(parts: UIMessage["parts"]): string {
+  return parts
+    .map((part) => (part.type === "text" ? part.text : ""))
+    .join("\n");
+}
+
+export async function buildCompletionRequestBody(
+  messages: UIMessage[],
+): Promise<{
+  messages: ChatMessage[];
+  context?: (null | components["schemas"]["AiCompletionContext"]) | undefined;
+  includeOtherCode: string;
+  selectedText?: string | null | undefined;
+}> {
+  const input = stringifyTextParts(messages.flatMap((m) => m.parts));
+  const completionBody = await getAICompletionBodyWithAttachments({ input });
+
+  // Map from UIMessage to our ChatMessage type
+  // If it's the last message, add the attachments from the completion body
+  function toChatMessage(message: UIMessage, isLast: boolean): ChatMessage {
+    // Clone parts to avoid mutating the original message
+    const parts = [...message.parts];
+    if (isLast) {
+      parts.push(...completionBody.attachments);
+    }
+    return {
+      role: message.role,
+      content: stringifyTextParts(message.parts), // This is no longer used in the backend
+      parts,
+    };
+  }
+
+  return {
+    ...completionBody.body,
+    messages: messages.map((m, idx) =>
+      toChatMessage(m, idx === messages.length - 1),
+    ),
+  };
+}
+
+interface AddToolResult {
+  tool: string;
+  toolCallId: string;
+  output: unknown;
+}
+
+export async function handleToolCall({
+  invokeAiTool,
+  addToolResult,
+  toolCall,
+}: {
+  invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
+  addToolResult: (result: AddToolResult) => Promise<void>;
+  toolCall: {
+    toolName: string;
+    toolCallId: string;
+    input: Record<string, never>;
+  };
+}) {
+  try {
+    const response = await invokeAiTool({
+      toolName: toolCall.toolName,
+      arguments: toolCall.input,
+    });
+    addToolResult({
+      tool: toolCall.toolName,
+      toolCallId: toolCall.toolCallId,
+      output: response.result || response.error,
+    });
+  } catch (error) {
+    Logger.error("Tool call failed:", error);
+    addToolResult({
+      tool: toolCall.toolName,
+      toolCallId: toolCall.toolCallId,
+      output: `Error: ${error instanceof Error ? error.message : String(error)}`,
+    });
+  }
+}
+
+/**
+ * Checks if we should send a message automatically based on the messages.
+ * We only want to send a message if we have completed tool calls and there is no reply yet.
+ */
+export function hasPendingToolCalls(messages: UIMessage[]): boolean {
+  if (messages.length === 0) {
+    return false;
+  }
+
+  const lastMessage = messages[messages.length - 1];
+  const parts = lastMessage.parts;
+
+  if (parts.length === 0) {
+    return false;
+  }
+
+  // Only auto-send if the last message is an assistant message
+  // Because assistant messages are the ones that can have tool calls
+  if (lastMessage.role !== "assistant") {
+    return false;
+  }
+
+  const toolParts = parts.filter((part) =>
+    part.type.startsWith("tool-"),
+  ) as ToolUIPart[];
+
+  const hasCompletedToolCalls = toolParts.some(
+    (part) => part.state === "output-available",
+  );
+
+  // Check if the last part has any text content
+  const lastPart = parts[parts.length - 1];
+  const hasTextContent =
+    lastPart.type === "text" && lastPart.text?.trim().length > 0;
+
+  // Only auto-send if we have completed tool calls and there is no reply yet
+  return hasCompletedToolCalls && !hasTextContent;
+}

package/src/components/chat/markdown-renderer.css

@@ -1,4 +1,10 @@
+@source "../../../node_modules/streamdown/dist/index.js";
+
 .mo-markdown-renderer {
+  /* Make headers a bit smaller */
+  --text-2xl: 1rem;
+  --text-3xl: 1.2rem;
+
   h1,
   h2,
   h3,
@@ -12,4 +18,37 @@
     margin-top: 8px;
     margin-bottom: 12px;
   }
+
+  pre {
+    width: 100%;
+  }
+
+  [data-code-block-container="true"] {
+    margin-top: 0 !important;
+  }
+
+  [data-code-block-header="true"] {
+    padding: 4px !important;
+  }
+
+  code {
+    font-size: inherit !important;
+  }
+
+  li.task-list-item {
+    list-style-type: none;
+  }
+
+  li {
+    padding-left: 6px;
+
+    > input {
+      margin-right: 6px;
+    }
+  }
+
+  ul,
+  ol {
+    padding-left: 10px;
+  }
 }

package/src/components/chat/markdown-renderer.tsx

@@ -3,10 +3,8 @@
 import { EditorView } from "@codemirror/view";
 import { useAtomValue } from "jotai";
 import { BetweenHorizontalStartIcon } from "lucide-react";
-import { marked } from "marked";
-import { memo, Suspense, useEffect, useMemo, useState } from "react";
-import Markdown, { type Components } from "react-markdown";
-import remarkGfm from "remark-gfm";
+import { memo, Suspense, useState } from "react";
+import { Streamdown, type StreamdownProps } from "streamdown";
 import { Button, type ButtonProps } from "@/components/ui/button";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
@@ -116,9 +114,9 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
   const { theme } = useTheme();
   const [value, setValue] = useState(code);
 
-  useEffect(() => {
+  if (value !== code) {
     setValue(code);
-  }, [code]);
+  }
 
   const handleCopyCode = async () => {
     await copyToClipboard(value);
@@ -152,18 +150,13 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
 const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
   const [copied, setCopied] = useState(false);
 
-  useEffect(() => {
-    if (copied) {
-      setTimeout(() => setCopied(false), 1000);
-    }
-  }, [copied]);
-
   return (
     <Button
       {...props}
       onClick={(e) => {
         onClick?.(e);
         setCopied(true);
+        setTimeout(() => setCopied(false), 1000);
       }}
     >
       {copied ? "Copied" : "Copy"}
@@ -171,14 +164,17 @@ const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
   );
 };
 
+type Components = StreamdownProps["components"];
+
 const COMPONENTS: Components = {
   code: ({ children, className }) => {
     const language = className?.replace("language-", "");
     if (language && typeof children === "string") {
+      const code = children.trim();
       return (
         <div>
           <div className="text-xs text-muted-foreground pl-1">{language}</div>
-          <CodeBlock code={children.trim()} language={language} />
+          <CodeBlock code={code} language={language} />
         </div>
       );
     }
@@ -186,42 +182,11 @@ const COMPONENTS: Components = {
   },
 };
 
-function parseMarkdownIntoBlocks(markdown: string): string[] {
-  const tokens = marked.lexer(markdown);
-  return tokens.map((token) => token.raw);
-}
-
-const PLUGINS = [remarkGfm];
-
-const MemoizedMarkdownBlock = memo(
-  ({ content }: { content: string }) => {
-    return (
-      <Markdown
-        components={COMPONENTS}
-        remarkPlugins={PLUGINS}
-        className="mo-markdown-renderer prose dark:prose-invert max-w-none prose-pre:pl-0"
-      >
-        {content}
-      </Markdown>
-    );
-  },
-  (prevProps, nextProps) => prevProps.content === nextProps.content,
-);
-
-MemoizedMarkdownBlock.displayName = "MemoizedMarkdownBlock";
-
 export const MarkdownRenderer = memo(({ content }: { content: string }) => {
-  const blocks = useMemo(() => parseMarkdownIntoBlocks(content), [content]);
-
   return (
-    <>
-      {blocks.map((block, index) => (
-        <MemoizedMarkdownBlock
-          content={block}
-          key={`markdown-block-${index}`}
-        />
-      ))}
-    </>
+    <Streamdown components={COMPONENTS} className="mo-markdown-renderer">
+      {content}
+    </Streamdown>
   );
 });
 MarkdownRenderer.displayName = "MarkdownRenderer";