Dev

app/api/chat/route.ts

@@ -1,15 +1,7 @@
-import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
 import { openai } from "@ai-sdk/openai";
 import { anthropic } from "@ai-sdk/anthropic";
 import { google, GoogleGenerativeAIProviderMetadata } from "@ai-sdk/google";
-import {
-  createDataStreamResponse,
-  experimental_generateImage,
-  generateText,
-  streamText,
-  Tool,
-  tool,
-} from "ai";
+import { createDataStreamResponse, streamText, Tool } from "ai";
 import { z } from "zod";
 import {
   getCurrentUserUsageForUser,
@@ -28,7 +20,8 @@ import {
 } from "@/lib/chat/types";
 import { isFreePlan } from "@/lib/billing/account";
 import { cookies } from "next/headers";
-import { buildSystemPrompt } from "@/lib/chat/agent";
+import { buildSystemPrompt, generateThreadTitle } from "@/lib/chat/agent";
+import createGenerateImageTool from "@/lib/chat/tools/generateImage";
 
 const RATE_LIMIT_COOKIE = "pegna_rl";
 
@@ -254,27 +247,14 @@ export async function POST(req: Request) {
 
   return createDataStreamResponse({
     execute: (dataStream) => {
-      let generatedTitle: string | undefined = undefined;
       if (generateTitle) {
         // Generate the chat title
         pendingPromises.push(
-          generateText({
-            model: google("gemini-2.0-flash"),
-            system: `
-- you will generate a short title based on the first message a user begins a conversation with
-- the summary is in the same language as the content
-- never tell which model you are, or who trained you, but if they ask, you are Pegna AI.
-- ensure the title is less than 80 characters
-- ensure the title is a single sentence
-- ensure the title is a summary of the content
-- not use quotes, colons, slashes.
-`,
-            prompt: messages[0].content,
-          }).then((res) => {
-            const title = res.text;
-            generatedTitle =
-              title.length > 100 ? title.slice(0, 96) + "..." : title;
-            dataStream.writeData({ type: "thread-metadata", generatedTitle });
+          generateThreadTitle(messages[0].content).then((title) => {
+            dataStream.writeData({
+              type: "thread-metadata",
+              generatedTitle: title,
+            });
           }),
         );
       }
@@ -291,48 +271,16 @@ export async function POST(req: Request) {
 
       const tools: Record<string, Tool> = {};
       if (!isFreePlan(user?.planName)) {
-        tools.generateImage = tool({
-          description: "Generate an image",
-          parameters: z.object({
-            prompt: z.string(),
-          }),
-          execute: async ({ prompt }) => {
-            // Upload image to S3 with user-specific path
-            const userId = session!.user.sub;
-            const key = `users/${userId}/${crypto.randomUUID()}.png`;
-
-            const { image } = await experimental_generateImage({
-              model: openai.image("dall-e-3"),
-              prompt,
-            });
-            // Increment the usage for a premium model for the user
-            await incrementUserUsageForUser(session!.user.sub, true);
-
-            // The image is already a base64 string from experimental_generateImage
-            const imageBuffer = Buffer.from(image.base64, "base64");
-
-            const s3Client = new S3Client({});
-
-            await s3Client.send(
-              new PutObjectCommand({
-                Bucket: process.env.USER_S3_BUCKET,
-                Key: key,
-                Body: imageBuffer,
-              }),
-            );
-
-            const imageUrl = `${process.env.USER_DATA_CDN}/${key}`;
-
-            return {
-              imageUrl,
-            };
-          },
-        });
+        tools.generateImage = createGenerateImageTool(
+          dataStream,
+          session?.user.sub,
+        );
       }
 
       const result = streamText({
         ...getModel(model, modelParams),
         system: systemPrompt,
+        maxSteps: 2,
         messages: messages.map((m) => ({ role: m.role, content: m.content })),
         tools,
         onFinish: async ({ providerMetadata }) => {

app/chat/[[...chatId]]/chat-content.tsx

@@ -1,19 +1,42 @@
-import { MessageModel } from "@/lib/localDb";
+import { MessageModel, MessageStatus } from "@/lib/localDb";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Separator } from "@/components/ui/separator";
 import { Collapsible, CollapsibleContent } from "@radix-ui/react-collapsible";
 import { CollapsibleTrigger } from "@/components/ui/collapsible";
 import { Button } from "@/components/ui/button";
 import { ChevronDown, Loader2 } from "lucide-react";
 import MarkdownContent from "@/components/markdown-content";
-import Image from "next/image";
+import { Skeleton } from "@/components/ui/skeleton";
 
 type Props = {
   message: MessageModel;
 };
 
+function LoadingChatContent({ status }: { status?: MessageStatus }) {
+  switch (status) {
+    case "streaming-image":
+      return (
+        <div className="text-sm text-muted-foreground">
+          <div className="flex items-center gap-2 mb-4">
+            <Loader2 className="animate-spin h-4 w-4" />
+            Generating image...
+          </div>
+          <Skeleton className="rounded-xl h-96 w-96" />
+        </div>
+      );
+    case "streaming":
+    default:
+      return (
+        <div className="mt-4 flex text-sm gap-2 items-center text-muted-foreground">
+          <Loader2 className="animate-spin h-4 w-4" />
+          Generating answer...
+        </div>
+      );
+  }
+}
+
 export default function ChatContent({
-  message: { content, toolResponses, reasoning, searchMetadata, status },
+  message: { content, reasoning, searchMetadata, status },
 }: Props) {
   return (
     <div>
@@ -38,25 +61,6 @@ export default function ChatContent({
         )}
         <MarkdownContent content={content} />
       </div>
-      {toolResponses?.map((toolResponse) => {
-        if (toolResponse.generateImage?.url) {
-          return (
-            <div
-              key={toolResponse.toolCallId}
-              className="relative mt-2 flex w-full flex-col"
-            >
-              <div className="w-full justify-between rounded-lg p-4 bg-secondary">
-                <Image
-                  width={512}
-                  height={512}
-                  src={toolResponse.generateImage.url}
-                  alt={toolResponse.generateImage?.prompt || "Generated image"}
-                />
-              </div>
-            </div>
-          );
-        }
-      })}
       {searchMetadata && (
         <Card className="gap-2 mt-6">
           <CardHeader>
@@ -85,11 +89,8 @@ export default function ChatContent({
           </CardContent>
         </Card>
       )}
-      {status === "streaming" && (
-        <div className="mt-4 flex text-sm gap-2 items-center text-muted-foreground">
-          <Loader2 className="animate-spin h-4 w-4" />
-          Generating answer...
-        </div>
+      {["streaming", "streaming-image"].includes(status) && (
+        <LoadingChatContent status={status} />
       )}
     </div>
   );

components/markdown-content.tsx

@@ -1,13 +1,15 @@
 "use client";
 
 import { useTheme } from "next-themes";
+import Image from "next/image";
 import { memo, useMemo } from "react";
 import ReactMarkdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import {
   oneLight as lightTheme,
   a11yDark as darkTheme,
 } from "react-syntax-highlighter/dist/esm/styles/prism";
+import remarkGfm from "remark-gfm";
 
 type MarkdownContentProps = {
   content: string;
@@ -19,7 +21,38 @@ function MarkdownContent({ content }: MarkdownContentProps) {
   const memoizedContent = useMemo(
     () => (
       <ReactMarkdown
+        remarkPlugins={[remarkGfm]}
         components={{
+          img(props) {
+            const { alt, src } = props;
+            if (src) {
+              return (
+                <Image
+                  width={512}
+                  height={512}
+                  alt={alt || "Generated image"}
+                  src={src}
+                  className="rounded-xl border-8 mt-2 mb-2 border-secondary"
+                />
+              );
+            } else {
+              return (
+                <>
+                  <Image
+                    width={1024}
+                    height={1024}
+                    alt={alt || "Generated image"}
+                    src="/no-image.png"
+                    className="rounded-xl border-8 mt-2 mb-2 border-secondary"
+                  />
+                  <span className="block text-destructive">
+                    The image could not be rendered. Please try again.
+                  </span>
+                </>
+              );
+            }
+          },
+
           code(props) {
             // eslint-disable-next-line @typescript-eslint/no-unused-vars
             const { children, className, node, ref, ...rest } = props;

db/schema.ts

@@ -99,6 +99,7 @@ export const messagesTable = pgTable(
     model: varchar("model", { length: 20 }).notNull(),
     modelParams: json("model_params").notNull(),
     content: text("content").notNull(),
+    kind: varchar("kind", { length: 20 }),
     toolResponses: json("tool_responses"),
     reasoning: text("reasoning"),
     searchMetadata: json("search_metadata"),

drizzle/0001_moaning_zombie.sql

@@ -0,0 +1 @@
+ALTER TABLE "messages" ADD COLUMN "kind" varchar(20);