@marimo-team/frontend 0.14.18-dev24 → 0.14.18-dev26

package/src/components/app-config/incorrect-model-id.tsx

@@ -1,14 +1,16 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import React from "react";
+import { AiModelId, type QualifiedModelId } from "@/core/ai/ids/ids";
 import { Banner } from "@/plugins/impl/common/error-banner";
-import { AiModelId, type QualifiedModelId } from "@/utils/ai/ids";
 
 interface IncorrectModelIdProps {
   value: string | null | undefined;
+  includeSuggestion?: boolean;
 }
 
 export const IncorrectModelId: React.FC<IncorrectModelIdProps> = ({
   value,
+  includeSuggestion = true,
 }) => {
   if (!value) {
     return null;
@@ -31,7 +33,7 @@ export const IncorrectModelId: React.FC<IncorrectModelIdProps> = ({
         provider.
       </span>
       <br />
-      {suggestion && (
+      {includeSuggestion && suggestion && (
         <span>
           Did you mean <code className="font-bold">{suggestion}</code>?
         </span>

package/src/components/app-config/user-config-form.tsx

@@ -12,7 +12,7 @@ import {
   MonitorIcon,
   PackageIcon,
 } from "lucide-react";
-import React, { useRef } from "react";
+import React, { useId, useRef } from "react";
 import { useForm } from "react-hook-form";
 import { Button } from "@/components/ui/button";
 import { Checkbox } from "@/components/ui/checkbox";
@@ -130,6 +130,8 @@ export const UserConfigForm: React.FC = () => {
   };
 
   const isWasmRuntime = isWasm();
+  const htmlCheckboxId = useId();
+  const ipynbCheckboxId = useId();
 
   const renderBody = () => {
     switch (activeCategory) {
@@ -168,6 +170,7 @@ export const UserConfigForm: React.FC = () => {
                     <FormLabel>Autosave delay (seconds)</FormLabel>
                     <FormControl>
                       <NumberField
+                        aria-label="Autosave delay"
                         data-testid="autosave-delay-input"
                         className="m-0 w-24"
                         isDisabled={
@@ -205,7 +208,7 @@ export const UserConfigForm: React.FC = () => {
                         <div className="flex gap-4">
                           <div className="flex items-center space-x-2">
                             <Checkbox
-                              id="html-checkbox"
+                              id={htmlCheckboxId}
                               checked={
                                 Array.isArray(field.value) &&
                                 field.value.includes("html")
@@ -219,11 +222,11 @@ export const UserConfigForm: React.FC = () => {
                                 );
                               }}
                             />
-                            <FormLabel htmlFor="html-checkbox">HTML</FormLabel>
+                            <FormLabel htmlFor={htmlCheckboxId}>HTML</FormLabel>
                           </div>
                           <div className="flex items-center space-x-2">
                             <Checkbox
-                              id="ipynb-checkbox"
+                              id={ipynbCheckboxId}
                               checked={
                                 Array.isArray(field.value) &&
                                 field.value.includes("ipynb")
@@ -237,7 +240,7 @@ export const UserConfigForm: React.FC = () => {
                                 );
                               }}
                             />
-                            <FormLabel htmlFor="ipynb-checkbox">
+                            <FormLabel htmlFor={ipynbCheckboxId}>
                               IPYNB
                             </FormLabel>
                           </div>
@@ -295,6 +298,7 @@ export const UserConfigForm: React.FC = () => {
                       <FormLabel>Line length</FormLabel>
                       <FormControl>
                         <NumberField
+                          aria-label="Line length"
                           data-testid="line-length-input"
                           className="m-0 w-24"
                           {...field}
@@ -716,6 +720,7 @@ export const UserConfigForm: React.FC = () => {
                     <FormControl>
                       <span className="inline-flex mr-2">
                         <NumberField
+                          aria-label="Code editor font size"
                           data-testid="code-editor-font-size-input"
                           className="m-0 w-24"
                           {...field}
@@ -848,6 +853,7 @@ export const UserConfigForm: React.FC = () => {
                       <FormLabel>Default table page size</FormLabel>
                       <FormControl>
                         <NumberField
+                          aria-label="Default table page size"
                           data-testid="default-table-page-size-input"
                           className="m-0 w-24"
                           {...field}
@@ -884,6 +890,7 @@ export const UserConfigForm: React.FC = () => {
                       <FormLabel>Default table max columns</FormLabel>
                       <FormControl>
                         <NumberField
+                          aria-label="Default table max columns"
                           data-testid="default-table-max-columns-input"
                           className="m-0 w-24"
                           {...field}

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

@@ -36,6 +36,7 @@ import {
   SelectValue,
 } from "@/components/ui/select";
 import { addMessageToChat } from "@/core/ai/chat-utils";
+import type { QualifiedModelId } from "@/core/ai/ids/ids";
 import {
   activeChatAtom,
   type Chat,
@@ -54,7 +55,7 @@ import { cn } from "@/utils/cn";
 import { timeAgo } from "@/utils/dates";
 import { Logger } from "@/utils/Logger";
 import { generateUUID } from "@/utils/uuid";
-import { KNOWN_AI_MODELS } from "../app-config/constants";
+import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
 import { PromptInput } from "../editor/ai/add-cell-with-ai";
 import { getAICompletionBody } from "../editor/ai/completion-utils";
@@ -276,7 +277,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
       saveConfig(newConfig);
     };
 
-    const handleModelChange = async (newModel: string) => {
+    const handleModelChange = async (newModel: QualifiedModelId) => {
       const newConfig: UserConfig = {
         ...userConfig,
         ai: {
@@ -324,30 +325,15 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
               </SelectContent>
             </Select>
           </FeatureFlagged>
-          <Select value={currentModel} onValueChange={handleModelChange}>
-            <SelectTrigger className="h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1">
-              <SelectValue placeholder="Model" />
-            </SelectTrigger>
-            <SelectContent>
-              {/* Show current model if it's not in the known models list */}
-              {!(KNOWN_AI_MODELS as readonly string[]).includes(
-                currentModel,
-              ) && (
-                <SelectItem
-                  key={currentModel}
-                  value={currentModel}
-                  className="text-sm"
-                >
-                  {currentModel}
-                </SelectItem>
-              )}
-              {KNOWN_AI_MODELS.map((model) => (
-                <SelectItem key={model} value={model} className="text-sm">
-                  {model}
-                </SelectItem>
-              ))}
-            </SelectContent>
-          </Select>
+          <AIModelDropdown
+            value={currentModel}
+            placeholder="Model"
+            onSelect={handleModelChange}
+            triggerClassName="h-6 text-xs shadow-none! ring-0! bg-muted hover:bg-muted/30 rounded-sm"
+            iconSize="small"
+            showAddCustomModelDocs={true}
+            forRole="chat"
+          />
         </div>
         <Button
           variant="ghost"

package/src/components/slides/slides.css

@@ -1,6 +1,5 @@
 @import "swiper/css";
 @import "swiper/css/virtual";
-@import "swiper/css/keyboard";
 @import "swiper/css/navigation";
 @import "swiper/css/pagination";
 @import "swiper/css/scrollbar";

package/src/core/ai/__tests__/model-registry.test.ts

@@ -0,0 +1,357 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+// Mock the models.json import
+vi.mock("@marimo-team/llm-info/models.json", () => {
+  const models: AiModel[] = [
+    {
+      name: "GPT-4",
+      model: "gpt-4",
+      description: "OpenAI GPT-4 model",
+      providers: ["openai"],
+      roles: ["chat", "edit"],
+      thinking: false,
+    },
+    {
+      name: "Claude 3",
+      model: "claude-3-sonnet",
+      description: "Anthropic Claude 3 Sonnet",
+      providers: ["anthropic"],
+      roles: ["chat", "edit"],
+      thinking: false,
+    },
+    {
+      name: "Gemini Pro",
+      model: "gemini-pro",
+      description: "Google Gemini Pro model",
+      providers: ["google"],
+      roles: ["chat", "edit"],
+      thinking: false,
+    },
+    {
+      name: "Multi Provider Model",
+      model: "multi-model",
+      description: "Model available on multiple providers",
+      providers: ["openai", "anthropic"],
+      roles: ["chat", "edit"],
+      thinking: false,
+    },
+  ];
+
+  return {
+    models: models,
+  };
+});
+
+import type { AiModel } from "@marimo-team/llm-info";
+import { AiModelRegistry } from "../model-registry";
+
+describe("AiModelRegistry", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe("create", () => {
+    it("should create registry with no custom or displayed models", () => {
+      const registry = AiModelRegistry.create({});
+
+      expect(registry).toBeInstanceOf(AiModelRegistry);
+      expect(registry.getCustomModels()).toEqual(new Set());
+      expect(registry.getDisplayedModels()).toEqual(new Set());
+    });
+
+    it("should create registry with custom models", () => {
+      const customModels = ["openai/custom-gpt", "anthropic/custom-claude"];
+      const registry = AiModelRegistry.create({ customModels });
+
+      expect(registry.getCustomModels()).toEqual(new Set(customModels));
+      expect(registry.getDisplayedModels()).toEqual(new Set());
+    });
+
+    it("should create registry with displayed models", () => {
+      const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
+      const registry = AiModelRegistry.create({ displayedModels });
+
+      expect(registry.getCustomModels()).toEqual(new Set());
+      expect(registry.getDisplayedModels()).toEqual(new Set(displayedModels));
+    });
+
+    it("should create registry with both custom and displayed models", () => {
+      const customModels = ["openai/custom-gpt"];
+      const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
+      const registry = AiModelRegistry.create({
+        customModels,
+        displayedModels,
+      });
+
+      expect(registry.getCustomModels()).toEqual(new Set(customModels));
+      expect(registry.getDisplayedModels()).toEqual(new Set(displayedModels));
+    });
+  });
+
+  describe("getModelsByProvider", () => {
+    it("should return models for a specific provider", () => {
+      const registry = AiModelRegistry.create({});
+      const openaiModels = registry.getModelsByProvider("openai");
+
+      expect(openaiModels).toHaveLength(2); // gpt-4 and multi-model
+      expect(
+        openaiModels.every((model) => model.providers.includes("openai")),
+      ).toBe(true);
+    });
+
+    it("should return empty array for provider with no models", () => {
+      const registry = AiModelRegistry.create({});
+      const azureModels = registry.getModelsByProvider("azure");
+
+      expect(azureModels).toEqual([]);
+    });
+
+    it("should include custom models for the provider", () => {
+      const customModels = ["openai/custom-gpt"];
+      const registry = AiModelRegistry.create({ customModels });
+      const openaiModels = registry.getModelsByProvider("openai");
+
+      const customModel = openaiModels.find((model) => model.custom);
+      expect(customModel).toBeDefined();
+      expect(customModel?.name).toBe("openai/custom-gpt");
+      expect(customModel?.model).toBe("custom-gpt");
+      expect(customModel?.description).toBe("Custom model");
+      expect(customModel?.providers).toEqual(["openai"]);
+      expect(customModel?.roles).toEqual([]);
+      expect(customModel?.thinking).toBe(false);
+    });
+
+    it("should filter models based on displayed models", () => {
+      const displayedModels = ["openai/gpt-4"];
+      const registry = AiModelRegistry.create({ displayedModels });
+      const openaiModels = registry.getModelsByProvider("openai");
+
+      expect(openaiModels).toHaveLength(1);
+      expect(openaiModels[0].model).toBe("gpt-4");
+    });
+
+    it("should filter custom models based on displayed models", () => {
+      const customModels = ["openai/custom-gpt", "anthropic/custom-claude"];
+      const displayedModels = ["openai/custom-gpt"];
+      const registry = AiModelRegistry.create({
+        customModels,
+        displayedModels,
+      });
+
+      const openaiModels = registry.getModelsByProvider("openai");
+      const anthropicModels = registry.getModelsByProvider("anthropic");
+
+      expect(
+        openaiModels.some(
+          (model) => model.custom && model.model === "custom-gpt",
+        ),
+      ).toBe(true);
+      expect(
+        anthropicModels.some(
+          (model) => model.custom && model.model === "custom-claude",
+        ),
+      ).toBe(false);
+    });
+  });
+
+  describe("getGroupedModelsByProvider", () => {
+    it("should return all models grouped by provider", () => {
+      const registry = AiModelRegistry.create({});
+      const groupedModels = registry.getGroupedModelsByProvider();
+
+      expect(groupedModels.has("openai")).toBe(true);
+      expect(groupedModels.has("anthropic")).toBe(true);
+      expect(groupedModels.has("google")).toBe(true);
+
+      const openaiModels = groupedModels.get("openai") || [];
+      const anthropicModels = groupedModels.get("anthropic") || [];
+      const googleModels = groupedModels.get("google") || [];
+
+      expect(openaiModels.length).toEqual(2);
+      expect(anthropicModels.length).toEqual(2);
+      expect(googleModels.length).toEqual(1);
+    });
+
+    it("should include custom models in the grouped results", () => {
+      const customModels = ["openai/custom-gpt", "anthropic/custom-claude"];
+      const registry = AiModelRegistry.create({ customModels });
+      const groupedModels = registry.getGroupedModelsByProvider();
+
+      const openaiModels = groupedModels.get("openai") || [];
+      const anthropicModels = groupedModels.get("anthropic") || [];
+
+      expect(
+        openaiModels.some(
+          (model) => model.custom && model.model === "custom-gpt",
+        ),
+      ).toBe(true);
+      expect(
+        anthropicModels.some(
+          (model) => model.custom && model.model === "custom-claude",
+        ),
+      ).toBe(true);
+    });
+
+    it("should respect displayed models filter", () => {
+      const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
+      const registry = AiModelRegistry.create({ displayedModels });
+      const groupedModels = registry.getGroupedModelsByProvider();
+
+      const openaiModels = groupedModels.get("openai") || [];
+      const anthropicModels = groupedModels.get("anthropic") || [];
+      const googleModels = groupedModels.get("google") || [];
+
+      expect(openaiModels.length).toBe(1);
+      expect(openaiModels[0].model).toBe("gpt-4");
+      expect(anthropicModels.length).toBe(1);
+      expect(anthropicModels[0].model).toBe("claude-3-sonnet");
+      expect(googleModels.length).toBe(0);
+    });
+  });
+
+  describe("getCustomModels", () => {
+    it("should return empty set when no custom models", () => {
+      const registry = AiModelRegistry.create({});
+      expect(registry.getCustomModels()).toEqual(new Set());
+    });
+
+    it("should return set of custom model IDs", () => {
+      const customModels = ["openai/custom-gpt", "anthropic/custom-claude"];
+      const registry = AiModelRegistry.create({ customModels });
+      expect(registry.getCustomModels()).toEqual(new Set(customModels));
+    });
+  });
+
+  describe("getDisplayedModels", () => {
+    it("should return empty set when no displayed models specified", () => {
+      const registry = AiModelRegistry.create({});
+      expect(registry.getDisplayedModels()).toEqual(new Set());
+    });
+
+    it("should return set of displayed model IDs", () => {
+      const displayedModels = ["openai/gpt-4", "anthropic/claude-3-sonnet"];
+      const registry = AiModelRegistry.create({ displayedModels });
+      expect(registry.getDisplayedModels()).toEqual(new Set(displayedModels));
+    });
+  });
+
+  describe("edge cases", () => {
+    it("should handle empty arrays for custom and displayed models", () => {
+      const registry = AiModelRegistry.create({
+        customModels: [],
+        displayedModels: [],
+      });
+
+      expect(registry.getCustomModels()).toEqual(new Set());
+      expect(registry.getDisplayedModels()).toEqual(new Set());
+
+      // Should still load default models
+      const openaiModels = registry.getModelsByProvider("openai");
+      expect(openaiModels.length).toBeGreaterThan(0);
+    });
+
+    it("should handle models with multiple providers", () => {
+      const registry = AiModelRegistry.create({});
+
+      const openaiModels = registry.getModelsByProvider("openai");
+      const anthropicModels = registry.getModelsByProvider("anthropic");
+
+      // The multi-model should appear in both providers
+      const multiModelInOpenai = openaiModels.find(
+        (model) => model.model === "multi-model",
+      );
+      const multiModelInAnthropic = anthropicModels.find(
+        (model) => model.model === "multi-model",
+      );
+
+      expect(multiModelInOpenai).toBeDefined();
+      expect(multiModelInAnthropic).toBeDefined();
+      expect(multiModelInOpenai).toEqual(multiModelInAnthropic);
+    });
+
+    it("should handle displayed models filter with non-existent models", () => {
+      const displayedModels = [
+        "openai/non-existent-model",
+        "anthropic/claude-3-sonnet",
+      ];
+      const registry = AiModelRegistry.create({ displayedModels });
+
+      const openaiModels = registry.getModelsByProvider("openai");
+      const anthropicModels = registry.getModelsByProvider("anthropic");
+
+      // Should only show the existing model
+      expect(openaiModels.length).toBe(0);
+      expect(anthropicModels.length).toBe(1);
+      expect(anthropicModels[0].model).toBe("claude-3-sonnet");
+    });
+  });
+
+  describe("model properties", () => {
+    it("should ensure all models have required properties", () => {
+      const registry = AiModelRegistry.create({});
+      const groupedModels = registry.getGroupedModelsByProvider();
+
+      for (const [provider, models] of groupedModels.entries()) {
+        for (const model of models) {
+          expect(model).toHaveProperty("name");
+          expect(model).toHaveProperty("model");
+          expect(model).toHaveProperty("description");
+          expect(model).toHaveProperty("providers");
+          expect(model).toHaveProperty("roles");
+          expect(model).toHaveProperty("thinking");
+          expect(model).toHaveProperty("custom");
+
+          expect(typeof model.name).toBe("string");
+          expect(typeof model.model).toBe("string");
+          expect(typeof model.description).toBe("string");
+          expect(Array.isArray(model.providers)).toBe(true);
+          expect(Array.isArray(model.roles)).toBe(true);
+          expect(typeof model.thinking).toBe("boolean");
+          expect(typeof model.custom).toBe("boolean");
+
+          expect(model.providers).toContain(provider);
+        }
+      }
+    });
+
+    it("should ensure custom models have correct custom property", () => {
+      const customModels = ["openai/custom-gpt"];
+      const registry = AiModelRegistry.create({ customModels });
+      const openaiModels = registry.getModelsByProvider("openai");
+
+      const customModel = openaiModels.find((model) => model.custom);
+      const defaultModel = openaiModels.find((model) => !model.custom);
+
+      expect(customModel).toMatchInlineSnapshot(`
+        {
+          "custom": true,
+          "description": "Custom model",
+          "model": "custom-gpt",
+          "name": "openai/custom-gpt",
+          "providers": [
+            "openai",
+          ],
+          "roles": [],
+          "thinking": false,
+        }
+      `);
+      expect(defaultModel).toMatchInlineSnapshot(`
+        {
+          "custom": false,
+          "description": "OpenAI GPT-4 model",
+          "model": "gpt-4",
+          "name": "GPT-4",
+          "providers": [
+            "openai",
+          ],
+          "roles": [
+            "chat",
+            "edit",
+          ],
+          "thinking": false,
+        }
+      `);
+    });
+  });
+});

package/src/{utils/ai → core/ai/ids}/__tests__/ids.test.ts

@@ -1,7 +1,8 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
-import { AiModelId, type ProviderId, type ShortModelId } from "../ids";
+import type { ProviderId } from "../ids";
+import { AiModelId, type ShortModelId } from "../ids";
 
 describe("AiModelId", () => {
   describe("constructor", () => {

package/src/{utils/ai → core/ai/ids}/ids.ts

@@ -1,16 +1,17 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { TypedString } from "../typed";
+import type { TypedString } from "@/utils/typed";
 
-/**
- * Supported providers by the marimo server.
- */
-export type ProviderId =
-  | "openai"
-  | "anthropic"
-  | "google"
-  | "ollama"
-  | "bedrock";
+export const PROVIDERS = [
+  "openai",
+  "anthropic",
+  "google",
+  "ollama",
+  "bedrock",
+  "deepseek",
+  "azure",
+] as const;
+export type ProviderId = (typeof PROVIDERS)[number];
 
 export type ShortModelId = TypedString<"ShortModelId">;
 
@@ -57,5 +58,12 @@ function guessProviderId(id: string): ProviderId {
   if (id.startsWith("gemini") || id.startsWith("google")) {
     return "google";
   }
+  if (id.startsWith("deepseek")) {
+    return "deepseek";
+  }
   return "ollama";
 }
+
+export function isKnownAIProvider(providerId: ProviderId): boolean {
+  return PROVIDERS.includes(providerId);
+}