feat: dynamic model discovery from vercel ai gateway (#1353)

components/multimodal-input.tsx

@@ -1,15 +1,14 @@
 "use client";
 
 import type { UseChatHelpers } from "@ai-sdk/react";
-import { Trigger } from "@radix-ui/react-select";
 import type { UIMessage } from "ai";
 import equal from "fast-deep-equal";
+import { CheckIcon } from "lucide-react";
 import {
   type ChangeEvent,
   type Dispatch,
   memo,
   type SetStateAction,
-  startTransition,
   useCallback,
   useEffect,
   useMemo,
@@ -18,34 +17,45 @@ import {
 } from "react";
 import { toast } from "sonner";
 import { useLocalStorage, useWindowSize } from "usehooks-ts";
-import { saveChatModelAsCookie } from "@/app/(chat)/actions";
-import { SelectItem } from "@/components/ui/select";
-import { chatModels } from "@/lib/ai/models";
+import {
+  ModelSelector,
+  ModelSelectorContent,
+  ModelSelectorGroup,
+  ModelSelectorInput,
+  ModelSelectorItem,
+  ModelSelectorList,
+  ModelSelectorLogo,
+  ModelSelectorName,
+  ModelSelectorTrigger,
+} from "@/components/ai-elements/model-selector";
+import {
+  chatModels,
+  DEFAULT_CHAT_MODEL,
+  modelsByProvider,
+} from "@/lib/ai/models";
 import type { Attachment, ChatMessage } from "@/lib/types";
 import type { AppUsage } from "@/lib/usage";
 import { cn } from "@/lib/utils";
 import { Context } from "./elements/context";
 import {
   PromptInput,
-  PromptInputModelSelect,
-  PromptInputModelSelectContent,
   PromptInputSubmit,
   PromptInputTextarea,
   PromptInputToolbar,
   PromptInputTools,
 } from "./elements/prompt-input";
-import {
-  ArrowUpIcon,
-  ChevronDownIcon,
-  CpuIcon,
-  PaperclipIcon,
-  StopIcon,
-} from "./icons";
+import { ArrowUpIcon, PaperclipIcon, StopIcon } from "./icons";
 import { PreviewAttachment } from "./preview-attachment";
 import { SuggestedActions } from "./suggested-actions";
 import { Button } from "./ui/button";
 import type { VisibilityType } from "./visibility-selector";
 
+function setCookie(name: string, value: string) {
+  const maxAge = 60 * 60 * 24 * 365; // 1 year
+  // biome-ignore lint/suspicious/noDocumentCookie: needed for client-side cookie setting
+  document.cookie = `${name}=${encodeURIComponent(value)}; path=/; max-age=${maxAge}`;
+}
+
 function PureMultimodalInput({
   chatId,
   input,
@@ -367,7 +377,7 @@ function PureMultimodalInput({
           />{" "}
           <Context {...contextProps} />
         </div>
-        <PromptInputToolbar className="!border-top-0 border-t-0! p-0 shadow-none dark:border-0 dark:border-transparent!">
+        <PromptInputToolbar className="border-top-0! border-t-0! p-0 shadow-none dark:border-0 dark:border-transparent!">
           <PromptInputTools className="gap-0 sm:gap-0.5">
             <AttachmentsButton
               fileInputRef={fileInputRef}
@@ -430,7 +440,8 @@ function PureAttachmentsButton({
   status: UseChatHelpers<ChatMessage>["status"];
   selectedModelId: string;
 }) {
-  const isReasoningModel = selectedModelId === "chat-model-reasoning";
+  const isReasoningModel =
+    selectedModelId.includes("reasoning") || selectedModelId.includes("think");
 
   return (
     <Button
@@ -457,52 +468,66 @@ function PureModelSelectorCompact({
   selectedModelId: string;
   onModelChange?: (modelId: string) => void;
 }) {
-  const [optimisticModelId, setOptimisticModelId] = useState(selectedModelId);
+  const [open, setOpen] = useState(false);
 
-  useEffect(() => {
-    setOptimisticModelId(selectedModelId);
-  }, [selectedModelId]);
+  const selectedModel =
+    chatModels.find((m) => m.id === selectedModelId) ??
+    chatModels.find((m) => m.id === DEFAULT_CHAT_MODEL) ??
+    chatModels[0];
+  const [provider] = selectedModel.id.split("/");
 
-  const selectedModel = chatModels.find(
-    (model) => model.id === optimisticModelId
-  );
+  // Provider display names
+  const providerNames: Record<string, string> = {
+    anthropic: "Anthropic",
+    openai: "OpenAI",
+    google: "Google",
+    xai: "xAI",
+    reasoning: "Reasoning",
+  };
 
   return (
-    <PromptInputModelSelect
-      onValueChange={(modelName) => {
-        const model = chatModels.find((m) => m.name === modelName);
-        if (model) {
-          setOptimisticModelId(model.id);
-          onModelChange?.(model.id);
-          startTransition(() => {
-            saveChatModelAsCookie(model.id);
-          });
-        }
-      }}
-      value={selectedModel?.name}
-    >
-      <Trigger asChild>
-        <Button className="h-8 px-2" variant="ghost">
-          <CpuIcon size={16} />
-          <span className="hidden font-medium text-xs sm:block">
-            {selectedModel?.name}
-          </span>
-          <ChevronDownIcon size={16} />
+    <ModelSelector onOpenChange={setOpen} open={open}>
+      <ModelSelectorTrigger asChild>
+        <Button className="h-8 w-[200px] justify-between px-2" variant="ghost">
+          {provider && <ModelSelectorLogo provider={provider} />}
+          <ModelSelectorName>{selectedModel.name}</ModelSelectorName>
         </Button>
-      </Trigger>
-      <PromptInputModelSelectContent className="min-w-[260px] p-0">
-        <div className="flex flex-col gap-px">
-          {chatModels.map((model) => (
-            <SelectItem key={model.id} value={model.name}>
-              <div className="truncate font-medium text-xs">{model.name}</div>
-              <div className="mt-px truncate text-[10px] text-muted-foreground leading-tight">
-                {model.description}
-              </div>
-            </SelectItem>
-          ))}
-        </div>
-      </PromptInputModelSelectContent>
-    </PromptInputModelSelect>
+      </ModelSelectorTrigger>
+      <ModelSelectorContent>
+        <ModelSelectorInput placeholder="Search models..." />
+        <ModelSelectorList>
+          {Object.entries(modelsByProvider).map(
+            ([providerKey, providerModels]) => (
+              <ModelSelectorGroup
+                heading={providerNames[providerKey] ?? providerKey}
+                key={providerKey}
+              >
+                {providerModels.map((model) => {
+                  const logoProvider = model.id.split("/")[0];
+                  return (
+                    <ModelSelectorItem
+                      key={model.id}
+                      onSelect={() => {
+                        onModelChange?.(model.id);
+                        setCookie("chat-model", model.id);
+                        setOpen(false);
+                      }}
+                      value={model.id}
+                    >
+                      <ModelSelectorLogo provider={logoProvider} />
+                      <ModelSelectorName>{model.name}</ModelSelectorName>
+                      {model.id === selectedModel.id && (
+                        <CheckIcon className="ml-auto size-4" />
+                      )}
+                    </ModelSelectorItem>
+                  );
+                })}
+              </ModelSelectorGroup>
+            )
+          )}
+        </ModelSelectorList>
+      </ModelSelectorContent>
+    </ModelSelector>
   );
 }