@marimo-team/frontend 0.14.18-dev4 → 0.14.18-dev41

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

@@ -0,0 +1,208 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
+import type { DataTable } from "@/core/kernel/messages";
+import type { AIContextItem } from "../registry";
+import { AIContextProvider } from "../registry";
+import { Boosts } from "./common";
+
+export interface TableContextItem extends AIContextItem {
+  type: "data";
+  data: DataTable;
+}
+
+export class TableContextProvider extends AIContextProvider<TableContextItem> {
+  readonly title = "Tables";
+  readonly mentionPrefix = "@";
+  readonly contextType = "data";
+
+  constructor(private tablesMap: DatasetTablesMap) {
+    super();
+  }
+
+  getItems(): TableContextItem[] {
+    return [...this.tablesMap.entries()].map(([tableName, table]) => ({
+      uri: this.asURI(tableName),
+      name: tableName,
+      type: this.contextType,
+      description: table.source === "memory" ? "in-memory" : table.source,
+      data: table,
+    }));
+  }
+
+  formatContext(item: TableContextItem): string {
+    const { data, uri: id } = item;
+    const { columns, source, num_rows, num_columns } = data;
+    const shape = [
+      num_rows == null ? undefined : `${num_rows} rows`,
+      num_columns == null ? undefined : `${num_columns} columns`,
+    ]
+      .filter(Boolean)
+      .join(", ");
+
+    let context = `Table: ${id}\nSource: ${source || "unknown"}`;
+    if (shape) {
+      context += `\nShape: ${shape}`;
+    }
+
+    if (columns && columns.length > 0) {
+      context += `\nColumns:\n${columns.map((col) => `  - ${col.name}: ${col.type}`).join("\n")}`;
+    }
+
+    return context;
+  }
+
+  formatCompletion(item: TableContextItem): Completion {
+    const tableName = item.data.name;
+    const table = item.data;
+    return {
+      label: `@${tableName}`,
+      displayLabel: tableName,
+      detail: table.source === "memory" ? "in-memory" : table.source,
+      boost:
+        table.source_type === "local"
+          ? Boosts.LOCAL_TABLE
+          : Boosts.REMOTE_TABLE,
+      type: table.variable_name ? "dataframe" : "table",
+      apply: `@${tableName}`,
+      section: table.variable_name ? "Dataframe" : "Table",
+      info: () => this.createTableInfoElement(tableName, table),
+    };
+  }
+
+  private createTableInfoElement(
+    tableName: string,
+    table: DataTable,
+  ): HTMLElement {
+    const infoContainer = document.createElement("div");
+    infoContainer.classList.add(
+      "mo-cm-tooltip",
+      "docs-documentation",
+      "min-w-[200px]",
+      "flex",
+      "flex-col",
+      "gap-1",
+      "p-2",
+    );
+
+    // Table header with name and source
+    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 = tableName;
+    headerDiv.append(nameDiv);
+
+    if (table.source) {
+      const sourceDiv = document.createElement("div");
+      sourceDiv.classList.add("text-xs", "text-muted-foreground");
+      sourceDiv.textContent = `Source: ${table.source}`;
+      headerDiv.append(sourceDiv);
+    }
+
+    infoContainer.append(headerDiv);
+
+    // Table shape info
+    const shape = [
+      table.num_rows == null ? undefined : `${table.num_rows} rows`,
+      table.num_columns == null ? undefined : `${table.num_columns} columns`,
+    ]
+      .filter(Boolean)
+      .join(", ");
+
+    if (shape) {
+      const shapeDiv = document.createElement("div");
+      shapeDiv.classList.add("text-xs", "bg-muted", "px-2", "py-1", "rounded");
+      shapeDiv.textContent = shape;
+      infoContainer.append(shapeDiv);
+    }
+
+    // Columns table
+    if (table.columns && table.columns.length > 0) {
+      const columnsDiv = document.createElement("div");
+      columnsDiv.classList.add("overflow-auto", "max-h-60");
+
+      const columnsTable = document.createElement("table");
+      columnsTable.classList.add("w-full", "text-xs", "border-collapse");
+
+      // Table header
+      const headerRow = columnsTable.insertRow();
+      headerRow.classList.add("bg-muted");
+
+      const nameHeader = headerRow.insertCell();
+      nameHeader.classList.add("p-2", "font-medium", "text-left");
+      nameHeader.textContent = "Column";
+
+      const typeHeader = headerRow.insertCell();
+      typeHeader.classList.add("p-2", "font-medium", "text-left");
+      typeHeader.textContent = "Type";
+
+      // Check if any column has metadata before adding the header
+      const hasAnyMetadata = table.columns.some(
+        (column) => this.getItemMetadata(table, column) !== undefined,
+      );
+
+      let metadataHeader: HTMLTableCellElement | undefined;
+      if (hasAnyMetadata) {
+        metadataHeader = headerRow.insertCell();
+        metadataHeader.classList.add("p-2", "font-medium", "text-left");
+        metadataHeader.textContent = "Metadata";
+      }
+
+      // Table rows
+      table.columns.forEach((column, index) => {
+        const row = columnsTable.insertRow();
+        row.classList.add(index % 2 === 0 ? "bg-background" : "bg-muted/30");
+
+        const nameCell = row.insertCell();
+        nameCell.classList.add("p-2", "font-mono");
+        nameCell.textContent = column.name;
+
+        const typeCell = row.insertCell();
+        typeCell.classList.add("p-2", "text-muted-foreground");
+        typeCell.textContent = column.type;
+
+        if (hasAnyMetadata) {
+          const metadataCell = row.insertCell();
+          metadataCell.classList.add("p-2");
+
+          const itemMetadata = this.getItemMetadata(table, column);
+          if (itemMetadata) {
+            metadataCell.append(itemMetadata);
+          }
+        }
+      });
+
+      columnsDiv.append(columnsTable);
+      infoContainer.append(columnsDiv);
+    }
+
+    return infoContainer;
+  }
+
+  private getItemMetadata(
+    table: DataTable,
+    column: DataTable["columns"][0],
+  ): HTMLSpanElement | undefined {
+    const isPrimaryKey = table.primary_keys?.includes(column.name);
+    const isIndexed = table.indexes?.includes(column.name);
+    if (isPrimaryKey || isIndexed) {
+      const badge = document.createElement("span");
+      badge.textContent = isPrimaryKey ? "PK" : "IDX";
+      const color = isPrimaryKey
+        ? "bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-200"
+        : "bg-amber-100 text-amber-800 dark:bg-amber-900 dark:text-amber-200";
+      badge.classList.add(
+        "text-xs",
+        "px-1.5",
+        "py-0.5",
+        "rounded-full",
+        "font-medium",
+        ...color.split(" "),
+      );
+      return badge;
+    }
+  }
+}

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

@@ -0,0 +1,72 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import { createVariableInfoElement } from "@/core/codemirror/completion/variable-completions";
+import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
+import type { Variable, Variables } from "@/core/variables/types";
+import { type AIContextItem, AIContextProvider } from "../registry";
+import { Boosts } from "./common";
+
+export interface VariableContextItem extends AIContextItem {
+  type: "variable";
+  data: {
+    variable: Variable;
+  };
+}
+
+export class VariableContextProvider extends AIContextProvider<VariableContextItem> {
+  readonly title = "Variables";
+  readonly mentionPrefix = "@";
+  readonly contextType = "variable";
+
+  constructor(
+    private variables: Variables,
+    private tablesMap: DatasetTablesMap,
+  ) {
+    super();
+  }
+
+  getItems(): VariableContextItem[] {
+    const ignore = new Set(this.tablesMap.keys());
+
+    return Object.entries(this.variables).flatMap(([name, variable]) => {
+      if (ignore.has(name)) {
+        return [];
+      }
+      return [
+        {
+          uri: this.asURI(name),
+          name: name,
+          type: this.contextType,
+          description: variable.dataType ?? "",
+          dataType: variable.dataType,
+          data: {
+            variable,
+          },
+        },
+      ];
+    });
+  }
+
+  formatCompletion(item: VariableContextItem): Completion {
+    const { data } = item;
+    const { variable } = data;
+    return {
+      label: `@${variable.name}`,
+      displayLabel: variable.name,
+      detail: variable.dataType ?? "",
+      boost: Boosts.VARIABLE,
+      type: this.contextType,
+      section: "Variable",
+      info: () => {
+        return createVariableInfoElement(variable);
+      },
+    };
+  }
+
+  formatContext(item: VariableContextItem): string {
+    const { uri: id, data } = item;
+    const { variable } = data;
+    return `Variable: ${id}\nType: ${variable.dataType || "unknown"}\nPreview: ${JSON.stringify(variable.value)}`;
+  }
+}

package/src/core/ai/context/registry.ts

@@ -0,0 +1,180 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import type { Resource } from "@marimo-team/codemirror-mcp";
+import type { TypedString } from "@/utils/typed";
+
+/**
+ * Unique identifier for a context item in the format "type:id"
+ * e.g., "variable://my_var", "data://users", "file://config.py"
+ */
+export type ContextLocatorId = TypedString<"ContextLocatorId">;
+
+/**
+ * Base interface for context items that can be mentioned in AI prompts
+ */
+export interface AIContextItem extends Resource<Record<string, unknown>> {
+  data: Record<string, unknown>;
+}
+
+/**
+ * Abstract base class for AI context providers with shared implementation
+ */
+export abstract class AIContextProvider<
+  T extends AIContextItem = AIContextItem,
+> {
+  /** Human-readable title for this context type */
+  abstract readonly title: string;
+
+  /** Prefix used for mentions (e.g., "@" for variables, "#" for files) */
+  abstract readonly mentionPrefix: "@" | "#";
+
+  /** Context type identifier used in ContextLocatorId */
+  abstract readonly contextType: string;
+
+  /** Get all available items of this context type */
+  abstract getItems(): T[];
+
+  /** Format the context for inclusion in AI prompt */
+  abstract formatContext(item: T): string;
+
+  /** Format completion */
+  abstract formatCompletion(item: T): Completion;
+
+  asURI(id: string): ContextLocatorId {
+    return `${this.contextType}://${id}` as ContextLocatorId;
+  }
+
+  /** Parse context IDs from input text using the provider's mention prefix */
+  parseContextIds(input: string): ContextLocatorId[] {
+    // Match @type://id, e.g., @data://users
+    const regex = new RegExp(
+      `${this.mentionPrefix}([\\w-]+):\\/\\/([\\w./-]+)`,
+      "g",
+    );
+    const matches = [...input.matchAll(regex)];
+    const justURI = (match: string) => match.slice(1);
+
+    const results = matches
+      .filter(([, type]) => type === this.contextType)
+      .map(([matchString]) => justURI(matchString) as ContextLocatorId);
+
+    return [...new Set(results)];
+  }
+
+  /** Create a basic completion object - can be used by subclasses */
+  protected createBasicCompletion(
+    item: T,
+    options?: {
+      boost?: number;
+      type?: string;
+      section?: string;
+      detail?: string;
+    },
+  ): Completion {
+    return {
+      label: `${this.mentionPrefix}${item.uri.split("://")[1]}`,
+      displayLabel: item.name,
+      detail: options?.detail || item.description,
+      boost: options?.boost || 1,
+      type: options?.type || this.contextType,
+      apply: `${this.mentionPrefix}${item.uri}`,
+      section: options?.section || this.title,
+    };
+  }
+}
+
+/**
+ * Registry for managing different AI context providers
+ */
+export class AIContextRegistry<T extends AIContextItem> {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private providers = new Set<AIContextProvider<any>>();
+
+  /**
+   * Register a new context provider
+   */
+  register<U extends AIContextItem>(
+    provider: AIContextProvider<U>,
+    // eslint-disable-next-line @typescript-eslint/prefer-return-this-type
+  ): AIContextRegistry<U | T> {
+    this.providers.add(provider);
+    return this;
+  }
+
+  /**
+   * Get all registered providers
+   */
+  getProviders(): Set<AIContextProvider<T>> {
+    return this.providers;
+  }
+
+  /**
+   * Get a specific provider by type
+   */
+  getProvider(type: string): AIContextProvider | undefined {
+    return [...this.providers].find(
+      (provider) => provider.contextType === type,
+    );
+  }
+
+  getAllItems(): T[] {
+    return [...this.providers].flatMap((provider) => provider.getItems());
+  }
+
+  /**
+   * Parse context IDs from input across all providers
+   */
+  parseAllContextIds(input: string): ContextLocatorId[] {
+    return [...this.providers].flatMap((provider) =>
+      provider.parseContextIds(input),
+    );
+  }
+
+  /**
+   * Get context information for mentioned items
+   */
+  getContextInfo(contextIds: ContextLocatorId[]): T[] {
+    const contextInfo: T[] = [];
+    const allItems = new Map<ContextLocatorId, T>(
+      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
+    );
+
+    for (const contextId of contextIds) {
+      const item = allItems.get(contextId);
+      if (item) {
+        contextInfo.push(item);
+      }
+    }
+
+    return contextInfo;
+  }
+
+  /**
+   * Format context for AI prompt inclusion
+   */
+  formatContextForAI(contextIds: ContextLocatorId[]): string {
+    const allItems = new Map<ContextLocatorId, T>(
+      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
+    );
+
+    const contextInfo: T[] = [];
+    for (const contextId of contextIds) {
+      const item = allItems.get(contextId);
+      if (item) {
+        contextInfo.push(item);
+      }
+    }
+
+    if (contextInfo.length === 0) {
+      return "";
+    }
+
+    return contextInfo
+      .map((item) => {
+        const provider = this.getProvider(item.type);
+        return provider?.formatContext(item) || "";
+      })
+      .join("\n\n");
+  }
+}

package/src/core/ai/ids/__tests__/ids.test.ts

@@ -0,0 +1,193 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type { ProviderId } from "../ids";
+import { AiModelId, type ShortModelId } from "../ids";
+
+describe("AiModelId", () => {
+  describe("constructor", () => {
+    it("should create an instance with providerId and shortModelId", () => {
+      const modelId = new AiModelId("openai", "gpt-4" as ShortModelId);
+      expect(modelId.providerId).toBe("openai");
+      expect(modelId.shortModelId).toBe("gpt-4");
+    });
+  });
+
+  describe("id getter", () => {
+    it("should return qualified model id", () => {
+      const modelId = new AiModelId("anthropic", "claude-3" as ShortModelId);
+      expect(modelId.id).toBe("anthropic/claude-3");
+    });
+
+    it("should handle all provider types", () => {
+      const providers: ProviderId[] = [
+        "openai",
+        "anthropic",
+        "google",
+        "ollama",
+        "bedrock",
+        "github",
+      ];
+
+      providers.forEach((provider) => {
+        const modelId = new AiModelId(provider, "test-model" as ShortModelId);
+        expect(modelId.id).toBe(`${provider}/test-model`);
+      });
+    });
+  });
+
+  describe("parse", () => {
+    describe("with qualified ids (containing '/')", () => {
+      it("should parse openai qualified id", () => {
+        const modelId = AiModelId.parse("openai/gpt-4");
+        expect(modelId.providerId).toBe("openai");
+        expect(modelId.shortModelId).toBe("gpt-4");
+      });
+
+      it("should parse anthropic qualified id", () => {
+        const modelId = AiModelId.parse("anthropic/claude-3-sonnet");
+        expect(modelId.providerId).toBe("anthropic");
+        expect(modelId.shortModelId).toBe("claude-3-sonnet");
+      });
+
+      it("should parse google qualified id", () => {
+        const modelId = AiModelId.parse("google/gemini-pro");
+        expect(modelId.providerId).toBe("google");
+        expect(modelId.shortModelId).toBe("gemini-pro");
+      });
+
+      it("should parse ollama qualified id", () => {
+        const modelId = AiModelId.parse("ollama/llama2");
+        expect(modelId.providerId).toBe("ollama");
+        expect(modelId.shortModelId).toBe("llama2");
+      });
+
+      it("should parse bedrock qualified id", () => {
+        const modelId = AiModelId.parse("bedrock/titan-text");
+        expect(modelId.providerId).toBe("bedrock");
+        expect(modelId.shortModelId).toBe("titan-text");
+      });
+
+      it("should parse github qualified id", () => {
+        const modelId = AiModelId.parse("github/gpt-4o");
+        expect(modelId.providerId).toBe("github");
+        expect(modelId.shortModelId).toBe("gpt-4o");
+      });
+
+      it("should handle multiple slashes", () => {
+        const modelId = AiModelId.parse("openai/gpt-4/turbo");
+        expect(modelId.providerId).toBe("openai");
+        expect(modelId.shortModelId).toBe("gpt-4/turbo");
+      });
+    });
+
+    describe("without qualified ids (no '/')", () => {
+      describe("openai models", () => {
+        it("should guess openai for gpt models", () => {
+          const modelId = AiModelId.parse("gpt-4");
+          expect(modelId.providerId).toBe("openai");
+          expect(modelId.shortModelId).toBe("gpt-4");
+        });
+
+        it("should guess openai for gpt-3.5 models", () => {
+          const modelId = AiModelId.parse("gpt-3.5-turbo");
+          expect(modelId.providerId).toBe("openai");
+          expect(modelId.shortModelId).toBe("gpt-3.5-turbo");
+        });
+
+        it("should guess openai for o3 models", () => {
+          const modelId = AiModelId.parse("o3-mini");
+          expect(modelId.providerId).toBe("openai");
+          expect(modelId.shortModelId).toBe("o3-mini");
+        });
+
+        it("should guess openai for o1 models", () => {
+          const modelId = AiModelId.parse("o1-preview");
+          expect(modelId.providerId).toBe("openai");
+          expect(modelId.shortModelId).toBe("o1-preview");
+        });
+      });
+
+      describe("anthropic models", () => {
+        it("should guess anthropic for claude models", () => {
+          const modelId = AiModelId.parse("claude-3-sonnet");
+          expect(modelId.providerId).toBe("anthropic");
+          expect(modelId.shortModelId).toBe("claude-3-sonnet");
+        });
+
+        it("should guess anthropic for claude-instant", () => {
+          const modelId = AiModelId.parse("claude-instant");
+          expect(modelId.providerId).toBe("anthropic");
+          expect(modelId.shortModelId).toBe("claude-instant");
+        });
+      });
+
+      describe("google models", () => {
+        it("should guess google for gemini models", () => {
+          const modelId = AiModelId.parse("gemini-pro");
+          expect(modelId.providerId).toBe("google");
+          expect(modelId.shortModelId).toBe("gemini-pro");
+        });
+
+        it("should guess google for google-prefixed models", () => {
+          const modelId = AiModelId.parse("google-palm");
+          expect(modelId.providerId).toBe("google");
+          expect(modelId.shortModelId).toBe("google-palm");
+        });
+      });
+
+      describe("ollama fallback", () => {
+        it("should default to ollama for unknown models", () => {
+          const modelId = AiModelId.parse("llama2");
+          expect(modelId.providerId).toBe("ollama");
+          expect(modelId.shortModelId).toBe("llama2");
+        });
+
+        it("should default to ollama for custom models", () => {
+          const modelId = AiModelId.parse("my-custom-model");
+          expect(modelId.providerId).toBe("ollama");
+          expect(modelId.shortModelId).toBe("my-custom-model");
+        });
+
+        it("should default to ollama for empty string", () => {
+          const modelId = AiModelId.parse("");
+          expect(modelId.providerId).toBe("ollama");
+          expect(modelId.shortModelId).toBe("");
+        });
+      });
+    });
+  });
+
+  describe("round-trip parsing", () => {
+    it("should maintain consistency when parsing generated ids", () => {
+      const original = new AiModelId(
+        "anthropic",
+        "claude-3-opus" as ShortModelId,
+      );
+      const parsed = AiModelId.parse(original.id);
+
+      expect(parsed.providerId).toBe(original.providerId);
+      expect(parsed.shortModelId).toBe(original.shortModelId);
+      expect(parsed.id).toBe(original.id);
+    });
+
+    it("should work for all providers", () => {
+      const providers: ProviderId[] = [
+        "openai",
+        "anthropic",
+        "google",
+        "ollama",
+        "bedrock",
+        "github",
+      ];
+
+      providers.forEach((provider) => {
+        const original = new AiModelId(provider, "test-model" as ShortModelId);
+        const parsed = AiModelId.parse(original.id);
+
+        expect(parsed.providerId).toBe(provider);
+        expect(parsed.shortModelId).toBe("test-model");
+      });
+    });
+  });
+});