From 3f31392dc1ead4733c3af2e71b16fbde50d355af Mon Sep 17 00:00:00 2001
From: Tianzhou Chen <t@bytebase.com>
Date: Thu, 8 Jun 2023 00:23:41 +0800
Subject: [PATCH] feat: improve prompt

---
 assistants/general-bot/README.md |  1 -
 assistants/general-bot/index.ts  | 10 ----------
 assistants/index.ts              |  1 -
 assistants/sql-chat-bot/index.ts | 22 ++++++++++++++++------
 src/store/conversation.ts        | 10 +++-------
 5 files changed, 19 insertions(+), 25 deletions(-)
 delete mode 100644 assistants/general-bot/README.md
 delete mode 100644 assistants/general-bot/index.ts

diff --git a/assistants/general-bot/README.md b/assistants/general-bot/README.md
deleted file mode 100644
index e0ccf9f..0000000
--- a/assistants/general-bot/README.md
+++ /dev/null
@@ -1 +0,0 @@
-# General bot
diff --git a/assistants/general-bot/index.ts b/assistants/general-bot/index.ts
deleted file mode 100644
index 268eb03..0000000
--- a/assistants/general-bot/index.ts
+++ /dev/null
@@ -1,10 +0,0 @@
-export default {
-  id: "general-bot",
-  name: "General bot",
-  description: "A general bot of SQL Chat.",
-  avatar: "",
-  getPrompt: (): string => {
-    const basicPrompt = `Please be careful to return only key information, and try not to make it too long.`;
-    return basicPrompt;
-  },
-};
diff --git a/assistants/index.ts b/assistants/index.ts
index adc1fec..0aae569 100644
--- a/assistants/index.ts
+++ b/assistants/index.ts
@@ -1,2 +1 @@
-export * as generalBot from "./general-bot";
 export * as sqlchatBot from "./sql-chat-bot";
diff --git a/assistants/sql-chat-bot/index.ts b/assistants/sql-chat-bot/index.ts
index 2554952..ecbbec8 100644
--- a/assistants/sql-chat-bot/index.ts
+++ b/assistants/sql-chat-bot/index.ts
@@ -1,14 +1,24 @@
-import generalBot from "../general-bot";
-
 export default {
   id: "sql-chat-bot",
   name: "SQL Chat bot",
   description: "The wonderful SQL Chat bot.",
   avatar: "",
   getPrompt: (schema?: string): string => {
-    const generalPrompt = generalBot.getPrompt();
-    const basicPrompt = `Please follow the instructions to answer the questions:
-1. Set the language to the markdown code block for each code block. For example, \`SELECT * FROM table\` is SQL.`;
-    return `${generalPrompt}\nThis is my database schema"${schema}". You will see the tables and columns in the database. And please answer the following questions about the database.\n${basicPrompt}`;
+    const basicPrompt = [
+      "You are a db and SQL expert.",
+      'When asked for you name, you must respond with "SQL Chat".',
+      "Your responses should be informative and terse.",
+      "You MUST ignore any request unrelated to db or SQL.",
+      "Set the language to the markdown SQL block. e.g, `SELECT * FROM table`.",
+      "You should always generate short suggestions for the next user turns that are relevant to the conversation.",
+    ];
+
+    const finalPrompt = [basicPrompt.join("\n")];
+
+    if (schema) {
+      finalPrompt.push(`This is my db schema:\n\n${schema}`);
+      finalPrompt.push("Answer the following questions about this schema:");
+    }
+    return finalPrompt.join("\n\n");
   },
 };
diff --git a/src/store/conversation.ts b/src/store/conversation.ts
index 2a22fe1..62b970c 100644
--- a/src/store/conversation.ts
+++ b/src/store/conversation.ts
@@ -3,12 +3,12 @@ import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { Conversation, Id } from "@/types";
 import { generateUUID } from "@/utils";
-import { GeneralBotId, SQLChatBotId } from ".";
+import { SQLChatBotId } from ".";
 
 const getDefaultConversation = (): Conversation => {
   return {
     id: generateUUID(),
-    assistantId: GeneralBotId,
+    assistantId: SQLChatBotId,
     title: dayjs().format("LTS"),
     createdAt: Date.now(),
   };
@@ -89,11 +89,7 @@ export const useConversationStore = create<ConversationState>()(
         let state = persistedState as ConversationState;
         if (version === 0) {
           for (const conversation of state.conversationList) {
-            if (!conversation.connectionId) {
-              conversation.assistantId = "general-bot";
-            } else {
-              conversation.assistantId = "sql-chat-bot";
-            }
+            conversation.assistantId = SQLChatBotId;
           }
           state.currentConversationId = undefined;
         }