From a2e4866d65a52b6d828e2dff3350b71c1960b06d Mon Sep 17 00:00:00 2001
From: boojack <stevenlgtm@gmail.com>
Date: Wed, 29 Mar 2023 15:06:32 +0800
Subject: [PATCH] chore: update bot avatar (#8)

---
 components/ChatView/MessageView.tsx | 18 ++++++++++++------
 components/MessageLoader.tsx        | 14 +++++++++++---
 store/chat.ts                       |  4 ++++
 store/connection.ts                 |  4 ++++
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/components/ChatView/MessageView.tsx b/components/ChatView/MessageView.tsx
index f6d33ce..d0cda29 100644
--- a/components/ChatView/MessageView.tsx
+++ b/components/ChatView/MessageView.tsx
@@ -3,10 +3,11 @@ import { ReactElement, useState } from "react";
 import { toast } from "react-hot-toast";
 import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
-import { useMessageStore, useUserStore } from "@/store";
+import { useChatStore, useConnectionStore, useMessageStore, useUserStore } from "@/store";
 import { Message } from "@/types";
 import Icon from "../Icon";
 import { CodeBlock } from "../CodeBlock";
+import EngineIcon from "../EngineIcon";
 
 interface Props {
   message: Message;
@@ -15,10 +16,13 @@ interface Props {
 const MessageView = (props: Props) => {
   const message = props.message;
   const userStore = useUserStore();
+  const chatStore = useChatStore();
+  const connectionStore = useConnectionStore();
   const messageStore = useMessageStore();
   const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLButtonElement | null>(null);
   const isCurrentUser = message.creatorId === userStore.currentUser.id;
   const showMenu = Boolean(menuAnchorEl);
+  const connection = connectionStore.getConnectionById(chatStore.getChatById(message.chatId)?.connectionId || "");
 
   const handleMoreMenuClick = (event: React.MouseEvent<HTMLButtonElement>) => {
     if (menuAnchorEl) {
@@ -47,17 +51,19 @@ const MessageView = (props: Props) => {
     >
       {isCurrentUser ? (
         <>
-          <div className="w-auto max-w-full bg-indigo-600 text-white px-4 py-2 rounded-lg whitespace-pre-wrap">
-            {message.content}
-          </div>
+          <div className="w-auto max-w-full bg-indigo-600 text-white px-4 py-2 rounded-lg whitespace-pre-wrap">{message.content}</div>
           <div className="w-10 h-10 p-1 border rounded-full flex justify-center items-center ml-2 shrink-0">
             <Icon.AiOutlineUser className="w-6 h-6" />
           </div>
         </>
       ) : (
         <>
-          <div className="w-10 h-10 p-1 flex justify-center items-center mr-2 shrink-0">
-            <img src="/chat-logo-bot.webp" alt="" />
+          <div className="flex justify-center items-center mr-2 shrink-0">
+            {connection ? (
+              <EngineIcon className="w-10 h-auto p-1 border rounded-full" engine={connection.engineType} />
+            ) : (
+              <img className="w-10 h-auto p-1" src="/chat-logo-bot.webp" alt="" />
+            )}
           </div>
           <ReactMarkdown
             className="w-auto max-w-[calc(100%-4rem)] bg-gray-100 px-4 py-2 rounded-lg prose prose-neutral"
diff --git a/components/MessageLoader.tsx b/components/MessageLoader.tsx
index 8a6a9ad..603baa0 100644
--- a/components/MessageLoader.tsx
+++ b/components/MessageLoader.tsx
@@ -1,11 +1,19 @@
 import { ThreeDots } from "react-loader-spinner";
-import Icon from "./Icon";
+import { useConnectionStore } from "@/store";
+import EngineIcon from "./EngineIcon";
 
 const MessageLoader = () => {
+  const connectionStore = useConnectionStore();
+  const connection = connectionStore.currentConnectionCtx?.connection;
+
   return (
     <div className={`w-full max-w-full flex flex-row justify-start items-start my-4 pr-8 sm:pr-24`}>
-      <div className="w-10 h-10 p-1 flex justify-center items-center mr-2 shrink-0">
-        <img src="/chat-logo-bot.webp" alt="" />
+      <div className="flex justify-center items-center mr-2 shrink-0">
+        {connection ? (
+          <EngineIcon className="w-10 h-auto p-1 border rounded-full" engine={connection.engineType} />
+        ) : (
+          <img className="w-10 h-auto p-1" src="/chat-logo-bot.webp" alt="" />
+        )}
       </div>
       <div className="mt-0.5 w-12 bg-gray-100 px-4 py-2 rounded-lg">
         <ThreeDots wrapperClass="opacity-80" width="24" height="24" color="" />
diff --git a/store/chat.ts b/store/chat.ts
index da83239..54b838d 100644
--- a/store/chat.ts
+++ b/store/chat.ts
@@ -19,6 +19,7 @@ interface ChatState {
   getState: () => ChatState;
   createChat: (connectionId?: Id, databaseName?: string) => Chat;
   setCurrentChat: (chat: Chat | undefined) => void;
+  getChatById: (chatId: Id) => Chat | undefined;
   updateChat: (chatId: Id, chat: Partial<Chat>) => void;
   clearChat: (filter: (chat: Chat) => boolean) => void;
 }
@@ -41,6 +42,9 @@ export const useChatStore = create<ChatState>()(
         return chat;
       },
       setCurrentChat: (chat: Chat | undefined) => set(() => ({ currentChat: chat })),
+      getChatById: (chatId: Id) => {
+        return get().chatList.find((item) => item.id === chatId);
+      },
       updateChat: (chatId: Id, chat: Partial<Chat>) => {
         set((state) => ({
           ...state,
diff --git a/store/connection.ts b/store/connection.ts
index 95df435..ff84f61 100644
--- a/store/connection.ts
+++ b/store/connection.ts
@@ -29,6 +29,7 @@ interface ConnectionState {
   setCurrentConnectionCtx: (connectionCtx: ConnectionContext | undefined) => void;
   getOrFetchDatabaseList: (connection: Connection, skipCache?: boolean) => Promise<Database[]>;
   getOrFetchDatabaseSchema: (database: Database) => Promise<Table[]>;
+  getConnectionById: (connectionId: string) => Connection | undefined;
   updateConnection: (connectionId: string, connection: Partial<Connection>) => void;
   clearConnection: (filter: (connection: Connection) => boolean) => void;
 }
@@ -97,6 +98,9 @@ export const useConnectionStore = create<ConnectionState>()(
         });
         return data;
       },
+      getConnectionById: (connectionId: string) => {
+        return get().connectionList.find((connection) => connection.id === connectionId);
+      },
       updateConnection: (connectionId: string, connection: Partial<Connection>) => {
         set((state) => ({
           ...state,