@marimo-team/islands 0.15.1 → 0.15.3

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

@@ -0,0 +1,50 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { FileUIPart, UIMessage } from "ai";
+import { blobToString } from "@/utils/fileToBase64";
+
+export function generateChatTitle(message: string): string {
+  return message.length > 50 ? `${message.slice(0, 50)}...` : message;
+}
+
+export async function convertToFileUIPart(
+  files: File[],
+): Promise<FileUIPart[]> {
+  const fileUIParts = await Promise.all(
+    files.map(async (file) => {
+      const part: FileUIPart = {
+        type: "file" as const,
+        mediaType: file.type,
+        filename: file.name,
+        url: await blobToString(file, "dataUrl"),
+      };
+      return part;
+    }),
+  );
+
+  return fileUIParts;
+}
+
+export function isLastMessageReasoning(messages: UIMessage[]): boolean {
+  if (messages.length === 0) {
+    return false;
+  }
+
+  const lastMessage = messages.at(-1);
+  if (!lastMessage) {
+    return false;
+  }
+
+  if (lastMessage.role !== "assistant" || !lastMessage.parts) {
+    return false;
+  }
+
+  const parts = lastMessage.parts;
+  if (parts.length === 0) {
+    return false;
+  }
+
+  // Check if the last part is reasoning
+  const lastPart = parts[parts.length - 1];
+  return lastPart.type === "reasoning";
+}

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

@@ -172,21 +172,17 @@ const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
 };
 
 const COMPONENTS: Components = {
-  code: ({ children, className, key }) => {
+  code: ({ children, className }) => {
     const language = className?.replace("language-", "");
     if (language && typeof children === "string") {
       return (
-        <div key={key}>
+        <div>
           <div className="text-xs text-muted-foreground pl-1">{language}</div>
           <CodeBlock code={children.trim()} language={language} />
         </div>
       );
     }
-    return (
-      <code key={key} className={className}>
-        {children}
-      </code>
-    );
+    return <code className={className}>{children}</code>;
   },
 };
 

package/src/components/chat/tool-call-accordion.tsx

@@ -1,5 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { ToolUIPart } from "ai";
 import {
   CheckCircleIcon,
   Loader2,
@@ -17,10 +18,10 @@ import { cn } from "@/utils/cn";
 
 interface ToolCallAccordionProps {
   toolName: string;
-  result?: string | null;
+  result: unknown;
   error?: string;
   index?: number;
-  state?: "partial-call" | "call" | "result";
+  state?: ToolUIPart["state"];
 }
 
 export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
@@ -30,7 +31,7 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
   index = 0,
   state,
 }) => {
-  const hasResult = state === "result" && (result || error);
+  const hasResult = state === "output-available" && (result || error);
   const status = error ? "error" : hasResult ? "success" : "loading";
 
   const getStatusIcon = () => {
@@ -84,8 +85,7 @@ export const ToolCallAccordion: React.FC<ToolCallAccordionProps> = ({
           {/* Only show content when tool is complete */}
           {hasResult && (
             <div className="space-y-3">
-              {/* Result */}
-              {result && (
+              {result !== undefined && result !== null && (
                 <div>
                   <div className="text-xs font-medium text-muted-foreground mt-2 mb-1">
                     Result:

package/src/components/datasources/__tests__/utils.test.ts

@@ -13,6 +13,12 @@ describe("sqlCode", () => {
     source: "local",
     source_type: "local",
     type: "table",
+    engine: null,
+    indexes: null,
+    num_columns: null,
+    num_rows: null,
+    variable_name: null,
+    primary_keys: null,
   };
   const mockColumn = {
     name: "email" as const,

package/src/components/datasources/column-preview.tsx

@@ -173,9 +173,7 @@ export function renderPreviewError({
 }
 
 export function renderStats(
-  stats: Partial<
-    Record<ColumnHeaderStatsKey, string | number | boolean | null | undefined>
-  >,
+  stats: Partial<Record<ColumnHeaderStatsKey, unknown>>,
   dataType: DataType,
 ) {
   return (

package/src/components/editor/actions/useNotebookActions.tsx

@@ -73,8 +73,8 @@ import { newNotebookURL } from "@/utils/urls";
 import { useRunAllCells } from "../cell/useRunCells";
 import { useChromeActions, useChromeState } from "../chrome/state";
 import { PANELS } from "../chrome/types";
-import { commandPaletteAtom } from "../controls/command-palette";
 import { keyboardShortcutsAtom } from "../controls/keyboard-shortcuts";
+import { commandPaletteAtom } from "../controls/state";
 import { AddDatabaseDialogContent } from "../database/add-database-form";
 import { displayLayoutName, getLayoutIcon } from "../renderers/layout-select";
 import { LAYOUT_TYPES } from "../renderers/types";

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -1,5 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { useCompletion } from "@ai-sdk/react";
 import {
   autocompletion,
   type Completion,
@@ -16,8 +17,7 @@ import ReactCodeMirror, {
   minimalSetup,
   type ReactCodeMirrorRef,
 } from "@uiw/react-codemirror";
-import { useCompletion } from "ai/react";
-import { useAtom, useAtomValue, useStore } from "jotai";
+import { useAtom, useStore } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import {
   ChevronsUpDown,
@@ -40,12 +40,9 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { toast } from "@/components/ui/use-toast";
-import { useModelChange } from "@/core/ai/config";
 import { resourceExtension } from "@/core/codemirror/ai/resources";
 import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";
 import { SQLLanguageAdapter } from "@/core/codemirror/language/languages/sql/sql";
-import { aiAtom } from "@/core/config/config";
-import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
@@ -58,6 +55,7 @@ import {
   createAiCompletionOnKeydown,
 } from "./completion-handlers";
 import {
+  CONTEXT_TRIGGER,
   getAICompletionBody,
   mentionsCompletionSource,
 } from "./completion-utils";
@@ -90,9 +88,6 @@ export const AddCellWithAI: React.FC<{
   const { theme } = useTheme();
   const runtimeManager = useRuntimeManager();
 
-  const ai = useAtomValue(aiAtom);
-  const editModel = ai?.models?.edit_model || DEFAULT_AI_MODEL;
-  const { saveModelChange } = useModelChange();
   const inputRef = useRef<ReactCodeMirrorRef>(null);
 
   const {
@@ -261,10 +256,6 @@ export const AddCellWithAI: React.FC<{
         <div className="ml-auto flex items-center gap-1">
           {languageDropdown}
           <AIModelDropdown
-            value={editModel}
-            onSelect={(model) => {
-              saveModelChange(model, "edit");
-            }}
             triggerClassName="h-7 text-xs max-w-64"
             iconSize="small"
             forRole="edit"
@@ -301,6 +292,7 @@ interface PromptInputProps {
   onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
   additionalCompletions?: AdditionalCompletions;
   maxHeight?: string;
+  onAddFiles?: (files: File[]) => void;
 }
 
 /**
@@ -318,6 +310,7 @@ export const PromptInput = ({
   onSubmit,
   onKeyDown,
   onClose,
+  onAddFiles,
   additionalCompletions,
   maxHeight,
 }: PromptInputProps) => {
@@ -346,7 +339,11 @@ export const PromptInput = ({
     return [
       autocompletion({}),
       markdownLanguage,
-      resourceExtension(markdownLanguage.language, store),
+      resourceExtension({
+        language: markdownLanguage.language,
+        store,
+        onAddFiles,
+      }),
       markdownLanguage.language.data.of({
         autocomplete: additionalCompletionsSource,
       }),
@@ -411,7 +408,13 @@ export const PromptInput = ({
         },
       ]),
     ];
-  }, [store, additionalCompletionsSource, handleSubmit, handleEscape]);
+  }, [
+    store,
+    onAddFiles,
+    additionalCompletionsSource,
+    handleSubmit,
+    handleEscape,
+  ]);
 
   return (
     <ReactCodeMirror
@@ -425,7 +428,9 @@ export const PromptInput = ({
       onChange={onChange}
       onKeyDown={onKeyDown}
       theme={theme === "dark" ? "dark" : "light"}
-      placeholder={placeholder || "Generate with AI"}
+      placeholder={
+        placeholder || `Generate with AI, ${CONTEXT_TRIGGER} to include context`
+      }
     />
   );
 };

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -2,7 +2,7 @@
 
 import { useCompletion } from "@ai-sdk/react";
 import { EditorView } from "@codemirror/view";
-import { Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
+import { AtSignIcon, Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
 import React, { useEffect, useId, useState } from "react";
 import CodeMirrorMerge from "react-codemirror-merge";
 import { Button } from "@/components/ui/button";
@@ -12,6 +12,7 @@ import "./merge-editor.css";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { useAtom } from "jotai";
+import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
 import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
 import { Tooltip } from "@/components/ui/tooltip";
@@ -30,7 +31,7 @@ import {
   CompletionActions,
   createAiCompletionOnKeydown,
 } from "./completion-handlers";
-import { getAICompletionBody } from "./completion-utils";
+import { addContextCompletion, getAICompletionBody } from "./completion-utils";
 
 const Original = CodeMirrorMerge.Original;
 const Modified = CodeMirrorMerge.Modified;
@@ -164,7 +165,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
             <SparklesIcon className="text-(--blue-10) shrink-0" size={16} />
             <PromptInput
               inputRef={inputRef}
-              className="h-full my-0 py-1.5"
+              className="h-full my-0 py-2"
               onClose={() => {
                 declineChange();
                 setCompletion("");
@@ -209,6 +210,24 @@ export const AiCompletionEditor: React.FC<Props> = ({
                 size="xs"
               />
             )}
+            <div className="flex flex-row items-center gap-0.5 -ml-1.5 -mr-2">
+              <Tooltip content="Add context">
+                <Button
+                  variant="text"
+                  size="icon"
+                  onClick={() => addContextCompletion(inputRef)}
+                >
+                  <AtSignIcon className="h-3 w-3" />
+                </Button>
+              </Tooltip>
+              <AIModelDropdown
+                triggerClassName="h-7"
+                iconSize="small"
+                displayIconOnly={true}
+                forRole="edit"
+              />
+            </div>
+
             <div className="h-full w-px bg-border mx-2" />
             <Tooltip content="Include code from other cells">
               <div className="flex flex-row items-start gap-1 overflow-hidden">

package/src/components/editor/ai/completion-handlers.tsx

@@ -25,10 +25,8 @@ export const createAiCompletionOnKeydown = (opts: {
     const metaKey = isPlatformMac() ? e.metaKey : e.ctrlKey;
 
     // Mod+Enter should accept the completion, if there is one
-    if (metaKey && e.key === "Enter") {
-      if (!isLoading && completion) {
-        handleAcceptCompletion();
-      }
+    if (metaKey && e.key === "Enter" && !isLoading && completion) {
+      handleAcceptCompletion();
     }
 
     // Mod+Shift+Delete should decline the completion

package/src/components/editor/ai/completion-utils.ts

@@ -1,28 +1,41 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type {
-  Completion,
-  CompletionContext,
-  CompletionSource,
+import {
+  type Completion,
+  type CompletionContext,
+  type CompletionSource,
+  startCompletion,
 } from "@codemirror/autocomplete";
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
+import type { FileUIPart } from "ai";
 import { getAIContextRegistry } from "@/core/ai/context/context";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
 
+export const CONTEXT_TRIGGER = "@";
+
+interface Opts {
+  input: string;
+}
+
+interface AICompletionBodyWithAttachments {
+  body: Omit<AiCompletionRequest, "language" | "prompt" | "code">;
+  attachments: FileUIPart[];
+}
+
 /**
  * Gets the request body for the AI completion API.
  */
 export function getAICompletionBody({
   input,
-}: {
-  input: string;
-}): Omit<AiCompletionRequest, "language" | "prompt" | "code"> {
+}: Opts): Omit<AiCompletionRequest, "language" | "prompt" | "code"> {
   let contextString = "";
+
   // Skip if no '@' in the input
   if (input.includes("@")) {
-    contextString = extractDatasetsAndVariables(input);
+    contextString = extractTaggedContext(input);
     Logger.debug("Included context", contextString);
   }
 
@@ -37,11 +50,49 @@ export function getAICompletionBody({
 }
 
 /**
- * Extracts datasets and variables from the input.
+ * Gets the request body and attachments for the AI completion API.
+ */
+export async function getAICompletionBodyWithAttachments({
+  input,
+}: Opts): Promise<AICompletionBodyWithAttachments> {
+  let contextString = "";
+  let attachments: FileUIPart[] = [];
+
+  // Skip if no '@' in the input
+  if (input.includes("@")) {
+    const registry = getAIContextRegistry(store);
+    const contextIds = registry.parseAllContextIds(input);
+
+    // Get context string
+    contextString = registry.formatContextForAI(contextIds);
+
+    // Get attachments
+    try {
+      attachments = await registry.getAttachmentsForContext(contextIds);
+      Logger.debug("Included attachments", attachments.length);
+    } catch (error) {
+      Logger.error("Error getting attachments:", error);
+    }
+  }
+
+  return {
+    body: {
+      includeOtherCode: getCodes(""),
+      context: {
+        plainText: contextString,
+        schema: [],
+        variables: [],
+      },
+    },
+    attachments,
+  };
+}
+
+/**
+ * Extracts datasets, variables and other context from the input.
  * References are with @<name> in the input.
- * Prioritizes datasets over variables if there's a name conflict.
  */
-function extractDatasetsAndVariables(input: string): string {
+function extractTaggedContext(input: string): string {
   const registry = getAIContextRegistry(store);
   const contextIds = registry.parseAllContextIds(input);
   return registry.formatContextForAI(contextIds);
@@ -71,3 +122,26 @@ export function mentionsCompletionSource(
     };
   };
 }
+
+export function addContextCompletion(
+  inputRef: React.RefObject<ReactCodeMirrorRef | null>,
+) {
+  if (inputRef.current?.view) {
+    const pos = inputRef.current.view.state.selection.main.from;
+    // Insert @ at the cursor position
+    inputRef.current.view.dispatch({
+      changes: {
+        from: pos,
+        to: pos,
+        insert: CONTEXT_TRIGGER,
+      },
+      selection: {
+        anchor: pos + 1,
+        head: pos + 1,
+      },
+    });
+    inputRef.current.view.focus();
+    // Trigger completion
+    startCompletion(inputRef.current.view);
+  }
+}

package/src/components/editor/alerts/startup-logs-alert.tsx

@@ -0,0 +1,72 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { BoxIcon, XIcon } from "lucide-react";
+import { useEffect, useState } from "react";
+import { useEvent } from "react-use-event-hook";
+import { useAlertActions, useAlerts } from "@/core/alerts/state";
+import { Banner } from "@/plugins/impl/common/error-banner";
+import { Button } from "../../ui/button";
+
+export const StartupLogsAlert: React.FC = () => {
+  const { startupLogsAlert } = useAlerts();
+  const { clearStartupLogsAlert } = useAlertActions();
+  const [hasCleared, setHasCleared] = useState(false);
+  const status = startupLogsAlert?.status;
+
+  const handleClear = useEvent(() => {
+    clearStartupLogsAlert();
+    setHasCleared(true);
+  });
+
+  const isDone = status === "done";
+  const isRunning = status === "start" || status === "append";
+
+  useEffect(() => {
+    let timeout: number | undefined;
+
+    if (isDone) {
+      // Dismiss after 5 seconds
+      timeout = window.setTimeout(() => handleClear(), 5000);
+    }
+
+    return () => {
+      if (timeout) {
+        window.clearTimeout(timeout);
+      }
+    };
+  }, [isDone]);
+
+  if (startupLogsAlert === null || hasCleared) {
+    return null;
+  }
+
+  return (
+    <div className="flex flex-col gap-4 mb-5 fixed top-5 left-12 min-w-[400px] max-w-[600px] z-200 opacity-95">
+      <Banner
+        kind="info"
+        className="flex flex-col rounded py-3 px-5 animate-in slide-in-from-left"
+      >
+        <div className="flex justify-between">
+          <span className="font-bold text-lg flex items-center mb-2">
+            <BoxIcon className="w-5 h-5 inline-block mr-2" />
+            {isRunning ? "Initializing..." : "Initialized"}
+          </span>
+          <Button
+            variant="text"
+            data-testid="remove-startup-logs-button"
+            size="icon"
+            onClick={handleClear}
+          >
+            <XIcon className="w-5 h-5" />
+          </Button>
+        </div>
+        <div className="flex flex-col gap-4 justify-between items-start text-muted-foreground text-base">
+          <div className="w-full">
+            <pre className="bg-muted p-3 rounded text-sm font-mono overflow-auto max-h-64 whitespace-pre-wrap">
+              {startupLogsAlert.content || "Starting startup script..."}
+            </pre>
+          </div>
+        </div>
+      </Banner>
+    </div>
+  );
+};

package/src/components/editor/chrome/panels/datasources-panel.tsx

@@ -1,6 +1,8 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { DataSources } from "@/components/datasources/datasources";
 
-export const DataSourcesPanel: React.FC = () => {
+const DataSourcesPanel: React.FC = () => {
   return <DataSources />;
 };
+
+export default DataSourcesPanel;

package/src/components/editor/chrome/panels/dependency-graph-panel.tsx

@@ -6,7 +6,7 @@ import { useVariables } from "@/core/variables/state";
 import { cn } from "@/utils/cn";
 import { DependencyGraph } from "../../../dependency-graph/dependency-graph";
 
-export const DependencyGraphPanel: React.FC = () => {
+const DependencyGraphPanel: React.FC = () => {
   const variables = useVariables();
   const cellIds = useCellIds();
   const [cells] = useCellDataAtoms();
@@ -21,3 +21,5 @@ export const DependencyGraphPanel: React.FC = () => {
     </div>
   );
 };
+
+export default DependencyGraphPanel;

package/src/components/editor/chrome/panels/documentation-panel.tsx

@@ -8,7 +8,7 @@ import "../../documentation.css";
 import { TextSearchIcon } from "lucide-react";
 import { PanelEmptyState } from "./empty-state";
 
-export const DocumentationPanel: React.FC = () => {
+const DocumentationPanel: React.FC = () => {
   const { documentation } = useAtomValue(documentationAtom);
 
   if (!documentation) {
@@ -27,3 +27,5 @@ export const DocumentationPanel: React.FC = () => {
     </div>
   );
 };
+
+export default DocumentationPanel;

package/src/components/editor/chrome/panels/error-panel.tsx

@@ -7,7 +7,7 @@ import { useCellErrors } from "../../../../core/cells/cells";
 import { MarimoErrorOutput } from "../../output/MarimoErrorOutput";
 import { PanelEmptyState } from "./empty-state";
 
-export const ErrorsPanel: React.FC = () => {
+const ErrorsPanel: React.FC = () => {
   const errors = useCellErrors();
 
   if (errors.length === 0) {
@@ -33,3 +33,5 @@ export const ErrorsPanel: React.FC = () => {
     </div>
   );
 };
+
+export default ErrorsPanel;

package/src/components/editor/chrome/panels/file-explorer-panel.tsx

@@ -6,7 +6,7 @@ import { cn } from "@/utils/cn";
 import { FileExplorer } from "../../file-tree/file-explorer";
 import { useFileExplorerUpload } from "../../file-tree/upload";
 
-export const FileExplorerPanel: React.FC = () => {
+const FileExplorerPanel: React.FC = () => {
   const { ref, height = 1 } = useResizeObserver<HTMLDivElement>();
   const { getRootProps, getInputProps, isDragActive } = useFileExplorerUpload({
     noClick: true,
@@ -31,3 +31,5 @@ export const FileExplorerPanel: React.FC = () => {
     </div>
   );
 };
+
+export default FileExplorerPanel;

package/src/components/editor/chrome/panels/logs-panel.tsx

@@ -14,7 +14,7 @@ interface Props {
   logs: CellLog[];
 }
 
-export const LogsPanel: React.FC = () => {
+const LogsPanel: React.FC = () => {
   const logs = useCellLogs();
   const { clearLogs } = useCellActions();
 
@@ -46,6 +46,8 @@ export const LogsPanel: React.FC = () => {
   );
 };
 
+export default LogsPanel;
+
 export const LogViewer: React.FC<Props> = ({ logs, className }) => {
   const hover = "opacity-70 group-hover:bg-(--gray-3) group-hover:opacity-100";
   return (

package/src/components/editor/chrome/panels/outline-panel.tsx

@@ -13,7 +13,7 @@ import {
   useActiveOutline,
 } from "./outline/useActiveOutline";
 
-export const OutlinePanel: React.FC = () => {
+const OutlinePanel: React.FC = () => {
   const { items } = useAtomValue(notebookOutline);
   const headerElements = useMemo(() => findOutlineElements(items), [items]);
   const { activeHeaderId, activeOccurrences } =
@@ -37,3 +37,5 @@ export const OutlinePanel: React.FC = () => {
     />
   );
 };
+
+export default OutlinePanel;

package/src/components/editor/chrome/panels/packages-panel.tsx

@@ -65,7 +65,7 @@ const PackageActionButton: React.FC<{
   );
 };
 
-export const PackagesPanel: React.FC = () => {
+const PackagesPanel: React.FC = () => {
   const [config] = useResolvedMarimoConfig();
   const packageManager = config.package_management.manager;
   const { getDependencyTree, getPackageList } = useRequestClient();
@@ -162,6 +162,8 @@ export const PackagesPanel: React.FC = () => {
   );
 };
 
+export default PackagesPanel;
+
 const InstallPackageForm: React.FC<{
   packageManager: string;
   onSuccess: () => void;
@@ -411,7 +413,7 @@ const RemoveButton: React.FC<{
 };
 
 const DependencyTree: React.FC<{
-  tree?: DependencyTreeNode;
+  tree: DependencyTreeNode | null;
   error?: Error | null;
   onSuccess: () => void;
 }> = ({ tree, error, onSuccess }) => {