From 980644f13c9420f58328d18d12fc1e51a31a1234 Mon Sep 17 00:00:00 2001
From: arcticfly <41524992+arcticfly@users.noreply.github.com>
Date: Thu, 17 Aug 2023 21:02:27 -0700
Subject: [PATCH] Support vicuna system message (#167)

* Support vicuna system message

* Change tags to USER and ASSISTANT
---
 .../modelProviders/openpipe-chat/frontend.ts  |  4 +-
 .../openpipe-chat/templatePrompt.ts           | 40 ++++++++++++++-----
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/app/src/modelProviders/openpipe-chat/frontend.ts b/app/src/modelProviders/openpipe-chat/frontend.ts
index 243c0c1..cbefd84 100644
--- a/app/src/modelProviders/openpipe-chat/frontend.ts
+++ b/app/src/modelProviders/openpipe-chat/frontend.ts
@@ -7,7 +7,7 @@ import {
   //   templateSystemUserAssistantPrompt,
   templateInstructionInputResponsePrompt,
   templateAiroborosPrompt,
-  templateHumanAssistantPrompt,
+  templateVicunaPrompt,
 } from "./templatePrompt";
 
 const frontendModelProvider: FrontendModelProvider<SupportedModel, OpenpipeChatOutput> = {
@@ -67,7 +67,7 @@ const frontendModelProvider: FrontendModelProvider<SupportedModel, OpenpipeChatO
       speed: "medium",
       provider: "openpipe/Chat",
       learnMoreUrl: "https://huggingface.co/lmsys/vicuna-13b-v1.5",
-      templatePrompt: templateHumanAssistantPrompt,
+      templatePrompt: templateVicunaPrompt,
     },
     "NousResearch/Nous-Hermes-llama-2-7b": {
       name: "Nous-Hermes-llama-2-7b",
diff --git a/app/src/modelProviders/openpipe-chat/templatePrompt.ts b/app/src/modelProviders/openpipe-chat/templatePrompt.ts
index 424f0fa..f9bf4a2 100644
--- a/app/src/modelProviders/openpipe-chat/templatePrompt.ts
+++ b/app/src/modelProviders/openpipe-chat/templatePrompt.ts
@@ -178,23 +178,41 @@ export const templateAiroborosPrompt = (messages: OpenpipeChatInput["messages"])
   return prompt;
 };
 
-// ### Human: your prompt here
-// ### Assistant:
-export const templateHumanAssistantPrompt = (messages: OpenpipeChatInput["messages"]) => {
+// A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.
+
+// USER: {prompt}
+// ASSISTANT:
+export const templateVicunaPrompt = (messages: OpenpipeChatInput["messages"]) => {
   const splitter = "\n";
 
-  const humanTag = "### Human: ";
-  const assistantTag = "### Assistant: ";
+  const humanTag = "USER: ";
+  const assistantTag = "ASSISTANT: ";
 
-  const formattedMessages = messages.map((message) => {
-    if (message.role === "system" || message.role === "user") {
-      return humanTag + message.content;
+  let combinedSystemMessage = "";
+  const conversationMessages = [];
+
+  for (const message of messages) {
+    if (message.role === "system") {
+      combinedSystemMessage += message.content;
+    } else if (message.role === "user") {
+      conversationMessages.push(humanTag + message.content);
     } else {
-      return assistantTag + message.content;
+      conversationMessages.push(assistantTag + message.content);
     }
-  });
+  }
 
-  let prompt = formattedMessages.join(splitter);
+  let systemMessage = "";
+
+  if (combinedSystemMessage) {
+    // If there is no user message, add a user tag to the system message
+    if (conversationMessages.find((message) => message.startsWith(humanTag))) {
+      systemMessage = `${combinedSystemMessage}\n\n`;
+    } else {
+      conversationMessages.unshift(humanTag + combinedSystemMessage);
+    }
+  }
+
+  let prompt = `${systemMessage}${conversationMessages.join(splitter)}`;
 
   // Ensure that the prompt ends with an assistant message
   const lastHumanIndex = prompt.lastIndexOf(humanTag);