feat: support messages to ask AI (#1380)

2024-12-24 20:01:48 +03:00 · 2023-03-18 22:07:14 +08:00 · 2023-03-18 22:07:14 +08:00 · 573f07ec82
commit 573f07ec82
parent 8b20cb9fd2
8 changed files with 99 additions and 108 deletions
--- a/api/openai.go
+++ b/api/openai.go
@ -1,5 +0,0 @@
-package api
-
-type OpenAICompletionRequest struct {
-	Prompt string `json:"prompt"`
-}
--- a/plugin/openai/chat_completion.go
+++ b/plugin/openai/chat_completion.go
@ -24,7 +24,7 @@ type ChatCompletionResponse struct {
 	Choices []ChatCompletionChoice `json:"choices"`
 }

-func PostChatCompletion(prompt string, apiKey string, apiHost string) (string, error) {
+func PostChatCompletion(messages []ChatCompletionMessage, apiKey string, apiHost string) (string, error) {
 	if apiHost == "" {
 		apiHost = "https://api.openai.com"
 	}
@ -34,8 +34,12 @@ func PostChatCompletion(prompt string, apiKey string, apiHost string) (string, e
 	}

 	values := map[string]interface{}{
-		"model":    "gpt-3.5-turbo",
-		"messages": []map[string]string{{"role": "user", "content": prompt}},
+		"model":             "gpt-3.5-turbo",
+		"messages":          messages,
+		"max_tokens":        2000,
+		"temperature":       0,
+		"frequency_penalty": 0.0,
+		"presence_penalty":  0.0,
 	}
 	jsonValue, err := json.Marshal(values)
 	if err != nil {
--- a/server/openai.go
+++ b/server/openai.go
@ -31,15 +31,15 @@ func (s *Server) registerOpenAIRoutes(g *echo.Group) {
 			return echo.NewHTTPError(http.StatusBadRequest, "OpenAI API key not set")
 		}

-		completionRequest := api.OpenAICompletionRequest{}
-		if err := json.NewDecoder(c.Request().Body).Decode(&completionRequest); err != nil {
+		messages := []openai.ChatCompletionMessage{}
+		if err := json.NewDecoder(c.Request().Body).Decode(&messages); err != nil {
 			return echo.NewHTTPError(http.StatusBadRequest, "Malformatted post chat completion request").SetInternal(err)
 		}
-		if completionRequest.Prompt == "" {
-			return echo.NewHTTPError(http.StatusBadRequest, "Prompt is required")
+		if len(messages) == 0 {
+			return echo.NewHTTPError(http.StatusBadRequest, "No messages provided")
 		}

-		result, err := openai.PostChatCompletion(completionRequest.Prompt, openAIConfig.Key, openAIConfig.Host)
+		result, err := openai.PostChatCompletion(messages, openAIConfig.Key, openAIConfig.Host)
 		if err != nil {
 			return echo.NewHTTPError(http.StatusInternalServerError, "Failed to post chat completion").SetInternal(err)
 		}
@ -47,42 +47,6 @@ func (s *Server) registerOpenAIRoutes(g *echo.Group) {
 		return c.JSON(http.StatusOK, composeResponse(result))
 	})

-	g.POST("/openai/text-completion", func(c echo.Context) error {
-		ctx := c.Request().Context()
-		openAIConfigSetting, err := s.Store.FindSystemSetting(ctx, &api.SystemSettingFind{
-			Name: api.SystemSettingOpenAIConfigName,
-		})
-		if err != nil && common.ErrorCode(err) != common.NotFound {
-			return echo.NewHTTPError(http.StatusInternalServerError, "Failed to find openai key").SetInternal(err)
-		}
-
-		openAIConfig := api.OpenAIConfig{}
-		if openAIConfigSetting != nil {
-			err = json.Unmarshal([]byte(openAIConfigSetting.Value), &openAIConfig)
-			if err != nil {
-				return echo.NewHTTPError(http.StatusInternalServerError, "Failed to unmarshal openai system setting value").SetInternal(err)
-			}
-		}
-		if openAIConfig.Key == "" {
-			return echo.NewHTTPError(http.StatusBadRequest, "OpenAI API key not set")
-		}
-
-		textCompletion := api.OpenAICompletionRequest{}
-		if err := json.NewDecoder(c.Request().Body).Decode(&textCompletion); err != nil {
-			return echo.NewHTTPError(http.StatusBadRequest, "Malformatted post text completion request").SetInternal(err)
-		}
-		if textCompletion.Prompt == "" {
-			return echo.NewHTTPError(http.StatusBadRequest, "Prompt is required")
-		}
-
-		result, err := openai.PostTextCompletion(textCompletion.Prompt, openAIConfig.Key, openAIConfig.Host)
-		if err != nil {
-			return echo.NewHTTPError(http.StatusInternalServerError, "Failed to post text completion").SetInternal(err)
-		}
-
-		return c.JSON(http.StatusOK, composeResponse(result))
-	})
-
 	g.GET("/openai/enabled", func(c echo.Context) error {
 		ctx := c.Request().Context()
 		openAIConfigSetting, err := s.Store.FindSystemSetting(ctx, &api.SystemSettingFind{
--- a/web/package.json
+++ b/web/package.json
@ -29,7 +29,8 @@
    "react-router-dom": "^6.8.2",
    "react-use": "^17.4.0",
    "semver": "^7.3.8",
-    "tailwindcss": "^3.2.4"
+    "tailwindcss": "^3.2.4",
+    "zustand": "^4.3.6"
  },
  "devDependencies": {
    "@types/lodash-es": "^4.17.5",
--- a/web/src/components/AskAIDialog.tsx
+++ b/web/src/components/AskAIDialog.tsx
@ -4,24 +4,21 @@ import { toast } from "react-hot-toast";
 import * as api from "../helpers/api";
 import useLoading from "../hooks/useLoading";
 import { marked } from "../labs/marked";
+import { useMessageStore } from "../store/zustand/message";
 import Icon from "./Icon";
 import { generateDialog } from "./Dialog";
 import showSettingDialog from "./SettingDialog";

 type Props = DialogProps;

-interface History {
-  question: string;
-  answer: string;
-}
-
 const AskAIDialog: React.FC<Props> = (props: Props) => {
  const { destroy, hide } = props;
  const fetchingState = useLoading(false);
-  const [historyList, setHistoryList] = useState<History[]>([]);
+  const messageStore = useMessageStore();
  const [isEnabled, setIsEnabled] = useState<boolean>(true);
  const [isInIME, setIsInIME] = useState(false);
  const [question, setQuestion] = useState<string>("");
+  const messageList = messageStore.messageList;

  useEffect(() => {
    api.checkOpenAIEnabled().then(({ data }) => {
@ -47,10 +44,18 @@ const AskAIDialog: React.FC<Props> = (props: Props) => {
  };

  const handleSendQuestionButtonClick = async () => {
+    if (!question) {
+      return;
+    }
+
    fetchingState.setLoading();
    setQuestion("");
+    messageStore.addMessage({
+      role: "user",
+      content: question,
+    });
    try {
-      await askQuestion(question);
+      await fetchChatCompletion();
    } catch (error: any) {
      console.error(error);
      toast.error(error.response.data.error);
@ -58,21 +63,15 @@ const AskAIDialog: React.FC<Props> = (props: Props) => {
    fetchingState.setFinish();
  };

-  const askQuestion = async (question: string) => {
-    if (question === "") {
-      return;
-    }
-
+  const fetchChatCompletion = async () => {
+    const messageList = messageStore.getState().messageList;
    const {
      data: { data: answer },
-    } = await api.postChatCompletion(question);
-    setHistoryList([
-      {
-        question,
-        answer: answer.replace(/^\n\n/, ""),
-      },
-      ...historyList,
-    ]);
+    } = await api.postChatCompletion(messageList);
+    messageStore.addMessage({
+      role: "assistant",
+      content: answer.replace(/^\n\n/, ""),
+    });
  };

  return (
@ -87,7 +86,36 @@ const AskAIDialog: React.FC<Props> = (props: Props) => {
        </button>
      </div>
      <div className="dialog-content-container !w-112 max-w-full">
-        <div className="w-full relative">
+        {messageList.map((message, index) => (
+          <div key={index} className="w-full flex flex-col justify-start items-start mt-4 space-y-2">
+            {message.role === "user" ? (
+              <div className="w-full flex flex-row justify-end items-start pl-6">
+                <span className="word-break shadow rounded-lg rounded-tr-none px-3 py-2 opacity-80 bg-gray-100 dark:bg-zinc-700">
+                  {message.content}
+                </span>
+              </div>
+            ) : (
+              <div className="w-full flex flex-row justify-start items-start pr-8 space-x-2">
+                <Icon.Bot className="mt-2 flex-shrink-0 mr-1 w-6 h-auto opacity-80" />
+                <div className="memo-content-wrapper !w-auto flex flex-col justify-start items-start shadow rounded-lg rounded-tl-none px-3 py-2 bg-gray-100 dark:bg-zinc-700">
+                  <div className="memo-content-text">{marked(message.content)}</div>
+                </div>
+              </div>
+            )}
+          </div>
+        ))}
+        {fetchingState.isLoading && (
+          <p className="w-full py-2 mt-4 flex flex-row justify-center items-center">
+            <Icon.Loader className="w-5 h-auto animate-spin" />
+          </p>
+        )}
+        {!isEnabled && (
+          <div className="w-full flex flex-col justify-center items-center mt-4 space-y-2">
+            <p>You have not set up your OpenAI API key.</p>
+            <Button onClick={() => handleGotoSystemSetting()}>Go to settings</Button>
+          </div>
+        )}
+        <div className="w-full relative mt-4">
          <Textarea
            className="w-full"
            placeholder="Ask anything…"
@ -104,32 +132,6 @@ const AskAIDialog: React.FC<Props> = (props: Props) => {
            onClick={handleSendQuestionButtonClick}
          />
        </div>
-        {fetchingState.isLoading && (
-          <p className="w-full py-2 mt-4 flex flex-row justify-center items-center">
-            <Icon.Loader className="w-5 h-auto animate-spin" />
-          </p>
-        )}
-        {historyList.map((history, index) => (
-          <div key={index} className="w-full flex flex-col justify-start items-start mt-4 space-y-2">
-            <div className="w-full flex flex-row justify-start items-start pr-6">
-              <span className="word-break rounded-lg rounded-tl-none px-3 py-2 opacity-80 bg-gray-100 dark:bg-zinc-700">
-                {history.question}
-              </span>
-            </div>
-            <div className="w-full flex flex-row justify-end items-start pl-8 space-x-2">
-              <div className="memo-content-wrapper !w-auto flex flex-col justify-start items-start rounded-lg rounded-tr-none px-3 py-2 bg-gray-100 dark:bg-zinc-700">
-                <div className="memo-content-text">{marked(history.answer)}</div>
-              </div>
-              <Icon.Bot className="mt-2 flex-shrink-0 mr-1 w-6 h-auto opacity-80" />
-            </div>
-          </div>
-        ))}
-        {!isEnabled && (
-          <div className="w-full flex flex-col justify-center items-center mt-4 space-y-2">
-            <p>You have not set up your OpenAI API key.</p>
-            <Button onClick={() => handleGotoSystemSetting()}>Go to settings</Button>
-          </div>
-        )}
      </div>
    </>
  );
--- a/web/src/helpers/api.ts
+++ b/web/src/helpers/api.ts
@ -250,16 +250,8 @@ export function deleteIdentityProvider(id: IdentityProviderId) {
  return axios.delete(`/api/idp/${id}`);
 }

-export function postChatCompletion(prompt: string) {
-  return axios.post<ResponseObject<string>>(`/api/openai/chat-completion`, {
-    prompt,
-  });
-}
-
-export function postTextCompletion(prompt: string) {
-  return axios.post<ResponseObject<string>>(`/api/openai/text-completion`, {
-    prompt,
-  });
+export function postChatCompletion(messages: any[]) {
+  return axios.post<ResponseObject<string>>(`/api/openai/chat-completion`, messages);
 }

 export function checkOpenAIEnabled() {
--- a/web/src/store/zustand/message.ts
+++ b/web/src/store/zustand/message.ts
@ -0,0 +1,26 @@
+import { create } from "zustand";
+import { persist } from "zustand/middleware";
+
+export interface Message {
+  role: "user" | "assistant";
+  content: string;
+}
+
+interface MessageState {
+  messageList: Message[];
+  getState: () => MessageState;
+  addMessage: (message: Message) => void;
+}
+
+export const useMessageStore = create<MessageState>()(
+  persist(
+    (set, get) => ({
+      messageList: [],
+      getState: () => get(),
+      addMessage: (message: Message) => set((state) => ({ messageList: [...state.messageList, message] })),
+    }),
+    {
+      name: "message-storage",
+    }
+  )
+);
--- a/web/yarn.lock
+++ b/web/yarn.lock
@ -3058,7 +3058,7 @@ uri-js@^4.2.2:
  dependencies:
    punycode "^2.1.0"

-use-sync-external-store@^1.0.0:
+use-sync-external-store@1.2.0, use-sync-external-store@^1.0.0:
  version "1.2.0"
  resolved "https://registry.yarnpkg.com/use-sync-external-store/-/use-sync-external-store-1.2.0.tgz#7dbefd6ef3fe4e767a0cf5d7287aacfb5846928a"
  integrity sha512-eEgnFxGQ1Ife9bzYs6VLi8/4X6CObHMw9Qr9tPY43iKwsPw8xE8+EFsf/2cFZ5S3esXgpWgtSCtLNS41F+sKPA==
@ -3144,3 +3144,10 @@ yocto-queue@^0.1.0:
  version "0.1.0"
  resolved "https://registry.yarnpkg.com/yocto-queue/-/yocto-queue-0.1.0.tgz#0294eb3dee05028d31ee1a5fa2c556a6aaf10a1b"
  integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
+
+zustand@^4.3.6:
+  version "4.3.6"
+  resolved "https://registry.yarnpkg.com/zustand/-/zustand-4.3.6.tgz#ce7804eb75361af0461a2d0536b65461ec5de86f"
+  integrity sha512-6J5zDxjxLE+yukC2XZWf/IyWVKnXT9b9HUv09VJ/bwGCpKNcaTqp7Ws28Xr8jnbvnZcdRaidztAPsXFBIqufiw==
+  dependencies:
+    use-sync-external-store "1.2.0"