From 642fc46afd760d51385260ed1b1dc29226d2d32a Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sat, 16 Mar 2024 20:47:54 +0800
Subject: [PATCH] feat: support virtual message model designation (#129)

---
 adapter/midjourney/chat.go                    | 15 +++++++------
 app/src/components/markdown/Link.tsx          |  4 ++--
 .../components/markdown/VirtualMessage.tsx    | 22 ++++++++++++++++---
 app/src/store/chat.ts                         |  4 ++--
 4 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/adapter/midjourney/chat.go b/adapter/midjourney/chat.go
index ef4e5a3c..27c988b1 100644
--- a/adapter/midjourney/chat.go
+++ b/adapter/midjourney/chat.go
@@ -133,14 +133,15 @@ func (c *ChatInstance) CreateStreamChatRequest(props *adaptercommon.ChatProps, c
 		return err
 	}
 
-	return c.CallbackActions(form, callback)
+	return c.CallbackActions(props, form, callback)
 }
 
-func toVirtualMessage(message string) string {
-	return "https://chatnio.virtual" + strings.Replace(message, " ", "-", -1)
+func toVirtualMessage(message string, model string) string {
+	prompt := strings.Replace(message, " ", "-", -1)
+	return fmt.Sprintf("https://chatnio.virtual%s::%s", prompt, model)
 }
 
-func (c *ChatInstance) CallbackActions(form *StorageForm, callback globals.Hook) error {
+func (c *ChatInstance) CallbackActions(props *adaptercommon.ChatProps, form *StorageForm, callback globals.Hook) error {
 	if form.Action == UpscaleAction {
 		return nil
 	}
@@ -148,14 +149,14 @@ func (c *ChatInstance) CallbackActions(form *StorageForm, callback globals.Hook)
 	actions := utils.Range(1, maxActions+1)
 
 	upscale := strings.Join(utils.Each(actions, func(index int) string {
-		return fmt.Sprintf("[U%d](%s)", index, toVirtualMessage(fmt.Sprintf("/UPSCALE %s %d", form.Task, index)))
+		return fmt.Sprintf("[U%d](%s)", index, toVirtualMessage(fmt.Sprintf("/UPSCALE %s %d", form.Task, index), props.OriginalModel))
 	}), " ")
 
 	variation := strings.Join(utils.Each(actions, func(index int) string {
-		return fmt.Sprintf("[V%d](%s)", index, toVirtualMessage(fmt.Sprintf("/VARIATION %s %d", form.Task, index)))
+		return fmt.Sprintf("[V%d](%s)", index, toVirtualMessage(fmt.Sprintf("/VARIATION %s %d", form.Task, index), props.OriginalModel))
 	}), " ")
 
-	reroll := fmt.Sprintf("[REROLL](%s)", toVirtualMessage(fmt.Sprintf("/REROLL %s", form.Task)))
+	reroll := fmt.Sprintf("[REROLL](%s)", toVirtualMessage(fmt.Sprintf("/REROLL %s", form.Task), props.OriginalModel))
 
 	return callback(&globals.Chunk{
 		Content: fmt.Sprintf("\n\n%s\n\n%s\n\n%s\n", upscale, variation, reroll),
diff --git a/app/src/components/markdown/Link.tsx b/app/src/components/markdown/Link.tsx
index c0c8c884..019c1b6a 100644
--- a/app/src/components/markdown/Link.tsx
+++ b/app/src/components/markdown/Link.tsx
@@ -30,8 +30,8 @@ export default function ({ href, children }: LinkProps) {
   const url: string = href?.toString() || "";
 
   if (url.startsWith("https://chatnio.virtual")) {
-    const message = url.slice(23).replace(/-/g, " ");
-    const prefix = message.split(" ")[0];
+    const message = url.slice(23);
+    const prefix = message.split("-")[0];
 
     return (
       <VirtualMessage message={message} prefix={prefix}>
diff --git a/app/src/components/markdown/VirtualMessage.tsx b/app/src/components/markdown/VirtualMessage.tsx
index c756b190..05aab1c7 100644
--- a/app/src/components/markdown/VirtualMessage.tsx
+++ b/app/src/components/markdown/VirtualMessage.tsx
@@ -1,6 +1,10 @@
 import React, { useState } from "react";
 import { useTranslation } from "react-i18next";
-import { useMessageActions, useWorking } from "@/store/chat.ts";
+import {
+  useConversationActions,
+  useMessageActions,
+  useWorking,
+} from "@/store/chat.ts";
 import {
   Dialog,
   DialogContent,
@@ -87,15 +91,24 @@ type VirtualMessageProps = {
   children: React.ReactNode;
 };
 
+function parseMessage(message: string): { prompt: string; model: string } {
+  const [prompt, ...rest] = message.split("::");
+  const model = rest.join(" ");
+  return { prompt: prompt.replace(/-/g, " "), model };
+}
+
 export function VirtualMessage({
   message,
   prefix,
   children,
 }: VirtualMessageProps) {
   const { t } = useTranslation();
+  const { selected } = useConversationActions();
   const { send: sendAction } = useMessageActions();
   const working = useWorking();
 
+  const { prompt, model } = parseMessage(message);
+
   return (
     <Dialog>
       <DialogTrigger asChild>
@@ -113,7 +126,7 @@ export function VirtualMessage({
           <DialogDescription className={`pb-2`}>
             {t("chat.send-message-desc")}
           </DialogDescription>
-          <VirtualPrompt message={message} prefix={prefix}>
+          <VirtualPrompt message={prompt} prefix={prefix}>
             {children}
           </VirtualPrompt>
         </DialogHeader>
@@ -123,7 +136,10 @@ export function VirtualMessage({
           </DialogClose>
           <DialogClose
             disabled={working}
-            onClick={async () => await sendAction(message)}
+            onClick={async () => {
+              selected(model);
+              await sendAction(prompt, model);
+            }}
             asChild
           >
             <Button variant={`default`}>
diff --git a/app/src/store/chat.ts b/app/src/store/chat.ts
index c823af31..a58bc049 100644
--- a/app/src/store/chat.ts
+++ b/app/src/store/chat.ts
@@ -514,7 +514,7 @@ export function useMessageActions() {
   const repetition_penalty = useSelector(repetitionPenaltySelector);
 
   return {
-    send: async (message: string) => {
+    send: async (message: string, using_model?: string) => {
       if (current === -1 && conversations[-1].messages.length === 0) {
         // preflight history if it's a new conversation
         dispatch(preflightHistory(message));
@@ -533,7 +533,7 @@ export function useMessageActions() {
         type: "chat",
         message,
         web,
-        model,
+        model: using_model || model,
         context: history,
         ignore_context: !context,
         max_tokens,