integrate ai gateway and improve ui components (#1145)

components/multimodal-input.tsx

@@ -24,7 +24,14 @@ import {
   PromptInputToolbar,
   PromptInputTools,
   PromptInputSubmit,
+  PromptInputModelSelect,
+  PromptInputModelSelectTrigger,
+  PromptInputModelSelectContent,
 } from './elements/prompt-input';
+import {
+  SelectItem,
+  SelectValue,
+} from '@/components/ui/select';
 import equal from 'fast-deep-equal';
 import type { UseChatHelpers } from '@ai-sdk/react';
 import { AnimatePresence, motion } from 'framer-motion';
@@ -32,6 +39,9 @@ import { ArrowDown } from 'lucide-react';
 import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom';
 import type { VisibilityType } from './visibility-selector';
 import type { Attachment, ChatMessage } from '@/lib/types';
+import { chatModels } from '@/lib/ai/models';
+import { saveChatModelAsCookie } from '@/app/(chat)/actions';
+import { startTransition } from 'react';
 
 function PureMultimodalInput({
   chatId,
@@ -46,6 +56,7 @@ function PureMultimodalInput({
   sendMessage,
   className,
   selectedVisibilityType,
+  selectedModelId,
 }: {
   chatId: string;
   input: string;
@@ -59,6 +70,7 @@ function PureMultimodalInput({
   sendMessage: UseChatHelpers<ChatMessage>['sendMessage'];
   className?: string;
   selectedVisibilityType: VisibilityType;
+  selectedModelId: string;
 }) {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { width } = useWindowSize();
@@ -318,6 +330,7 @@ function PureMultimodalInput({
         <PromptInputToolbar className="px-2 py-1">
           <PromptInputTools className="gap-2">
             <AttachmentsButton fileInputRef={fileInputRef} status={status} />
+            <ModelSelectorCompact selectedModelId={selectedModelId} />
           </PromptInputTools>
           {status === 'submitted' ? (
             <StopButton stop={stop} setMessages={setMessages} />
@@ -343,6 +356,7 @@ export const MultimodalInput = memo(
     if (!equal(prevProps.attachments, nextProps.attachments)) return false;
     if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
       return false;
+    if (prevProps.selectedModelId !== nextProps.selectedModelId) return false;
 
     return true;
   },
@@ -373,6 +387,50 @@ function PureAttachmentsButton({
 
 const AttachmentsButton = memo(PureAttachmentsButton);
 
+function PureModelSelectorCompact({
+  selectedModelId,
+}: {
+  selectedModelId: string;
+}) {
+  const [optimisticModelId, setOptimisticModelId] = useState(selectedModelId);
+
+  const selectedModel = chatModels.find(model => model.id === optimisticModelId);
+
+  return (
+    <PromptInputModelSelect
+      value={selectedModel?.name}
+      onValueChange={(modelName) => {
+        const model = chatModels.find(m => m.name === modelName);
+        if (model) {
+          setOptimisticModelId(model.id);
+          startTransition(() => {
+            saveChatModelAsCookie(model.id);
+          });
+        }
+      }}
+    >
+      <PromptInputModelSelectTrigger 
+        type="button"
+        className="text-xs focus:outline-none focus:ring-0 focus:ring-offset-0 focus-visible:ring-0 focus-visible:ring-offset-0 data-[state=open]:ring-0 data-[state=closed]:ring-0"
+      >
+        {selectedModel?.name || "Select model"}
+      </PromptInputModelSelectTrigger>
+      <PromptInputModelSelectContent>
+        {chatModels.map((model) => (
+          <SelectItem key={model.id} value={model.name}>
+            <div className="flex flex-col items-start gap-1 py-1">
+              <div className="font-medium">{model.name}</div>
+              <div className="text-xs text-muted-foreground">{model.description}</div>
+            </div>
+          </SelectItem>
+        ))}
+      </PromptInputModelSelectContent>
+    </PromptInputModelSelect>
+  );
+}
+
+const ModelSelectorCompact = memo(PureModelSelectorCompact);
+
 function PureStopButton({
   stop,
   setMessages,