Merge pull request #13 from e-roy/remove-last-message-and-refresh

feat: remove last message and reload request

src/app/api/gemini-pro/route.ts

@@ -6,6 +6,7 @@ import {
   HarmCategory,
   HarmBlockThreshold,
   GenerateContentRequest,
+  Content,
 } from "@google/generative-ai";
 
 import { GeneralSettings } from "@/types";
@@ -43,22 +44,28 @@ export async function POST(req: Request) {
   // console.log("safety_settings", safety_settings);
   // console.log("messages =================>", messages);
 
+  // consecutive user messages need to be merged into the same content, 2 consecutive Content objects will error with the Gemini api
   const reqContent: GenerateContentRequest = {
-    contents: messages.map((m: Message) => {
+    contents: messages.reduce((acc: Content[], m: Message) => {
       if (m.role === "user") {
-        return {
-          role: "user",
-          parts: [{ text: m.content }],
-        };
-      }
-      if (m.role === "assistant") {
-        return {
+        const lastContent = acc[acc.length - 1];
+        if (lastContent && lastContent.role === "user") {
+          lastContent.parts.push({ text: m.content });
+        } else {
+          acc.push({
+            role: "user",
+            parts: [{ text: m.content }],
+          });
+        }
+      } else if (m.role === "assistant") {
+        acc.push({
           role: "model",
           parts: [{ text: m.content }],
-        };
+        });
       }
-      return undefined;
-    }),
+
+      return acc;
+    }, []),
   };
 
   const incomingSafetySettings = safety_settings || defaultSafetySettings;

src/components/ChatContainer.tsx

@@ -1,41 +1,35 @@
 "use client";
 // components/ChatContainer.tsx
-import React, {
-  useRef,
-  useEffect,
-  useState,
-  FormEvent,
-  useCallback,
-} from "react";
+import React, { useRef, useEffect, FormEvent, useCallback } from "react";
 import { Message, useChat } from "ai/react";
 import { Card } from "./ui/card";
 
-import { MarkdownViewer } from "./markdown-viewer/MarkdownViewer";
 import { useControlContext } from "@/providers/ControlContext";
 import { CommonForm } from "./CommonForm";
 import { TypingBubble } from "./TypingBubble";
-import { MessageCircleX, User, Bot } from "lucide-react";
+import { MessageCircleX } from "lucide-react";
 import { Button } from "./ui/button";
+import { MessageItem } from "./MessageItem";
 
 export const ChatContainer = () => {
   const { generalSettings, safetySettings } = useControlContext();
-  const [loading, setLoading] = useState<boolean>(false);
 
-  const { messages, setMessages, input, handleInputChange, handleSubmit } =
-    useChat({
-      id: "gemini-pro",
-      api: `/api/gemini-pro`,
-      body: {
-        general_settings: generalSettings,
-        safety_settings: safetySettings,
-      },
-      onError: () => {
-        setLoading(false);
-      },
-      onFinish: () => {
-        setLoading(false);
-      },
-    });
+  const {
+    messages,
+    setMessages,
+    input,
+    handleInputChange,
+    handleSubmit,
+    reload,
+    isLoading,
+  } = useChat({
+    id: "gemini-pro",
+    api: `/api/gemini-pro`,
+    body: {
+      general_settings: generalSettings,
+      safety_settings: safetySettings,
+    },
+  });
 
   const messagesEndRef = useRef<HTMLDivElement>(null);
 
@@ -49,13 +43,32 @@ export const ChatContainer = () => {
 
   const handleFormSubmit = async (event: FormEvent<HTMLFormElement>) => {
     event.preventDefault();
-    setLoading(true);
     handleSubmit(event);
   };
 
-  const handleClearChat = () => {
+  const handleClearChat = useCallback(() => {
     setMessages([]);
-  };
+  }, [setMessages]);
+
+  const handleRefreshMessage = useCallback(() => {
+    reload();
+  }, [reload]);
+
+  const messagesRef = useRef<Message[]>(messages);
+
+  useEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
+
+  const handleRemoveMessage = useCallback(
+    (id: string) => {
+      const newMessages = messagesRef.current.filter(
+        (message) => message.id !== id
+      );
+      setMessages(newMessages);
+    },
+    [setMessages]
+  );
 
   return (
     <div className="flex flex-col h-[95vh]">
@@ -73,39 +86,24 @@ export const ChatContainer = () => {
           </div>
         )}
         <div className="flex-1 overflow-y-auto">
-          {messages.map((message: Message) => (
-            <div key={message.id} className={`flex`}>
-              <div
-                className={`${
-                  message.role === "user"
-                    ? ""
-                    : "bg-primary/10 dark:bg-primary/10"
-                } px-4 py-8 w-full`}
-              >
-                <div className={`my-4`}>
-                  {message.role === "user" ? (
-                    <div className={`flex space-x-4 font-medium`}>
-                      <User />
-                      <div>You</div>
-                    </div>
-                  ) : (
-                    <div className={`flex space-x-4 font-medium`}>
-                      <Bot />
-                      <div>Gemini Pro</div>
-                    </div>
-                  )}
-                </div>
-                <MarkdownViewer text={message.content} />
-              </div>
-            </div>
+          {messages.map((message, index) => (
+            <MessageItem
+              key={message.id}
+              message={message}
+              isLastMessage={messages.length === index + 1}
+              isLoading={isLoading}
+              onRefresh={handleRefreshMessage}
+              onRemove={() => handleRemoveMessage(message.id)}
+            />
           ))}
+
           <div ref={messagesEndRef} />
-          {loading && <TypingBubble />}
+          {isLoading && <TypingBubble />}
         </div>
         <CommonForm
           value={input}
           placeholder="Chat with Gemini Pro"
-          loading={loading}
+          loading={isLoading}
           onInputChange={handleInputChange}
           onFormSubmit={handleFormSubmit}
           isSubmittable={input.trim() !== ""}

src/components/MessageItem.tsx

@@ -0,0 +1,65 @@
+import React, { memo } from "react";
+import { Message } from "ai/react";
+import { Bot, RefreshCw, Trash, User } from "lucide-react";
+import { MarkdownViewer } from "./markdown-viewer/MarkdownViewer";
+import { Button } from "./ui/button";
+
+type MessageItemProps = {
+  message: Message;
+  isLastMessage: boolean;
+  isLoading: boolean;
+  onRefresh: () => void;
+  onRemove: () => void;
+};
+
+export const MessageItem: React.FC<MessageItemProps> = memo(
+  function MessageItem({
+    message,
+    isLastMessage,
+    isLoading,
+    onRefresh,
+    onRemove,
+  }) {
+    const isUser = message.role === "user";
+
+    return (
+      <div className="flex">
+        <div
+          className={`${
+            isUser ? "" : "bg-primary/10 dark:bg-primary/10"
+          } px-4 py-4 w-full`}
+        >
+          <div className="my-4 flex justify-between">
+            <div className="flex space-x-4 font-medium">
+              {isUser ? <User /> : <Bot />}
+              <div>{isUser ? "You" : "Gemini Pro"}</div>
+            </div>
+            <div className={`space-x-6`}>
+              {!isLoading && isLastMessage && isUser && (
+                <Button
+                  variant="icon"
+                  type="button"
+                  size="sm"
+                  onClick={onRefresh}
+                >
+                  <RefreshCw className="w-4 h-4" />
+                </Button>
+              )}
+              {!isLoading && isLastMessage && (
+                <Button
+                  variant="icon"
+                  type="button"
+                  size="sm"
+                  onClick={onRemove}
+                >
+                  <Trash className="w-4 h-4" />
+                </Button>
+              )}
+            </div>
+          </div>
+          <MarkdownViewer text={message.content} />
+        </div>
+      </div>
+    );
+  }
+);

src/components/control/ControlContainer.tsx

@@ -12,12 +12,12 @@ import {
   DropdownMenuItem,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
-import ImageUploadComponent from "../ImageUploadComponent";
 import { useControlContext } from "@/providers/ControlContext";
 
 import { ChevronDown } from "lucide-react";
 import { SettingsSelector } from "./SettingsSelector";
 import { ExtraButtons } from "./ExtraButtons";
+import ImageUploadComponent from "./ImageUploadComponent";
 
 import { settings_data } from "./settings_data";
 

src/components/ImageUploadComponent.tsx → src/components/control/ImageUploadComponent.tsx

@@ -1,7 +1,7 @@
 "use client";
 // components/ImageUploadComponent.tsx
 import React, { useState, useCallback, memo } from "react";
-import { Card } from "./ui/card";
+import { Card } from "../ui/card";
 import { useDropzone } from "react-dropzone";
 import { Import, Upload, XCircle } from "lucide-react";
 

src/components/markdown-viewer/code.tsx

@@ -50,22 +50,34 @@ interface CodeBlockProps {
   children?: string | React.ReactNode;
 }
 
-export const CodeBlock: React.FC<CodeBlockProps> = ({ ...props }) => {
-  const { children, className, ...rest } = props;
-  const match = /language-(\w+)/.exec(className as string) || [];
+export const CodeBlock: React.FC<CodeBlockProps> = ({
+  children,
+  className,
+  ...rest
+}) => {
+  const match = /language-(\w+)/.exec(className || "");
+  const language = match?.[1];
 
-  return match ? (
-    <SyntaxHighlighter
-      {...rest}
-      PreTag="div"
-      language={match[1]}
-      style={a11yDark}
-    >
-      {String(children).replace(/\n$/, "")}
-    </SyntaxHighlighter>
-  ) : (
-    <code {...rest} className={className}>
-      {children}
-    </code>
-  );
+  if (language) {
+    return (
+      <SyntaxHighlighter
+        {...rest}
+        PreTag="div"
+        language={language}
+        style={a11yDark}
+      >
+        {String(children).replace(/\n$/, "")}
+      </SyntaxHighlighter>
+    );
+  }
+
+  if (typeof children === "string") {
+    return (
+      <code {...rest} className={className}>
+        {children}
+      </code>
+    );
+  }
+
+  return null;
 };