@marimo-team/islands 0.15.2 → 0.15.3

package/src/core/ai/context/providers/cell-output.ts

@@ -0,0 +1,349 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import type { FileUIPart } from "ai";
+import { toPng } from "html-to-image";
+import { notebookAtom } from "@/core/cells/cells";
+import { type CellId, CellOutputId } from "@/core/cells/ids";
+import { displayCellName } from "@/core/cells/names";
+import { isOutputEmpty } from "@/core/cells/outputs";
+import type { OutputMessage } from "@/core/kernel/messages";
+import type { JotaiStore } from "@/core/state/jotai";
+import { Logger } from "@/utils/Logger";
+import { type AIContextItem, AIContextProvider } from "../registry";
+import { contextToXml } from "../utils";
+import { Boosts } from "./common";
+
+export interface CellOutputContextItem extends AIContextItem {
+  type: "cell-output";
+  data: {
+    cellId: CellId;
+    cellName: string;
+    cellCode: string;
+    output: OutputMessage;
+    outputType: "text" | "media";
+    processedContent?: string;
+    imageUrl?: string;
+    shouldDownloadImage?: boolean;
+  };
+}
+
+function isMediaMimetype(
+  mimetype: OutputMessage["mimetype"] | undefined,
+  htmlString: string,
+): boolean {
+  if (!mimetype) {
+    return false;
+  }
+
+  const mediaPrefixes = ["image/", "video/", "audio/", "application/pdf"];
+  if (mediaPrefixes.some((prefix) => mimetype.startsWith(prefix))) {
+    return true;
+  }
+  const mediaIncludes = ["svg", "vega"];
+  if (mediaIncludes.some((include) => mimetype.includes(include))) {
+    return true;
+  }
+
+  // If it is HTML, we need to check if it contains a media tag
+  if (mimetype === "text/html") {
+    const mediaTags = [
+      "<img",
+      "<video",
+      "<audio",
+      "<iframe",
+      "<canvas",
+      "<svg",
+      "<marimo-ui-element",
+    ];
+    if (mediaTags.some((tag) => htmlString.includes(tag))) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+function parseHtmlContent(htmlString: string): string {
+  try {
+    // Create a temporary DOM element to parse HTML
+    const tempDiv = document.createElement("div");
+    tempDiv.innerHTML = htmlString;
+
+    // Extract text content, removing HTML tags
+    const textContent = tempDiv.textContent || tempDiv.innerText || "";
+
+    // Clean up extra whitespace
+    return textContent.replaceAll(/\s+/g, " ").trim();
+  } catch (error) {
+    Logger.error("Error parsing HTML content:", error);
+    // If parsing fails, return the original string
+    return htmlString;
+  }
+}
+
+export class CellOutputContextProvider extends AIContextProvider<CellOutputContextItem> {
+  readonly title = "Cell Outputs";
+  readonly mentionPrefix = "@";
+  readonly contextType = "cell-output";
+
+  constructor(private store: JotaiStore) {
+    super();
+  }
+
+  getItems(): CellOutputContextItem[] {
+    const notebook = this.store.get(notebookAtom);
+    const items: CellOutputContextItem[] = [];
+
+    for (const [cellIndex, cellId] of notebook.cellIds.inOrderIds.entries()) {
+      const cellRuntime = notebook.cellRuntime[cellId];
+
+      // Filter to only cells with output
+      if (!cellRuntime?.output || isOutputEmpty(cellRuntime.output)) {
+        continue;
+      }
+
+      const cellData = notebook.cellData[cellId];
+      const output = cellRuntime.output;
+      const mimetype = output.mimetype;
+
+      // Determine output type
+      const isMedia = isMediaMimetype(mimetype, String(output.data));
+      const outputType = isMedia ? "media" : "text";
+
+      const cellName = displayCellName(cellData.name, cellIndex);
+
+      let processedContent: string | undefined;
+      let imageUrl: string | undefined;
+      let shouldDownloadImage = false;
+
+      // Process text content
+      if (outputType === "text" && typeof output.data === "string") {
+        processedContent =
+          mimetype === "text/html"
+            ? parseHtmlContent(output.data)
+            : output.data;
+      }
+
+      // Process media content - for now, we'll just note that it's media
+      if (outputType === "media") {
+        if (
+          typeof output.data === "string" &&
+          output.data.startsWith("data:")
+        ) {
+          imageUrl = output.data; // Data URL
+        } else if (
+          typeof output.data === "string" &&
+          output.data.startsWith("http")
+        ) {
+          imageUrl = output.data; // External URL
+        } else {
+          // Download the image from the DOM element
+          shouldDownloadImage = true;
+        }
+      }
+
+      items.push({
+        uri: this.asURI(cellId),
+        name: cellName,
+        type: this.contextType,
+        description: `Cell output (${mimetype || "unknown"})`,
+        data: {
+          cellId,
+          cellName,
+          cellCode: cellData?.code || "",
+          output,
+          outputType,
+          processedContent,
+          imageUrl,
+          shouldDownloadImage,
+        },
+      });
+    }
+
+    return items;
+  }
+
+  formatCompletion(item: CellOutputContextItem): Completion {
+    const { data } = item;
+    return {
+      label: `@${data.cellName}`,
+      displayLabel: data.cellName,
+      detail: `${data.outputType} output`,
+      boost: Boosts.CELL_OUTPUT,
+      type: this.contextType,
+      section: "Cell Output",
+      apply: `@${data.cellName}`,
+      info: () => {
+        const infoContainer = document.createElement("div");
+        infoContainer.classList.add(
+          "mo-cm-tooltip",
+          "docs-documentation",
+          "min-w-[300px]",
+          "max-w-[500px]",
+          "flex",
+          "flex-col",
+          "gap-2",
+        );
+
+        const headerDiv = document.createElement("div");
+        headerDiv.classList.add("flex", "flex-col", "gap-1");
+
+        const nameDiv = document.createElement("div");
+        nameDiv.classList.add("font-bold", "text-base");
+        nameDiv.textContent = data.cellName;
+        headerDiv.append(nameDiv);
+
+        const descriptionDiv = document.createElement("div");
+        descriptionDiv.classList.add("text-sm", "text-muted-foreground");
+        headerDiv.append(descriptionDiv);
+
+        infoContainer.append(headerDiv);
+
+        // Show cell code preview
+        if (data.cellCode) {
+          const codeHeaderDiv = document.createElement("div");
+          codeHeaderDiv.classList.add(
+            "text-xs",
+            "font-medium",
+            "text-muted-foreground",
+          );
+          codeHeaderDiv.textContent = "Code:";
+          infoContainer.append(codeHeaderDiv);
+
+          const codeDiv = document.createElement("div");
+          codeDiv.classList.add(
+            "text-xs",
+            "font-mono",
+            "bg-muted",
+            "p-2",
+            "rounded",
+            "max-h-20",
+            "overflow-y-auto",
+          );
+          codeDiv.textContent =
+            data.cellCode.slice(0, 200) +
+            (data.cellCode.length > 200 ? "..." : "");
+          infoContainer.append(codeDiv);
+        }
+
+        // Show output preview
+        if (data.processedContent) {
+          const outputHeaderDiv = document.createElement("div");
+          outputHeaderDiv.classList.add(
+            "text-xs",
+            "font-medium",
+            "text-muted-foreground",
+            "mt-2",
+          );
+          outputHeaderDiv.textContent = "Output Preview:";
+          infoContainer.append(outputHeaderDiv);
+
+          const outputDiv = document.createElement("div");
+          outputDiv.classList.add(
+            "text-xs",
+            "bg-muted",
+            "p-2",
+            "rounded",
+            "max-h-24",
+            "overflow-y-auto",
+            "mb-2",
+          );
+          outputDiv.textContent =
+            data.processedContent.slice(0, 300) +
+            (data.processedContent.length > 300 ? "..." : "");
+          infoContainer.append(outputDiv);
+        }
+
+        if (data.outputType === "media") {
+          const mediaDiv = document.createElement("div");
+          mediaDiv.classList.add(
+            "text-xs",
+            "text-muted-foreground",
+            "italic",
+            "mb-2",
+          );
+          mediaDiv.textContent =
+            "Contains media content (image, SVG, or canvas)";
+          infoContainer.append(mediaDiv);
+        }
+
+        return infoContainer;
+      },
+    };
+  }
+
+  formatContext(item: CellOutputContextItem): string {
+    const { data } = item;
+
+    const contextData = {
+      name: data.cellName,
+      cellId: data.cellId,
+      outputType: data.outputType,
+      mimetype: data.output.mimetype,
+    } as const;
+
+    let details = `Cell Code:\n${data.cellCode}\n\n`;
+
+    if (data.outputType === "text" && data.processedContent) {
+      details += `Output:\n${data.processedContent}`;
+    } else if (data.outputType === "media") {
+      details += `Media Output: Contains ${data.output.mimetype} content`;
+      if (data.imageUrl) {
+        details += `\nImage URL: ${data.imageUrl}`;
+      }
+    }
+
+    return contextToXml({
+      type: this.contextType,
+      data: contextData,
+      details,
+    });
+  }
+
+  /** Get attachments for cell output items that have shouldDownloadImage=true */
+  override async getAttachments(
+    items: CellOutputContextItem[],
+  ): Promise<FileUIPart[]> {
+    // Filter items that need image downloading
+    const itemsNeedingDownload = items.filter(
+      (item) =>
+        item.data.shouldDownloadImage && item.data.outputType === "media",
+    );
+
+    if (itemsNeedingDownload.length === 0) {
+      return [];
+    }
+
+    // Prepare download requests
+    const downloadRequests = itemsNeedingDownload.flatMap((item) => {
+      const outputElement = document.getElementById(
+        CellOutputId.create(item.data.cellId),
+      );
+      if (!outputElement) {
+        Logger.warn(`Output element not found for cell ${item.data.cellId}`);
+        return [];
+      }
+      return {
+        cellId: item.data.cellId,
+        cellName: item.data.cellName,
+        mimetype: item.data.output.mimetype,
+        element: outputElement,
+      };
+    });
+
+    try {
+      return await Promise.all(
+        downloadRequests.map(async (item) => ({
+          type: "file",
+          filename: `${item.cellName}-output-screenshot`,
+          mediaType: "image/png",
+          url: await toPng(item.element),
+        })),
+      );
+    } catch (error) {
+      Logger.error("Error downloading cell output images:", error);
+      return [];
+    }
+  }
+}

package/src/core/ai/context/providers/common.ts

@@ -3,6 +3,10 @@
 export const Boosts = {
   LOCAL_TABLE: 5,
   REMOTE_TABLE: 4,
+  HIGH: 4,
   VARIABLE: 3,
-  ERROR: 2,
+  MEDIUM: 3,
+  CELL_OUTPUT: 2,
+  LOW: 2,
+  ERROR: 1,
 } as const;

package/src/core/ai/context/providers/file.ts

@@ -0,0 +1,287 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import {
+  type Completion,
+  type CompletionContext,
+  type CompletionResult,
+  type CompletionSource,
+  closeCompletion,
+} from "@codemirror/autocomplete";
+import { toast } from "@/components/ui/use-toast";
+import { contextCallbacks } from "@/core/codemirror/ai/state";
+import type { EditRequests, FileInfo, RunRequests } from "@/core/network/types";
+import { deserializeBlob } from "@/utils/blob";
+import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
+import { Logger } from "@/utils/Logger";
+import { type AIContextItem, AIContextProvider } from "../registry";
+import { contextToXml } from "../utils";
+import { Boosts } from "./common";
+export interface FileContextItem extends AIContextItem {
+  type: "file";
+  data: {
+    path: string;
+    isDirectory: boolean;
+  };
+}
+
+export interface FileSearchConfig {
+  maxDepth: number;
+  maxResults: number;
+  defaultResultsLimit: number;
+  includeDirectories: boolean;
+}
+
+const DEFAULT_FILE_SEARCH_CONFIG: FileSearchConfig = {
+  maxDepth: 3,
+  maxResults: 20,
+  defaultResultsLimit: 5,
+  includeDirectories: false,
+};
+
+export class FileContextProvider extends AIContextProvider<FileContextItem> {
+  readonly title = "Files";
+  readonly mentionPrefix = "#";
+  readonly contextType = "file";
+
+  constructor(
+    private apiRequests: EditRequests & RunRequests,
+    private config: FileSearchConfig = DEFAULT_FILE_SEARCH_CONFIG,
+  ) {
+    super();
+  }
+
+  /**
+   * Create a dynamic completion source for file mentions
+   * This bypasses the standard registry system to enable dynamic searching
+   */
+  createCompletionSource(): CompletionSource {
+    return async (
+      context: CompletionContext,
+    ): Promise<CompletionResult | null> => {
+      // Look for # followed by any characters (including dots, slashes, etc.)
+      const match = context.matchBefore(/#[^\s#]*/);
+      if (!match) {
+        return null;
+      }
+
+      const matchText = match.text;
+      if (!matchText.startsWith("#")) {
+        return null;
+      }
+
+      const searchQuery = matchText.slice(1); // Remove the #
+      if (searchQuery.length === 0) {
+        // Show some popular files/directories even with no query
+        return this.getDefaultCompletions(match);
+      }
+
+      try {
+        const files = (await this.searchFiles(searchQuery)) || [];
+        const completions = files.map((file: FileInfo) => {
+          const item: FileContextItem = {
+            uri: this.asURI(file.path),
+            name: file.name,
+            type: this.contextType,
+            description: file.isDirectory ? "Directory" : "File",
+            data: {
+              path: file.path,
+              isDirectory: file.isDirectory,
+            },
+          };
+          return this.formatCompletion(item);
+        });
+
+        return {
+          from: match.from,
+          options: completions,
+        };
+      } catch (error) {
+        Logger.error("Failed to search files:", error);
+        return null;
+      }
+    };
+  }
+
+  private searchFiles = async (
+    query: string,
+    options: Partial<FileSearchConfig> = {},
+  ): Promise<FileInfo[]> => {
+    const { maxDepth, maxResults, includeDirectories } = {
+      ...this.config,
+      ...options,
+    };
+    const response = await this.apiRequests.sendSearchFiles({
+      query,
+      includeFiles: true,
+      includeDirectories: includeDirectories,
+      depth: maxDepth,
+      limit: maxResults,
+    });
+    return response.files;
+  };
+
+  private async getDefaultCompletions(match: {
+    from: number;
+  }): Promise<CompletionResult | null> {
+    try {
+      // Show common file types when no specific query is given
+      // Use broad searches for common file types
+      const searches = ["py", "md", "csv"];
+
+      // Try the first search that returns results
+      for (const search of searches) {
+        try {
+          const files = await this.searchFiles(search, {
+            maxDepth: 1,
+            maxResults: 5,
+          });
+
+          if (files && files.length > 0) {
+            const completions = files.map((file) => {
+              const item: FileContextItem = {
+                uri: this.asURI(file.path),
+                name: file.name,
+                type: this.contextType,
+                description: file.isDirectory ? "Directory" : "File",
+                data: {
+                  path: file.path,
+                  isDirectory: file.isDirectory,
+                },
+              };
+              return this.formatCompletion(item);
+            });
+
+            return {
+              from: match.from,
+              options: completions,
+            };
+          }
+        } catch (error) {
+          Logger.error("Failed to get default file completions:", error);
+        }
+      }
+
+      // If no searches return results, return empty
+      return {
+        from: match.from,
+        options: [],
+      };
+    } catch (error) {
+      Logger.error("Failed to get default file completions:", error);
+      return null;
+    }
+  }
+
+  getItems(): FileContextItem[] {
+    // Files are fetched dynamically, so return empty array
+    // This provider relies on dynamic fetching via createCompletionSource()
+    return [];
+  }
+
+  formatCompletion(item: FileContextItem): Completion {
+    const { data, name } = item;
+    const icon = data.isDirectory ? "📁" : "📄";
+
+    return {
+      ...this.createBasicCompletion(item),
+      type: "file",
+      section: "File",
+      boost: data.isDirectory ? Boosts.MEDIUM : Boosts.LOW,
+      detail: data.path,
+      displayLabel: `${icon} ${name}`,
+      apply: async (view, completion, from, to) => {
+        // First try to add the file as an attachment, if the callback is provided
+        // otherwise add it to the prompt
+        const addAttachment = view.state.facet(contextCallbacks)?.addAttachment;
+        if (!addAttachment) {
+          Logger.warn("No addAttachment callback provided");
+          return;
+        }
+
+        const fileDetails = await this.apiRequests
+          .sendFileDetails({ path: data.path })
+          .catch((error) => {
+            toast({
+              title: "Failed to get file details",
+              description: error.message,
+            });
+            return null;
+          });
+
+        if (!fileDetails) {
+          return;
+        }
+
+        const mimeType = fileDetails.mimeType || "text/plain";
+
+        // Handle binary vs text files
+        let blob: Blob;
+        if (
+          mimeType.startsWith("text/") ||
+          mimeType.includes("json") ||
+          mimeType.includes("xml")
+        ) {
+          // Text files - create blob directly from contents
+          blob = new Blob([fileDetails.contents || ""], { type: mimeType });
+        } else {
+          // Binary files - use blob utility to decode base64
+          if (fileDetails.contents) {
+            try {
+              // Create data URL using utility and deserialize blob
+              const dataURL = base64ToDataURL(
+                fileDetails.contents as Base64String,
+                mimeType,
+              );
+              blob = await deserializeBlob(dataURL);
+            } catch {
+              // Fallback to treating as text
+              blob = new Blob([fileDetails.contents], { type: mimeType });
+            }
+          } else {
+            blob = new Blob([""], { type: mimeType });
+          }
+        }
+
+        const file = new File([blob], name, { type: mimeType });
+        addAttachment(file);
+
+        // Close completion and delete the entire mention text (from # to cursor)
+        view.dispatch({
+          changes: { from, to, insert: "" },
+        });
+
+        closeCompletion(view);
+      },
+      info: () => {
+        const element = document.createElement("div");
+        element.classList.add("flex", "flex-col", "gap-1", "p-2");
+
+        const title = document.createElement("div");
+        title.classList.add("font-bold");
+        title.textContent = name;
+        element.append(title);
+
+        const path = document.createElement("div");
+        path.classList.add("text-xs", "text-muted-foreground");
+        path.textContent = data.path;
+        element.append(path);
+        return element;
+      },
+    };
+  }
+
+  formatContext(item: FileContextItem): string {
+    const { data, name } = item;
+    return contextToXml({
+      type: this.contextType,
+      data: {
+        name: name,
+        path: data.path,
+        isDirectory: data.isDirectory,
+      },
+      details: data.isDirectory
+        ? "Directory containing files and subdirectories"
+        : "File",
+    });
+  }
+}