Refactor AI Generation with AI SDK v5

app/ai/components/ai-chat-form.tsx

@@ -2,23 +2,27 @@
 
 import { AIChatFormBody } from "@/components/editor/ai/ai-chat-form-body";
 import { AlertBanner } from "@/components/editor/ai/alert-banner";
+import { EnhancePromptButton } from "@/components/editor/ai/enhance-prompt-button";
 import { ImageUploader } from "@/components/editor/ai/image-uploader";
 import ThemePresetSelect from "@/components/editor/theme-preset-select";
 import { Button } from "@/components/ui/button";
 import { useAIChatForm } from "@/hooks/use-ai-chat-form";
-import { useAIThemeGenerationCore } from "@/hooks/use-ai-theme-generation-core";
+import { useAIEnhancePrompt } from "@/hooks/use-ai-enhance-prompt";
 import { useGuards } from "@/hooks/use-guards";
 import { MAX_IMAGE_FILES } from "@/lib/constants";
 import { cn } from "@/lib/utils";
 import { AIPromptData } from "@/types/ai";
 import { ArrowUp, Loader, StopCircle } from "lucide-react";
 
 export function AIChatForm({
-  handleThemeGeneration,
+  onThemeGeneration,
+  isGeneratingTheme,
+  onCancelThemeGeneration,
 }: {
-  handleThemeGeneration: (promptData: AIPromptData | null) => void;
+  onThemeGeneration: (promptData: AIPromptData) => void;
+  isGeneratingTheme: boolean;
+  onCancelThemeGeneration: () => void;
 }) {
-  const { loading: aiGenerateLoading, cancelThemeGeneration } = useAIThemeGenerationCore();
   const {
     editorContentDraft,
     handleContentChange,
@@ -31,10 +35,23 @@ export function AIChatForm({
     handleImageRemove,
     isSomeImageUploading,
     isUserDragging,
+    isInitializing,
   } = useAIChatForm();
 
   const { checkValidSession, checkValidSubscription } = useGuards();
 
+  const { startEnhance, stopEnhance, enhancedPromptAsJsonContent, isEnhancingPrompt } =
+    useAIEnhancePrompt();
+
+  const handleEnhancePrompt = () => {
+    // TODO: Add subscription check, this should be a Pro only feature
+    if (!checkValidSession() || !checkValidSubscription()) return; // Act as an early return;
+
+    // Only send images that are not loading, and strip loading property
+    const images = uploadedImages.filter((img) => !img.loading).map(({ url }) => ({ url }));
+    startEnhance({ ...promptData, images });
+  };
+
   const handleGenerate = async () => {
     if (!checkValidSession() || !checkValidSubscription()) return; // Act as an early return
 
@@ -44,7 +61,7 @@ export function AIChatForm({
     // Proceed only if there is text, or at least one image
     if (isEmptyPrompt && images.length === 0) return;
 
-    handleThemeGeneration({
+    onThemeGeneration({
       ...promptData,
       content: promptData?.content ?? "",
       mentions: promptData?.mentions ?? [],
@@ -61,20 +78,28 @@ export function AIChatForm({
       <div className="bg-background relative z-10 flex size-full min-h-[100px] flex-1 flex-col gap-2 overflow-hidden rounded-lg border p-2 shadow-xs">
         <AIChatFormBody
           isUserDragging={isUserDragging}
-          aiGenerateLoading={aiGenerateLoading}
+          disabled={isEnhancingPrompt}
+          canSubmit={
+            !isGeneratingTheme &&
+            !isEnhancingPrompt &&
+            !isEmptyPrompt &&
+            !isSomeImageUploading &&
+            !isInitializing
+          }
           uploadedImages={uploadedImages}
           handleImagesUpload={handleImagesUpload}
           handleImageRemove={handleImageRemove}
           handleContentChange={handleContentChange}
           handleGenerate={handleGenerate}
           initialEditorContent={editorContentDraft ?? undefined}
           textareaKey={editorContentDraft ? "with-draft" : "no-draft"}
+          externalEditorContent={isEnhancingPrompt ? enhancedPromptAsJsonContent : undefined}
         />
 
         <div className="flex items-center justify-between gap-2">
           <div className="flex w-full max-w-64 items-center gap-2 overflow-hidden">
             <ThemePresetSelect
-              disabled={aiGenerateLoading}
+              disabled={isGeneratingTheme || isEnhancingPrompt || isInitializing}
               withCycleThemes={false}
               variant="outline"
               size="sm"
@@ -83,22 +108,34 @@ export function AIChatForm({
           </div>
 
           <div className="flex items-center gap-2">
+            {/* TODO: This should be a Pro only feature */}
+            {promptData?.content ? (
+              <EnhancePromptButton
+                isEnhancing={isEnhancingPrompt}
+                onStart={handleEnhancePrompt}
+                onStop={stopEnhance}
+                disabled={isGeneratingTheme || isInitializing}
+              />
+            ) : null}
+
             <ImageUploader
               fileInputRef={fileInputRef}
               onImagesUpload={handleImagesUpload}
               onClick={() => fileInputRef.current?.click()}
               disabled={
-                aiGenerateLoading ||
+                isGeneratingTheme ||
+                isEnhancingPrompt ||
+                isInitializing ||
                 uploadedImages.some((img) => img.loading) ||
                 uploadedImages.length >= MAX_IMAGE_FILES
               }
             />
 
-            {aiGenerateLoading ? (
+            {isGeneratingTheme ? (
               <Button
                 variant="destructive"
                 size="sm"
-                onClick={cancelThemeGeneration}
+                onClick={onCancelThemeGeneration}
                 className={cn("flex items-center gap-1", "@max-[350px]/form:w-8")}
               >
                 <StopCircle />
@@ -107,11 +144,17 @@ export function AIChatForm({
             ) : (
               <Button
                 size="icon"
-                className="size-8"
+                className="size-8 shadow-none"
                 onClick={handleGenerate}
-                disabled={isEmptyPrompt || isSomeImageUploading || aiGenerateLoading}
+                disabled={
+                  isEmptyPrompt ||
+                  isSomeImageUploading ||
+                  isGeneratingTheme ||
+                  isEnhancingPrompt ||
+                  isInitializing
+                }
               >
-                {aiGenerateLoading ? <Loader className="animate-spin" /> : <ArrowUp />}
+                {isGeneratingTheme ? <Loader className="animate-spin" /> : <ArrowUp />}
               </Button>
             )}
           </div>

app/ai/components/ai-chat-hero.tsx

@@ -1,10 +1,10 @@
 "use client";
 
 import { HorizontalScrollArea } from "@/components/horizontal-scroll-area";
-import { useAIGenerateTheme } from "@/hooks/use-ai-generate-theme";
+import { useChatContext } from "@/hooks/use-chat-context";
+import { useAIThemeGenerationCore } from "@/hooks/use-ai-theme-generation-core";
 import { useGuards } from "@/hooks/use-guards";
 import { usePostLoginAction } from "@/hooks/use-post-login-action";
-import { useAIChatStore } from "@/store/ai-chat-store";
 import { usePreferencesStore } from "@/store/preferences-store";
 import { AIPromptData } from "@/types/ai";
 import { useRouter } from "next/navigation";
@@ -13,23 +13,23 @@ import { ChatHeading } from "./chat-heading";
 import { SuggestedPillActions } from "./suggested-pill-actions";
 
 export function AIChatHero() {
-  const { generateTheme } = useAIGenerateTheme();
+  const { startNewChat } = useChatContext();
+  const { generateThemeCore, isGeneratingTheme, cancelThemeGeneration } =
+    useAIThemeGenerationCore();
   const { checkValidSession, checkValidSubscription } = useGuards();
   const router = useRouter();
 
-  const { clearMessages } = useAIChatStore();
   const { setChatSuggestionsOpen } = usePreferencesStore();
 
-  const handleRedirectAndThemeGeneration = async (promptData: AIPromptData | null) => {
+  const handleRedirectAndThemeGeneration = (promptData: AIPromptData) => {
     if (!checkValidSession("signup", "AI_GENERATE_FROM_PAGE", { promptData })) return;
     if (!checkValidSubscription()) return;
 
-    // Clear the messages and open the chat suggestions when the user starts a chat from the '/ai' page
-    clearMessages();
+    startNewChat();
     setChatSuggestionsOpen(true);
 
     router.push("/editor/theme?tab=ai");
-    generateTheme(promptData);
+    generateThemeCore(promptData);
   };
 
   usePostLoginAction("AI_GENERATE_FROM_PAGE", ({ promptData }) => {
@@ -39,17 +39,24 @@ export function AIChatHero() {
   return (
     <div className="relative isolate flex w-full flex-1">
       <div className="@container relative isolate z-1 mx-auto flex max-w-[49rem] flex-1 flex-col justify-center px-4">
-        <ChatHeading />
+        <ChatHeading isGeneratingTheme={isGeneratingTheme} />
 
         {/* Chat form input and suggestions */}
         <div className="relative mx-auto flex w-full flex-col gap-2">
           <div className="relative isolate z-10 w-full">
-            <AIChatForm handleThemeGeneration={handleRedirectAndThemeGeneration} />
+            <AIChatForm
+              onThemeGeneration={handleRedirectAndThemeGeneration}
+              isGeneratingTheme={isGeneratingTheme}
+              onCancelThemeGeneration={cancelThemeGeneration}
+            />
           </div>
 
           {/* Quick suggestions */}
           <HorizontalScrollArea className="mx-auto py-2">
-            <SuggestedPillActions handleThemeGeneration={handleRedirectAndThemeGeneration} />
+            <SuggestedPillActions
+              onThemeGeneration={handleRedirectAndThemeGeneration}
+              isGeneratingTheme={isGeneratingTheme}
+            />
           </HorizontalScrollArea>
         </div>
       </div>

app/ai/components/chat-heading.tsx

@@ -1,14 +1,10 @@
-import { useAIThemeGenerationCore } from "@/hooks/use-ai-theme-generation-core";
-
-export function ChatHeading() {
-  const { loading: aiGenerateLoading } = useAIThemeGenerationCore();
-
+export function ChatHeading({ isGeneratingTheme }: { isGeneratingTheme: boolean }) {
   return (
     <h1
       style={
         {
-          "--gradient-accent": aiGenerateLoading ? "var(--foreground)" : "var(--foreground)",
-          "--gradient-base": aiGenerateLoading ? "var(--muted-foreground)" : "var(--foreground)",
+          "--gradient-accent": isGeneratingTheme ? "var(--foreground)" : "var(--foreground)",
+          "--gradient-base": isGeneratingTheme ? "var(--muted-foreground)" : "var(--foreground)",
         } as React.CSSProperties
       }
       className="animate-text bg-gradient-to-r from-(--gradient-base) via-(--gradient-accent) to-(--gradient-base) bg-[200%_auto] bg-clip-text pb-4 text-center text-[clamp(24px,7cqw,46px)] font-semibold tracking-tighter text-pretty text-transparent"

app/ai/components/suggested-pill-actions.tsx

@@ -1,7 +1,6 @@
 "use client";
 
 import { PillActionButton } from "@/components/editor/ai/pill-action-button";
-import { useAIThemeGenerationCore } from "@/hooks/use-ai-theme-generation-core";
 import { useImageUpload } from "@/hooks/use-image-upload";
 import { imageUploadReducer } from "@/hooks/use-image-upload-reducer";
 import { MAX_IMAGE_FILE_SIZE } from "@/lib/constants";
@@ -12,12 +11,12 @@ import { ImageIcon, Sparkles } from "lucide-react";
 import { useEffect, useReducer } from "react";
 
 export function SuggestedPillActions({
-  handleThemeGeneration,
+  onThemeGeneration,
+  isGeneratingTheme,
 }: {
-  handleThemeGeneration: (promptData: AIPromptData | null) => void;
+  onThemeGeneration: (promptData: AIPromptData) => void;
+  isGeneratingTheme: boolean;
 }) {
-  const { loading: aiIsGenerating } = useAIThemeGenerationCore();
-
   const [uploadedImages, dispatch] = useReducer(imageUploadReducer, []);
 
   const { fileInputRef, handleImagesUpload, canUploadMore, isSomeImageUploading } = useImageUpload({
@@ -30,7 +29,7 @@ export function SuggestedPillActions({
   // Automatically send prompt when an image is selected and loaded
   useEffect(() => {
     if (uploadedImages.length > 0 && !isSomeImageUploading) {
-      handleThemeGeneration({
+      onThemeGeneration({
         content: "", // No text prompt
         mentions: [], // No mentions
         images: [uploadedImages[0]],
@@ -42,7 +41,7 @@ export function SuggestedPillActions({
 
   const handleSetPrompt = async (prompt: string) => {
     const promptData = createCurrentThemePrompt({ prompt });
-    handleThemeGeneration(promptData);
+    onThemeGeneration(promptData);
   };
 
   const handleImageButtonClick = () => {
@@ -61,14 +60,14 @@ export function SuggestedPillActions({
 
   return (
     <>
-      <PillActionButton onClick={handleImageButtonClick} disabled={aiIsGenerating}>
+      <PillActionButton onClick={handleImageButtonClick} disabled={isGeneratingTheme}>
         <input
           type="file"
           accept="image/*"
           multiple={false}
           ref={fileInputRef}
           onChange={handleImageUpload}
-          disabled={aiIsGenerating}
+          disabled={isGeneratingTheme}
           style={{ display: "none" }}
         />
         <ImageIcon /> From an Image
@@ -78,7 +77,7 @@ export function SuggestedPillActions({
         <PillActionButton
           key={key}
           onClick={() => handleSetPrompt(prompt)}
-          disabled={aiIsGenerating}
+          disabled={isGeneratingTheme}
         >
           <Sparkles /> {label}
         </PillActionButton>

app/ai/page.tsx

@@ -1,4 +1,3 @@
-import { Metadata } from "next";
 import { AIAnnouncement } from "./components/ai-announcement";
 import { AIChatHero } from "./components/ai-chat-hero";
 

app/api/enhance-prompt/route.ts

@@ -0,0 +1,68 @@
+import { AIPromptData } from "@/types/ai";
+import { buildUserContentPartsFromPromptData } from "@/utils/ai/message-converter";
+import { baseProviderOptions, MODELS } from "@/utils/ai/providers";
+import { smoothStream, streamText } from "ai";
+
+export async function POST(req: Request) {
+  // TODO: Add session and subscription check, this should be a Pro only feature
+  // TODO: Record AI usage, providing the model id to `recordAIUsage` function
+
+  const body = await req.json();
+  const { prompt: _prompt, promptData }: { prompt: string; promptData: AIPromptData } = body;
+  const userContentParts = buildUserContentPartsFromPromptData(promptData);
+
+  const result = streamText({
+    system: ENHANCE_PROMPT_SYSTEM,
+    messages: [
+      {
+        role: "user",
+        content: userContentParts,
+      },
+    ],
+    model: MODELS.promptEnhancement,
+    providerOptions: baseProviderOptions,
+    experimental_transform: smoothStream({
+      delayInMs: 10,
+      chunking: "word",
+    }),
+  });
+
+  return result.toUIMessageStreamResponse();
+}
+
+const ENHANCE_PROMPT_SYSTEM = `# Role
+You are a prompt refinement specialist for a shadcn/ui theme generator. Rewrite the user's input into a precise, ready-to-use prompt that preserves the original intent, language, and tone. Write as the requester of the theme, not as an assistant.
+
+# Core principles
+- Language matching: respond in the exact same language as the user
+- Cultural context: respect regional expressions, slang, and cultural references
+- Length limit: output must NOT exceed 500 characters
+- If when analyzing the user's prompt you recognize a vibe, style, or visual reference, include it in the output.
+
+# Mentions
+- User input may include mentions like @Current Theme or @PresetName. Mentions are always in the format of @[label].
+- Mentions are predefined styles that are intended to be used as the base or reference for the theme.
+- Preserve them verbatim in the output text (same labels, same order if possible).
+- Do not invent new mentions. Only keep and reposition mentions that appear in the user's prompt or in the provided mention list.
+- Avoid repeating the same mention multiple times.
+
+# Enhancement guidelines
+- If the prompt is vague, add concrete visual details (colors, mood, typography, style references)
+- If it mentions a brand or style, include relevant design characteristics
+- If it's already detailed, clarify ambiguous parts and tighten wording
+- Preserve any specific requests (colors, fonts, mood, etc.)
+- Add context that helps the theme generator understand the desired aesthetic
+
+# Output rules
+- Output a single line of plain text suitable to paste into the Generate Theme tool
+- No greetings or meta commentary; do not address the user
+- Do not narrate with phrases like "I'm seeing", "let's", "alright", or "what you want is"
+- No bullets, no quotes, no markdown, no JSON
+
+# What NOT to do
+- Don't change the user's core request
+- Don't add conflicting style directions
+- Don't exceed 500 characters
+- Don't use a different language than the user
+- Do not list the mentions' properties in the enhanced prompt.
+- Do not use em dashes (—)`;