feat: added db operations for creating chat and messages

Author: Praashh

prisma/migrations/20250617075116_added_role_in_message_table/migration.sql

@@ -0,0 +1,16 @@
+/*
+  Warnings:
+
+  - A unique constraint covering the columns `[userId]` on the table `Subscription` will be added. If there are existing duplicate values, this will fail.
+  - Changed the type of `role` on the `Message` table. No cast exists, the column would be dropped and recreated, which cannot be done if there is data, since the column is required.
+
+*/
+-- CreateEnum
+CREATE TYPE "MessageRole" AS ENUM ('USER', 'ASSISTANT');
+
+-- AlterTable
+ALTER TABLE "Message" DROP COLUMN "role",
+ADD COLUMN     "role" "MessageRole" NOT NULL;
+
+-- CreateIndex
+CREATE UNIQUE INDEX "Subscription_userId_key" ON "Subscription"("userId");

prisma/schema.prisma

@@ -102,7 +102,7 @@ model Message {
     chatId String
     chat Chat @relation(fields: [chatId], references: [id], onDelete: Cascade)
     content String
-    role String
+    role MessageRole
     createdAt DateTime @default(now())
     updatedAt DateTime @updatedAt
     @@index([chatId])
@@ -125,4 +125,9 @@ enum TransactionStatus {
   PENDING
   SUCCESS
   FAILED
+}
+
+enum MessageRole {
+  USER
+  ASSISTANT
 }

src/app/api/ask/route.ts

@@ -2,7 +2,7 @@ import { env } from "@/env";
 import { auth } from "@/server/auth";
 import { fetchChatCompletion } from "@/models/service";
 import { DEFAULT_MODEL_ID, getModelById } from "@/models/constants";
-
+import { db } from "@/server/db";
 interface ChatMessage {
   role: "user" | "assistant" | "system";
   content: string;
@@ -11,6 +11,7 @@ interface ChatMessage {
 interface TypeGPTPayload {
   messages: ChatMessage[];
   model?: string;
+  chatId: string;
 }
 
 interface TypeGPTErrorResponse {
@@ -30,9 +31,18 @@ export async function POST(req: Request): Promise<Response> {
       return new Response("Unauthorized", { status: 401 });
     }
 
-    const { messages, model = DEFAULT_MODEL_ID } =
+    const { messages, model = DEFAULT_MODEL_ID, chatId } =
       (await req.json()) as TypeGPTPayload;
 
+    // Save user message
+    await db.message.create({
+      data: {
+        chatId,
+        content: messages[messages.length - 1]?.content ?? "",
+        role: "USER",
+      },
+    });
+
     const modelInfo = getModelById(model);
     if (!modelInfo) {
       return new Response(
@@ -50,6 +60,8 @@ export async function POST(req: Request): Promise<Response> {
     const writer = writable.getWriter();
 
     void (async () => {
+      let accumulatedContent = "";
+      
       try {
         const response = await fetchChatCompletion({
           modelId: model,
@@ -89,11 +101,53 @@ export async function POST(req: Request): Promise<Response> {
             const { done, value } = await reader.read();
 
             if (done) {
+              // Save AI message to database after stream completes
+              if (accumulatedContent.trim()) {
+                await db.message.create({
+                  data: {
+                    chatId,
+                    content: accumulatedContent,
+                    role: "ASSISTANT",
+                  },
+                });
+              }
+              
               await writer.write(encoder.encode("data: [DONE]\n\n"));
               await writer.close();
               break;
             }
 
+            // Accumulate content for database storage
+            const chunk = new TextDecoder().decode(value);
+            const lines = chunk.split("\n");
+            
+            for (const line of lines) {
+              if (line.trim() === "") continue;
+              
+              if (line.startsWith("data: ")) {
+                const data = line.substring(6);
+                
+                if (data === "[DONE]") continue;
+                
+                try {
+                  const parsedData = JSON.parse(data) as {
+                    choices?: Array<{
+                      delta?: {
+                        content?: string;
+                      };
+                    }>;
+                  };
+                  
+                  const content = parsedData.choices?.[0]?.delta?.content;
+                  if (content) {
+                    accumulatedContent += content;
+                  }
+                } catch (e) {
+                  // Ignore parsing errors for non-JSON lines
+                }
+              }
+            }
+
             await writer.write(value);
           }
         } catch (error) {
@@ -112,7 +166,6 @@ export async function POST(req: Request): Promise<Response> {
       }
     })();
 
-
     return new Response(readable, { headers });
   } catch (error) {
     console.error("Error in chat API:", error);

src/components/ui/ui-input.tsx

@@ -3,24 +3,13 @@ import React, { useState, useRef, useEffect } from "react";
 import { useSession } from "next-auth/react";
 import { Textarea } from "@/components/ui/textarea";
 import { Button } from "@/components/ui/button";
-import { Badge } from "@/components/ui/badge";
 import {
   ChatCircleDotsIcon,
   MicrophoneIcon,
   SpinnerGapIcon,
 } from "@phosphor-icons/react";
-import {
-  Select,
-  SelectContent,
-  SelectGroup,
-  SelectItem,
-  SelectLabel,
-  SelectTrigger,
-  SelectValue,
-} from "@/components/ui/select";
 import ReactMarkdown from "react-markdown";
 import SyntaxHighlighter from "react-syntax-highlighter";
-import { vs2015 } from "react-syntax-highlighter/dist/esm/styles/hljs";
 import remarkGfm from "remark-gfm";
 import { Geist_Mono } from "next/font/google";
 import { cn } from "@/lib/utils";
@@ -61,7 +50,7 @@ const UIInput = () => {
     scrollToBottom();
   }, [messages]);
 
-  const saveChat = api.chat.createChat.useMutation({
+  const createChat = api.chat.createChat.useMutation({
     onError: (error) => {
       console.error("Error saving chat:", error);
     },
@@ -74,6 +63,8 @@ const UIInput = () => {
       return;
     }
 
+    const tempMessageId = `ai-${Date.now()}`;
+
     try {
       const reader = response.body?.getReader();
       if (!reader) {
@@ -82,8 +73,6 @@ const UIInput = () => {
         return;
       }
 
-      const tempMessageId = `ai-${Date.now()}`;
-
       setMessages((prev) => [
         ...prev,
         { id: tempMessageId, role: "assistant", content: "" },
@@ -111,6 +100,7 @@ const UIInput = () => {
         const { done, value } = await reader.read();
 
         if (done) {
+          // Final update with complete content
           setMessages((prev) =>
             prev.map((msg) =>
               msg.id === tempMessageId
@@ -152,8 +142,22 @@ const UIInput = () => {
                     content?: string;
                   };
                 }>;
+                error?: string;
               };
 
+              // Handle error responses
+              if (parsedData.error) {
+                console.error("Stream error:", parsedData.error);
+                setMessages((prev) =>
+                  prev.map((msg) =>
+                    msg.id === tempMessageId
+                      ? { ...msg, content: `Error: ${parsedData.error}` }
+                      : msg,
+                  ),
+                );
+                break;
+              }
+
               const content = parsedData.choices?.[0]?.delta?.content;
               if (content) {
                 accumulatedContent += content;
@@ -169,13 +173,16 @@ const UIInput = () => {
           updateMessage(accumulatedContent);
         }
       }
-
-      saveChat.mutate({
-        message: userMessage,
-        model: model,
-      });
     } catch (error) {
       console.error("Error processing stream:", error);
+      // Update message with error state
+      setMessages((prev) =>
+        prev.map((msg) =>
+          msg.id === tempMessageId
+            ? { ...msg, content: "Error: Failed to process response" }
+            : msg,
+        ),
+      );
     } finally {
       setIsLoading(false);
       abortControllerRef.current = null;
@@ -209,6 +216,7 @@ const UIInput = () => {
     abortControllerRef.current = new AbortController();
 
     try {
+      const {chatId} = await createChat.mutateAsync();
       setTimeout(() => {
         void (async () => {
           try {
@@ -220,6 +228,7 @@ const UIInput = () => {
               body: JSON.stringify({
                 messages: [{ role: "user", content: currentQuery }],
                 model: model,
+                chatId: chatId,
               }),
               signal: abortControllerRef.current?.signal,
             });
@@ -395,6 +404,12 @@ const UIInput = () => {
               <Textarea
                 value={query}
                 onChange={(e) => setQuery(e.target.value)}
+                onKeyDown={(e) => {
+                  if (e.key === "Enter" && !e.shiftKey) {
+                    e.preventDefault();
+                    void handleCreateChat(e as any);
+                  }
+                }}
                 placeholder="Ask whatever you want to be"
                 className="h-[2rem] resize-none rounded-none border-none bg-transparent px-0 py-1 shadow-none ring-0 focus-visible:ring-0 dark:bg-transparent"
                 disabled={isLoading}