From d11648f2bb61ff423ccbe6099fa05a3413a8fd43 Mon Sep 17 00:00:00 2001
From: Steven <stevenlgtm@gmail.com>
Date: Sun, 19 Mar 2023 23:19:08 +0800
Subject: [PATCH] feat: add prompt for assistant

---
 components/ChatView/MessageView.tsx |  2 +-
 components/ChatView/index.tsx       | 17 ++++++++++++-----
 pages/api/chat.ts                   |  8 --------
 store/assistant.ts                  | 20 ++++++++++++++++++++
 store/index.ts                      |  1 +
 store/user.ts                       | 12 +-----------
 6 files changed, 35 insertions(+), 25 deletions(-)
 create mode 100644 store/assistant.ts

diff --git a/components/ChatView/MessageView.tsx b/components/ChatView/MessageView.tsx
index 268c41d..4ef5467 100644
--- a/components/ChatView/MessageView.tsx
+++ b/components/ChatView/MessageView.tsx
@@ -15,7 +15,7 @@ const MessageView = (props: Props) => {
   return (
     <div className={`w-full flex flex-row justify-start items-start my-4 ${isCurrentUser ? "justify-end pl-8 sm:pl-24" : "pr-8 sm:pr-24"}`}>
       {isCurrentUser ? (
-        <div className="w-auto max-w-full bg-indigo-600 text-white px-4 py-2 rounded-lg rounded-tr-none shadow whitespace-pre">
+        <div className="w-auto max-w-full bg-indigo-600 text-white px-4 py-2 rounded-lg rounded-tr-none shadow whitespace-pre-wrap">
           {message.content}
         </div>
       ) : (
diff --git a/components/ChatView/index.tsx b/components/ChatView/index.tsx
index dea15e5..6ab3986 100644
--- a/components/ChatView/index.tsx
+++ b/components/ChatView/index.tsx
@@ -1,6 +1,6 @@
 import axios from "axios";
 import { useEffect, useRef, useState } from "react";
-import { defaultChat, useChatStore, useMessageStore, useUserStore } from "../../store";
+import { defaultChat, getPromptOfAssistant, useChatStore, useMessageStore, useUserStore } from "../../store";
 import { Chat, Message, UserRole } from "../../types";
 import { generateUUID } from "../../utils";
 import Icon from "../Icon";
@@ -47,11 +47,18 @@ const ChatView = () => {
 
     setIsRequesting(true);
     const messageList = messageStore.getState().messageList.filter((message) => message.chatId === currentChat.id);
+    const prompt = getPromptOfAssistant(userStore.getAssistantById(currentChat.assistantId)!);
     const { data } = await axios.post<string>("/api/chat", {
-      messages: messageList.map((message) => ({
-        role: message.creatorId === userStore.currentUser.id ? UserRole.User : UserRole.Assistant,
-        content: message.content,
-      })),
+      messages: [
+        {
+          role: "system",
+          content: prompt,
+        },
+        ...messageList.map((message) => ({
+          role: message.creatorId === userStore.currentUser.id ? UserRole.User : UserRole.Assistant,
+          content: message.content,
+        })),
+      ],
     });
     messageStore.addMessage({
       id: generateUUID(),
diff --git a/pages/api/chat.ts b/pages/api/chat.ts
index 7000a66..bc9ea52 100644
--- a/pages/api/chat.ts
+++ b/pages/api/chat.ts
@@ -14,12 +14,4 @@ const handler = async (req: NextApiRequest, res: NextApiResponse) => {
   res.status(200).json(completionResponse.data.choices[0].message?.content || "");
 };
 
-// TODO(steven): Implement a generic getChatPrompt function that takes in a
-// message and a robot identifier, then returns a string with specific prompts.
-const getChatPrompt = async (message: string) => {
-  return `
-  Question: ${message}
-  Answer:`;
-};
-
 export default handler;
diff --git a/store/assistant.ts b/store/assistant.ts
new file mode 100644
index 0000000..558c0b9
--- /dev/null
+++ b/store/assistant.ts
@@ -0,0 +1,20 @@
+import { User, UserRole } from "../types";
+
+// Assistant is a special user.
+export const assistantList: User[] = [
+  {
+    id: "assistant-dba",
+    name: "ChatDBA",
+    description: "🤖️ I am a chatbot that can help you with database administration.",
+    avatar: "",
+    role: UserRole.Assistant,
+  },
+];
+
+// getPromptOfAssistant define the special prompt for each assistant.
+export const getPromptOfAssistant = (assistant: User) => {
+  if (assistant.id === "assistant-dba") {
+    return `Remember that you are a DBA who is well versed in various databases. And you know everything about databases. You will answer some questions about databases.`;
+  }
+  return "";
+};
diff --git a/store/index.ts b/store/index.ts
index 5e5bd31..b4553f1 100644
--- a/store/index.ts
+++ b/store/index.ts
@@ -1,3 +1,4 @@
 export * from "./user";
+export * from "./assistant";
 export * from "./chat";
 export * from "./message";
diff --git a/store/user.ts b/store/user.ts
index c5e734e..e5f6d68 100644
--- a/store/user.ts
+++ b/store/user.ts
@@ -1,16 +1,7 @@
 import { create } from "zustand";
+import { assistantList } from ".";
 import { Id, User, UserRole } from "../types";
 
-export const assistantList: User[] = [
-  {
-    id: "assistant-dba",
-    name: "ChatDBA",
-    description: "🤖️ I am a chatbot that can help you with database administration.",
-    avatar: "",
-    role: UserRole.Assistant,
-  },
-];
-
 const localUser: User = {
   id: "local-user",
   name: "Local user",
@@ -20,7 +11,6 @@ const localUser: User = {
 };
 
 interface UserState {
-  // We can think assistants are special users.
   assistantList: User[];
   currentUser: User;
   getAssistantById: (id: string) => User | undefined;