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

package/src/components/ai/ai-model-dropdown.tsx

@@ -0,0 +1,288 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Role } from "@marimo-team/llm-info";
+import { capitalize } from "lodash-es";
+import { ChevronDownIcon, CircleHelpIcon } from "lucide-react";
+import {
+  AiModelId,
+  isKnownAIProvider,
+  type ProviderId,
+  type QualifiedModelId,
+} from "@/core/ai/ids/ids";
+import { type AiModel, AiModelRegistry } from "@/core/ai/model-registry";
+import { useResolvedMarimoConfig } from "@/core/config/config";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuPortal,
+  DropdownMenuSeparator,
+  DropdownMenuSub,
+  DropdownMenuSubContent,
+  DropdownMenuSubTrigger,
+  DropdownMenuTrigger,
+} from "../ui/dropdown-menu";
+import { AiProviderIcon } from "./ai-provider-icon";
+
+interface AIModelDropdownProps {
+  value?: string;
+  placeholder?: string;
+  onSelect: (modelId: QualifiedModelId) => void;
+  triggerClassName?: string;
+  customDropdownContent?: React.ReactNode;
+  iconSize?: "medium" | "small";
+  showAddCustomModelDocs?: boolean;
+  forRole?: Role;
+}
+
+export const AIModelDropdown = ({
+  value,
+  placeholder,
+  onSelect,
+  triggerClassName,
+  customDropdownContent,
+  iconSize = "medium",
+  showAddCustomModelDocs = false,
+  forRole,
+}: AIModelDropdownProps) => {
+  const currentValue = value ? AiModelId.parse(value) : undefined;
+
+  const selectModel = (modelId: QualifiedModelId) => {
+    onSelect(modelId);
+  };
+
+  const [marimoConfig] = useResolvedMarimoConfig();
+  const configModels = marimoConfig.ai?.models;
+
+  // Only include autocompleteModel if copilot is set to "custom"
+  const autocompleteModel =
+    marimoConfig.completion.copilot === "custom"
+      ? configModels?.autocomplete_model
+      : undefined;
+
+  const aiModelRegistry = AiModelRegistry.create({
+    // We add all the custom models and the models used in the editor.
+    // If they among the known models, they won't overwrite them.
+    customModels: [
+      ...(configModels?.custom_models ?? []),
+      configModels?.chat_model,
+      autocompleteModel,
+      configModels?.edit_model,
+    ].filter(Boolean),
+    displayedModels: configModels?.displayed_models,
+  });
+  const modelsByProvider = aiModelRegistry.getGroupedModelsByProvider();
+
+  const activeModel =
+    forRole === "autocomplete"
+      ? configModels?.autocomplete_model
+      : forRole === "chat"
+        ? configModels?.chat_model
+        : forRole === "edit"
+          ? configModels?.edit_model
+          : undefined;
+
+  const iconSizeClass = iconSize === "medium" ? "h-4 w-4" : "h-3 w-3";
+
+  const renderModelWithRole = (modelId: AiModelId, role: Role) => {
+    const maybeModelMatch = aiModelRegistry.getModel(modelId.id);
+
+    return (
+      <div className="flex items-center gap-2 w-full px-2 py-1">
+        <AiProviderIcon
+          provider={modelId.providerId}
+          className={iconSizeClass}
+        />
+        <div className="flex flex-col">
+          <span>{maybeModelMatch?.name || modelId.shortModelId}</span>
+          <span className="text-xs text-muted-foreground">{modelId.id}</span>
+        </div>
+
+        <div className="ml-auto flex gap-1">
+          <span
+            key={role}
+            className={`text-xs px-1.5 py-0.5 rounded font-medium ${getTagColour(role)}`}
+          >
+            {role}
+          </span>
+        </div>
+      </div>
+    );
+  };
+
+  return (
+    <DropdownMenu>
+      <DropdownMenuTrigger
+        className={`flex items-center justify-between px-2 py-0.5 border rounded-md
+            hover:bg-accent hover:text-accent-foreground ${triggerClassName}`}
+      >
+        <div className="flex items-center gap-2">
+          {currentValue ? (
+            <>
+              <AiProviderIcon
+                provider={currentValue.providerId}
+                className={iconSizeClass}
+              />
+              <span className="truncate">
+                {isKnownAIProvider(currentValue.providerId)
+                  ? currentValue.shortModelId
+                  : currentValue.id}
+              </span>
+            </>
+          ) : (
+            <span className="text-muted-foreground truncate">
+              {placeholder}
+            </span>
+          )}
+        </div>
+        <ChevronDownIcon className={`${iconSizeClass} ml-1`} />
+      </DropdownMenuTrigger>
+
+      <DropdownMenuContent className="w-[300px]">
+        {activeModel &&
+          forRole &&
+          renderModelWithRole(AiModelId.parse(activeModel), forRole)}
+        {activeModel && forRole && <DropdownMenuSeparator />}
+
+        {[...modelsByProvider.entries()].map(([provider, models]) => (
+          <ProviderDropdownContent
+            key={provider}
+            provider={provider}
+            onSelect={selectModel}
+            models={models}
+            iconSizeClass={iconSizeClass}
+          />
+        ))}
+
+        {customDropdownContent}
+
+        {showAddCustomModelDocs && (
+          <>
+            <DropdownMenuSeparator />
+            <DropdownMenuItem className="flex items-center gap-2">
+              <a
+                className="flex items-center gap-1"
+                href="https://docs.marimo.io/guides/editor_features/ai_completion/?h=models#other-ai-providers"
+                target="_blank"
+                rel="noreferrer"
+              >
+                <CircleHelpIcon className="h-3 w-3" />
+                <span>How to add a custom model</span>
+              </a>
+            </DropdownMenuItem>
+          </>
+        )}
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+};
+
+const ProviderDropdownContent = ({
+  provider,
+  onSelect,
+  models,
+  customModelIcon,
+  iconSizeClass,
+}: {
+  provider: ProviderId;
+  onSelect: (modelId: QualifiedModelId) => void;
+  models: AiModel[];
+  customModelIcon?: React.ReactNode;
+  iconSizeClass: string;
+}) => {
+  const iconProvider = isKnownAIProvider(provider)
+    ? provider
+    : "openai-compatible";
+
+  const maybeProviderInfo = AiModelRegistry.getProviderInfo(provider);
+
+  if (models.length === 0) {
+    return null;
+  }
+
+  return (
+    <DropdownMenuSub>
+      <DropdownMenuSubTrigger>
+        <p className="flex items-center gap-2">
+          <AiProviderIcon provider={iconProvider} className={iconSizeClass} />
+          {getProviderLabel(provider)}
+        </p>
+      </DropdownMenuSubTrigger>
+      <DropdownMenuPortal>
+        <DropdownMenuSubContent
+          className="max-h-[40vh] overflow-y-auto"
+          alignOffset={-90}
+        >
+          {maybeProviderInfo && (
+            <>
+              <p className="text-sm text-muted-foreground p-2 max-w-[300px]">
+                {maybeProviderInfo.description}
+                <br />
+              </p>
+
+              <p className="text-sm text-muted-foreground p-2 pt-0">
+                You can find more information about this provider{" "}
+                <a
+                  href={maybeProviderInfo.url}
+                  target="_blank"
+                  className="underline"
+                  rel="noreferrer"
+                >
+                  here
+                </a>
+                .
+              </p>
+              <DropdownMenuSeparator />
+            </>
+          )}
+          {models.map((model) => {
+            const qualifiedModelId =
+              `${provider}/${model.model}` as QualifiedModelId;
+            return (
+              <DropdownMenuItem
+                key={qualifiedModelId}
+                className="flex items-center gap-2"
+                onSelect={(e) => {
+                  onSelect(qualifiedModelId);
+                }}
+                onClick={(e) => {
+                  e.preventDefault();
+                  onSelect(qualifiedModelId);
+                }}
+              >
+                <AiProviderIcon provider={iconProvider} className="h-4 w-4" />
+                <div className="pl-1 flex flex-col">
+                  <span>{model.name}</span>
+                  <span className="text-xs text-muted-foreground">
+                    {model.model}
+                  </span>
+                </div>
+                {model.custom && customModelIcon}
+              </DropdownMenuItem>
+            );
+          })}
+        </DropdownMenuSubContent>
+      </DropdownMenuPortal>
+    </DropdownMenuSub>
+  );
+};
+
+function getProviderLabel(provider: ProviderId): string {
+  const providerInfo = AiModelRegistry.getProviderInfo(provider);
+  if (providerInfo) {
+    return providerInfo.name;
+  }
+  return capitalize(provider);
+}
+
+function getTagColour(role: Role): string {
+  switch (role) {
+    case "chat":
+      return "bg-[var(--purple-3)] text-[var(--purple-11)]";
+    case "autocomplete":
+      return "bg-[var(--green-3)] text-[var(--green-11)]";
+    case "edit":
+      return "bg-[var(--blue-3)] text-[var(--blue-11)]";
+  }
+  return "bg-[var(--mauve-3)] text-[var(--mauve-11)]";
+}

package/src/components/ai/ai-provider-icon.tsx

@@ -3,20 +3,23 @@
 import AnthropicIcon from "@marimo-team/llm-info/icons/anthropic.svg?inline";
 import BedrockIcon from "@marimo-team/llm-info/icons/aws.svg?inline";
 import AzureIcon from "@marimo-team/llm-info/icons/azure.svg?inline";
-import GoogleIcon from "@marimo-team/llm-info/icons/google.svg?inline";
+import DeepseekIcon from "@marimo-team/llm-info/icons/deepseek.svg?inline";
+import GeminiIcon from "@marimo-team/llm-info/icons/googlegemini.svg?inline";
 import OllamaIcon from "@marimo-team/llm-info/icons/ollama.svg?inline";
 import OpenAIIcon from "@marimo-team/llm-info/icons/openai.svg?inline";
 import { BotIcon } from "lucide-react";
 import * as React from "react";
+import type { ProviderId } from "@/core/ai/ids/ids";
 import { cn } from "@/utils/cn";
 
-const icons = {
+const icons: Record<ProviderId, string> = {
   openai: OpenAIIcon,
   anthropic: AnthropicIcon,
-  google: GoogleIcon,
+  google: GeminiIcon,
   ollama: OllamaIcon,
   azure: AzureIcon,
   bedrock: BedrockIcon,
+  deepseek: DeepseekIcon,
 };
 
 export interface AiProviderIconProps
@@ -30,7 +33,7 @@ export const AiProviderIcon: React.FC<AiProviderIconProps> = ({
   className = "",
   ...props
 }) => {
-  if (provider === "openai-compatible") {
+  if (provider === "openai-compatible" || !(provider in icons)) {
     return <BotIcon className={cn("h-4 w-4", className)} />;
   }
 

package/src/components/app-config/ai-config.tsx

@@ -1,10 +1,12 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import React, { useId } from "react";
+import { InfoIcon } from "lucide-react";
+import React from "react";
 import type { FieldPath, UseFormReturn } from "react-hook-form";
 import {
   FormControl,
   FormDescription,
+  FormErrorsBanner,
   FormField,
   FormItem,
   FormLabel,
@@ -14,9 +16,12 @@ import { Input } from "@/components/ui/input";
 import { Kbd } from "@/components/ui/kbd";
 import { NativeSelect } from "@/components/ui/native-select";
 import { Textarea } from "@/components/ui/textarea";
+import type { QualifiedModelId } from "@/core/ai/ids/ids";
 import { CopilotConfig } from "@/core/codemirror/copilot/copilot-config";
 import { DEFAULT_AI_MODEL, type UserConfig } from "@/core/config/config-schema";
 import { isWasm } from "@/core/wasm/utils";
+import { Events } from "@/utils/events";
+import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import {
   AiProviderIcon,
   type AiProviderIconProps,
@@ -27,10 +32,12 @@ import {
   AccordionItem,
   AccordionTrigger,
 } from "../ui/accordion";
+import { DropdownMenuSeparator } from "../ui/dropdown-menu";
 import { ExternalLink } from "../ui/links";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "../ui/tabs";
+import { Tooltip } from "../ui/tooltip";
 import { SettingSubtitle } from "./common";
-import { AWS_REGIONS, KNOWN_AI_MODELS } from "./constants";
+import { AWS_REGIONS } from "./constants";
 import { IncorrectModelId } from "./incorrect-model-id";
 import { IsOverridden } from "./is-overridden";
 
@@ -182,6 +189,7 @@ interface ModelSelectorProps {
   description?: React.ReactNode;
   disabled?: boolean;
   label: string;
+  onSubmit: (values: UserConfig) => Promise<void>;
 }
 
 export const ModelSelector: React.FC<ModelSelectorProps> = ({
@@ -193,42 +201,72 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
   description,
   disabled = false,
   label,
+  onSubmit,
 }) => {
-  const modelInputId = useId();
-
   return (
     <FormField
       control={form.control}
       name={name}
       disabled={disabled}
-      render={({ field }) => (
-        <div className="flex flex-col space-y-1">
+      render={({ field }) => {
+        const value = asStringOrUndefined(field.value);
+
+        const selectModel = (modelId: QualifiedModelId) => {
+          field.onChange(modelId);
+          onSubmit(form.getValues());
+        };
+
+        const renderFormItem = () => (
           <FormItem className={formItemClasses}>
             <FormLabel>{label}</FormLabel>
             <FormControl>
-              <Input
-                list={modelInputId}
-                data-testid={testId}
-                className="m-0 inline-flex"
+              <AIModelDropdown
+                value={value}
                 placeholder={placeholder}
-                {...field}
-                value={asStringOrUndefined(field.value)}
+                onSelect={selectModel}
+                triggerClassName="text-sm"
+                customDropdownContent={
+                  <>
+                    <DropdownMenuSeparator />
+                    <p className="px-2 py-1.5 text-sm text-muted-secondary flex items-center gap-1">
+                      Enter a custom model
+                      <Tooltip content="Models should include the provider prefix, e.g. 'openai/gpt-4o'">
+                        <InfoIcon className="h-3 w-3" />
+                      </Tooltip>
+                    </p>
+                    <div className="px-2 py-1">
+                      <Input
+                        data-testid={testId}
+                        className="w-full border-border shadow-none focus-visible:shadow-xs"
+                        placeholder={placeholder}
+                        {...field}
+                        value={asStringOrUndefined(field.value)}
+                        onKeyDown={Events.stopPropagation()}
+                      />
+                      {value && (
+                        <IncorrectModelId
+                          value={value}
+                          includeSuggestion={false}
+                        />
+                      )}
+                    </div>
+                  </>
+                }
               />
             </FormControl>
             <FormMessage />
-            <IsOverridden userConfig={config} name={name} />
           </FormItem>
-          <datalist id={modelInputId}>
-            {KNOWN_AI_MODELS.map((model) => (
-              <option value={model} key={model}>
-                {model}
-              </option>
-            ))}
-          </datalist>
-          <IncorrectModelId value={asStringOrUndefined(field.value)} />
-          {description && <FormDescription>{description}</FormDescription>}
-        </div>
-      )}
+        );
+
+        return (
+          <div className="flex flex-col space-y-1">
+            {renderFormItem()}
+            <IsOverridden userConfig={config} name={name} />
+            <IncorrectModelId value={value} />
+            {description && <FormDescription>{description}</FormDescription>}
+          </div>
+        );
+      }}
     />
   );
 };
@@ -297,6 +335,7 @@ export const ProviderSelect: React.FC<ProviderSelectProps> = ({
 
 const renderCopilotProvider = (
   form: UseFormReturn<UserConfig>,
+  onSubmit: (values: UserConfig) => Promise<void>,
   config: UserConfig,
 ) => {
   const copilot = form.getValues("completion.copilot");
@@ -331,27 +370,16 @@ const renderCopilotProvider = (
 
   if (copilot === "custom") {
     return (
-      <>
-        <p className="text-sm text-muted-secondary">
-          Configure your custom AI completion provider with the following
-          settings.
-        </p>
-        <ModelSelector
-          label="Autocomplete Model"
-          form={form}
-          config={config}
-          name="ai.models.autocomplete_model"
-          placeholder="ollama/qwen2.5-coder:1.5b"
-          testId="custom-model-input"
-          description={
-            <>
-              Model to use for code completion when using a custom provider.
-              Models should include the provider name and model name separated
-              by a slash.
-            </>
-          }
-        />
-      </>
+      <ModelSelector
+        label="Autocomplete Model"
+        form={form}
+        config={config}
+        name="ai.models.autocomplete_model"
+        placeholder="ollama/qwen2.5-coder:1.5b"
+        testId="custom-model-input"
+        description="Model to use for code completion when using a custom provider."
+        onSubmit={onSubmit}
+      />
     );
   }
 };
@@ -363,6 +391,7 @@ const SettingGroup = ({ children }: { children: React.ReactNode }) => {
 export const AiCodeCompletionConfig: React.FC<AiConfigProps> = ({
   form,
   config,
+  onSubmit,
 }) => {
   return (
     <SettingGroup>
@@ -380,7 +409,7 @@ export const AiCodeCompletionConfig: React.FC<AiConfigProps> = ({
         testId="copilot-select"
       />
 
-      {renderCopilotProvider(form, config)}
+      {renderCopilotProvider(form, onSubmit, config)}
     </SettingGroup>
   );
 };
@@ -645,17 +674,17 @@ export const AiProvidersConfig: React.FC<AiConfigProps> = ({
   );
 };
 
-export const AiAssistConfig: React.FC<AiConfigProps> = ({ form, config }) => {
+export const AiAssistConfig: React.FC<AiConfigProps> = ({
+  form,
+  config,
+  onSubmit,
+}) => {
   const isWasmRuntime = isWasm();
 
   return (
     <SettingGroup>
       <SettingSubtitle>AI Assistant</SettingSubtitle>
-      <p className="text-sm text-muted-secondary">
-        Use the Chat panel to talk to your codebase, or make edits using the{" "}
-        <Kbd className="inline">Generate with AI</Kbd> button.
-      </p>
-
+      <FormErrorsBanner />
       <ModelSelector
         label="Chat Model"
         form={form}
@@ -665,21 +694,10 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({ form, config }) => {
         testId="ai-chat-model-input"
         disabled={isWasmRuntime}
         description={
-          <>
-            <p>
-              Model to use for chat conversations in the Chat panel. Models
-              should include the provider name and model name separated by a
-              slash. For example, "anthropic/claude-3-5-sonnet-latest" or
-              "google/gemini-2.0-flash-exp".
-            </p>
-            <p className="pt-1">
-              Depending on the provider, we will use the respective API key and
-              additional configuration.
-            </p>
-          </>
+          <span>Model to use for chat conversations in the Chat panel.</span>
         }
+        onSubmit={onSubmit}
       />
-
       <ModelSelector
         label="Edit Model"
         form={form}
@@ -689,20 +707,26 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({ form, config }) => {
         testId="ai-edit-model-input"
         disabled={isWasmRuntime}
         description={
-          <>
-            <p>
-              Model to use for code editing with the{" "}
-              <Kbd className="inline">Generate with AI</Kbd> button. Models
-              should include the provider name and model name separated by a
-              slash.
-            </p>
-            <p className="pt-1">
-              You can use a faster, cheaper model for edits if desired.
-            </p>
-          </>
+          <span>
+            Model to use for code editing with the{" "}
+            <Kbd className="inline">Generate with AI</Kbd> button.
+          </span>
         }
+        onSubmit={onSubmit}
       />
 
+      <ul className="bg-muted p-2 rounded-md list-disc space-y-1 pl-6">
+        <li className="text-xs text-muted-secondary">
+          Models should include the provider name and model name separated by a
+          slash. For example, "anthropic/claude-3-5-sonnet-latest" or
+          "google/gemini-2.0-flash-exp"
+        </li>
+        <li className="text-xs text-muted-secondary">
+          Depending on the provider, we will use the respective API key and
+          additional configuration.
+        </li>
+      </ul>
+
       <FormField
         control={form.control}
         name="ai.rules"

package/src/components/app-config/app-config-button.tsx

@@ -2,6 +2,7 @@
 
 import { useAtom } from "jotai";
 import { SettingsIcon } from "lucide-react";
+import { VisuallyHidden } from "react-aria";
 import { AppConfigForm } from "@/components/app-config/app-config-form";
 import {
   Popover,
@@ -10,7 +11,12 @@ import {
 } from "@/components/ui/popover";
 import { Button as EditorButton } from "../editor/inputs/Inputs";
 import { Button } from "../ui/button";
-import { Dialog, DialogContent, DialogTrigger } from "../ui/dialog";
+import {
+  Dialog,
+  DialogContent,
+  DialogTitle,
+  DialogTrigger,
+} from "../ui/dialog";
 import { Tooltip } from "../ui/tooltip";
 import { settingDialogAtom } from "./state";
 import { UserConfigForm } from "./user-config-form";
@@ -46,6 +52,9 @@ export const ConfigButton: React.FC<Props> = ({
 
   const userSettingsDialog = (
     <DialogContent className="w-[80vw] h-[70vh] overflow-hidden sm:max-w-5xl top-[15vh] p-0">
+      <VisuallyHidden>
+        <DialogTitle>User settings</DialogTitle>
+      </VisuallyHidden>
       <UserConfigForm />
     </DialogContent>
   );

package/src/components/app-config/constants.ts

@@ -1,38 +1,4 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-export const KNOWN_AI_MODELS = [
-  // Anthropic
-  "anthropic/claude-opus-4-1-20250805",
-  "anthropic/claude-opus-4-20250514",
-  "anthropic/claude-sonnet-4-20250514",
-  "anthropic/claude-3-7-sonnet-latest",
-  "anthropic/claude-3-5-sonnet-latest",
-  "anthropic/claude-3-5-haiku-latest",
-
-  // DeepSeek
-  "deepseek/deepseek-v3",
-  "deepseek/deepseek-r1",
-
-  // Google
-  "google/gemini-2.5-flash-preview-05-20",
-  "google/gemini-2.5-pro-preview-06-05",
-  "google/gemini-2.0-flash",
-  "google/gemini-2.0-flash-lite",
-
-  // OpenAI
-  "openai/o3",
-  "openai/o4-mini",
-  "openai/gpt-4.5-preview",
-  "openai/gpt-4.1",
-  "openai/gpt-4o",
-  "openai/gpt-3.5-turbo",
-
-  // AWS Bedrock Models
-  "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
-  "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
-  "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
-  "bedrock/meta.llama3-3-70b-instruct-v1:0",
-  "bedrock/cohere.command-r-plus-v1",
-] as const;
 
 /**
  * AWS regions where the Bedrock service is available