From d03ce9fac1b809b0febb562a2adfedda3e66963f Mon Sep 17 00:00:00 2001 From: steven Date: Mon, 17 Apr 2023 16:14:39 +0800 Subject: [PATCH] feat: implement general assistant --- assistants/README.md | 1 + assistants/general-bot/README.md | 1 + assistants/general-bot/index.ts | 10 +++++++++ assistants/index.ts | 3 +++ assistants/migration-bot/index.ts | 10 +++++++++ assistants/sql-chat-bot/index.ts | 14 +++++++++++++ src/store/assistant.ts | 34 ++++++++++--------------------- src/store/conversation.ts | 16 +++++++++++++++ src/types/assistant.ts | 7 +++++++ src/types/index.ts | 1 + 10 files changed, 74 insertions(+), 23 deletions(-) create mode 100644 assistants/README.md create mode 100644 assistants/general-bot/README.md create mode 100644 assistants/general-bot/index.ts create mode 100644 assistants/index.ts create mode 100644 assistants/migration-bot/index.ts create mode 100644 assistants/sql-chat-bot/index.ts create mode 100644 src/types/assistant.ts diff --git a/assistants/README.md b/assistants/README.md new file mode 100644 index 0000000..aa2ab15 --- /dev/null +++ b/assistants/README.md @@ -0,0 +1 @@ +# Assistant bots diff --git a/assistants/general-bot/README.md b/assistants/general-bot/README.md new file mode 100644 index 0000000..e0ccf9f --- /dev/null +++ b/assistants/general-bot/README.md @@ -0,0 +1 @@ +# General bot diff --git a/assistants/general-bot/index.ts b/assistants/general-bot/index.ts new file mode 100644 index 0000000..268eb03 --- /dev/null +++ b/assistants/general-bot/index.ts @@ -0,0 +1,10 @@ +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 new file mode 100644 index 0000000..ff98c0b --- /dev/null +++ b/assistants/index.ts @@ -0,0 +1,3 @@ +export * as generalBot from "./general-bot"; +export * as sqlchatBot from "./sql-chat-bot"; +export * as migrationBot from "./migration-bot"; diff --git a/assistants/migration-bot/index.ts b/assistants/migration-bot/index.ts new file mode 100644 index 0000000..8a49075 --- /dev/null +++ b/assistants/migration-bot/index.ts @@ -0,0 +1,10 @@ +export default { + id: "migration-bot", + name: "Migration bot", + description: "A bot focused on database migration.", + avatar: "", + getPrompt: (input?: string): string => { + // TODO: update prompt. + return ""; + }, +}; diff --git a/assistants/sql-chat-bot/index.ts b/assistants/sql-chat-bot/index.ts new file mode 100644 index 0000000..2554952 --- /dev/null +++ b/assistants/sql-chat-bot/index.ts @@ -0,0 +1,14 @@ +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}`; + }, +}; diff --git a/src/store/assistant.ts b/src/store/assistant.ts index 0fc3028..4cf5961 100644 --- a/src/store/assistant.ts +++ b/src/store/assistant.ts @@ -1,30 +1,18 @@ import { first } from "lodash-es"; -import { Id, User } from "@/types"; +import { Assistant, Id } from "@/types"; +import * as customAssistantList from "../../assistants"; -// Assistant is a special user. -export const assistantList: User[] = [ - { - id: "sql-assistant", - name: "SQL Chat", - description: "🤖️ I'm an expert in SQL. I can answer your questions about databases and SQL.", - avatar: "", - }, -]; +const assistantList: Assistant[] = Object.keys(customAssistantList).map((name) => { + return { + ...((customAssistantList as any)[name].default as Assistant), + }; +}); export const getAssistantById = (id: Id) => { - const user = assistantList.find((user) => user.id === id); - return user || (first(assistantList) as User); + const assistant = assistantList.find((assistant) => assistant.id === id); + return assistant || (first(assistantList) as Assistant); }; -// getPromptOfAssistant define the special prompt for each assistant. -export const getPromptGeneratorOfAssistant = (assistant: User) => { - 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. -2. Please be careful to return only key information, and try not to make it too long. -`; - if (assistant.id === "sql-assistant") { - return (schema: string) => - `This 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}`; - } - return () => `\n${basicPrompt}`; +export const getPromptGeneratorOfAssistant = (assistant: Assistant) => { + return assistant.getPrompt; }; diff --git a/src/store/conversation.ts b/src/store/conversation.ts index 47be74e..2b05a19 100644 --- a/src/store/conversation.ts +++ b/src/store/conversation.ts @@ -60,6 +60,22 @@ export const useConversationStore = create()( }), { name: "conversation-storage", + version: 1, + migrate: (persistedState: any, version: number) => { + 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"; + } + } + state.currentConversation = undefined; + } + + return state; + }, } ) ); diff --git a/src/types/assistant.ts b/src/types/assistant.ts new file mode 100644 index 0000000..a31d90f --- /dev/null +++ b/src/types/assistant.ts @@ -0,0 +1,7 @@ +export interface Assistant { + id: string; + name: string; + description: string; + avatar: string; + getPrompt: (input?: string) => string; +} diff --git a/src/types/index.ts b/src/types/index.ts index 4f74f47..500c609 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -7,3 +7,4 @@ export * from "./message"; export * from "./setting"; export * from "./api"; export * from "./connector"; +export * from "./assistant";