fix: GeneralChatPromptWrapper output (#70)

* fix: `GeneralChatPromptWrapper` output
* test: add tests to chat prompt wrappers

src/chatWrappers/GeneralChatPromptWrapper.ts

@@ -17,9 +17,9 @@ export class GeneralChatPromptWrapper extends ChatPromptWrapper {
         systemPrompt: string, promptIndex: number, lastStopString: string | null, lastStopStringSuffix: string | null
     }) {
         if (promptIndex === 0)
-            return systemPrompt + `\n\n### ${this._instructionName}:\n\n` + prompt + `\n\n### ${this._responseName}:\n\n`;
+            return systemPrompt + `\n\n### ${this._instructionName}:\n` + prompt + `\n\n### ${this._responseName}:\n`;
 
-        return this._getPromptPrefix(lastStopString, lastStopStringSuffix) + prompt + `\n\n### ${this._responseName}:\n\n`;
+        return this._getPromptPrefix(lastStopString, lastStopStringSuffix) + prompt + `\n\n### ${this._responseName}:\n`;
     }
 
     public override getStopStrings(): string[] {
@@ -42,9 +42,9 @@ export class GeneralChatPromptWrapper extends ChatPromptWrapper {
                 ? lastStopStringSuffix
                 : ((lastStopString ?? "") + (lastStopStringSuffix ?? "")),
             [
-                `\n\n### ${this._instructionName}:\n\n`,
-                `### ${this._instructionName}:\n\n`
+                `\n\n### ${this._instructionName}:\n`,
+                `### ${this._instructionName}:\n`
             ]
-        ) ?? `\n\n### ${this._instructionName}:\n\n`;
+        ) ?? `\n\n### ${this._instructionName}:\n`;
     }
 }

test/standalone/chatWrappers/ChatMLChatPromptWrapper.test.ts

@@ -0,0 +1,127 @@
+import {describe, expect, test} from "vitest";
+import {generateContextTextFromConversationHistory} from "../../../src/chatWrappers/generateContextTextFromConversationHistory.js";
+import {ChatMLChatPromptWrapper, ConversationInteraction} from "../../../src/index.js";
+
+
+describe("ChatMLChatPromptWrapper", () => {
+    const conversationHistory: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }];
+    const conversationHistory2: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }, {
+        prompt: "How are you?",
+        response: "I'm good, how are you?"
+    }];
+
+    test("should generate valid output for default roles", () => {
+        const chatWrapper = new ChatMLChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>"
+        `);
+
+        const chatWrapper2 = new ChatMLChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>
+          <|im_start|>user
+          How are you?<|im_end|>
+          <|im_start|>assistant
+          I'm good, how are you?<|im_end|>"
+        `);
+
+        const chatWrapper3 = new ChatMLChatPromptWrapper();
+        const {text: response3, stopStringSuffix, stopString} = generateContextTextFromConversationHistory(chatWrapper3, conversationHistory);
+
+        const newPrompt = conversationHistory2[1].prompt;
+        const wrappedNewPrompt = chatWrapper3.wrapPrompt(newPrompt, {
+            systemPrompt: response3,
+            promptIndex: 1,
+            lastStopString: stopString,
+            lastStopStringSuffix: stopStringSuffix
+        });
+
+        expect(response3).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>"
+        `);
+
+        expect(wrappedNewPrompt).toMatchInlineSnapshot(`
+          "
+          <|im_start|>user
+          How are you?<|im_end|>
+          <|im_start|>assistant
+          "
+        `);
+
+        expect(response3 + wrappedNewPrompt).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>
+          <|im_start|>user
+          How are you?<|im_end|>
+          <|im_start|>assistant
+          "
+        `);
+    });
+
+    test("should generate valid output for custom system prompt", () => {
+        const chatWrapper = new ChatMLChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          Below is an instruction the describes a task, Write a response the appropriately completes the request.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>"
+        `);
+
+        const chatWrapper2 = new ChatMLChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response2).toMatchInlineSnapshot(`
+          "<|im_start|>system
+          Below is an instruction the describes a task, Write a response the appropriately completes the request.<|im_end|>
+          <|im_start|>user
+          Hi there!<|im_end|>
+          <|im_start|>assistant
+          Hello!<|im_end|>
+          <|im_start|>user
+          How are you?<|im_end|>
+          <|im_start|>assistant
+          I'm good, how are you?<|im_end|>"
+        `);
+    });
+});

test/standalone/chatWrappers/FalconChatPromptWrapper.test.ts

@@ -0,0 +1,140 @@
+import {describe, expect, test} from "vitest";
+import {generateContextTextFromConversationHistory} from "../../../src/chatWrappers/generateContextTextFromConversationHistory.js";
+import {ConversationInteraction, FalconChatPromptWrapper} from "../../../src/index.js";
+
+
+describe("FalconChatPromptWrapper", () => {
+    const conversationHistory: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }];
+    const conversationHistory2: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }, {
+        prompt: "How are you?",
+        response: "I'm good, how are you?"
+    }];
+
+    test("should generate valid output for default roles", () => {
+        const chatWrapper = new FalconChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          User: Hi there!
+          Assistant: Hello!
+          User: "
+        `);
+
+        const chatWrapper2 = new FalconChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          User: Hi there!
+          Assistant: Hello!
+          User: How are you?
+          Assistant: I'm good, how are you?
+          User: "
+        `);
+
+        const chatWrapper3 = new FalconChatPromptWrapper();
+        const {text: response3, stopStringSuffix, stopString} = generateContextTextFromConversationHistory(chatWrapper3, conversationHistory);
+
+        const newPrompt = conversationHistory2[1].prompt;
+        const wrappedNewPrompt = chatWrapper3.wrapPrompt(newPrompt, {
+            systemPrompt: response3,
+            promptIndex: 1,
+            lastStopString: stopString,
+            lastStopStringSuffix: stopStringSuffix
+        });
+
+        expect(response3).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          User: Hi there!
+          Assistant: Hello!
+          User: "
+        `);
+
+        expect(wrappedNewPrompt).toMatchInlineSnapshot(`
+          "How are you?
+          Assistant: "
+        `);
+
+        expect(response3 + wrappedNewPrompt).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          User: Hi there!
+          Assistant: Hello!
+          User: How are you?
+          Assistant: "
+        `);
+    });
+
+    test("should generate valid output for custom roles", () => {
+        const chatWrapper = new FalconChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          Instruction: Hi there!
+          Response: Hello!
+          Instruction: "
+        `);
+
+        const chatWrapper2 = new FalconChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          Instruction: Hi there!
+          Response: Hello!
+          Instruction: How are you?
+          Response: I'm good, how are you?
+          Instruction: "
+        `);
+    });
+
+    test("should generate valid output for custom system prompt", () => {
+        const chatWrapper = new FalconChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response).toMatchInlineSnapshot(`
+          "Below is an instruction the describes a task, Write a response the appropriately completes the request.
+          User: Hi there!
+          Assistant: Hello!
+          User: "
+        `);
+
+        const chatWrapper2 = new FalconChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response2).toMatchInlineSnapshot(`
+          "Below is an instruction the describes a task, Write a response the appropriately completes the request.
+          Instruction: Hi there!
+          Response: Hello!
+          Instruction: How are you?
+          Response: I'm good, how are you?
+          Instruction: "
+        `);
+    });
+});

test/standalone/chatWrappers/GeneralChatPromptWrapper.test.ts

@@ -0,0 +1,198 @@
+import {describe, expect, test} from "vitest";
+import {generateContextTextFromConversationHistory} from "../../../src/chatWrappers/generateContextTextFromConversationHistory.js";
+import {ConversationInteraction, GeneralChatPromptWrapper} from "../../../src/index.js";
+
+
+describe("GeneralChatPromptWrapper", () => {
+    const conversationHistory: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }];
+    const conversationHistory2: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }, {
+        prompt: "How are you?",
+        response: "I'm good, how are you?"
+    }];
+
+    test("should generate valid output for default roles", () => {
+        const chatWrapper = new GeneralChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Human:
+          Hi there!
+
+          ### Assistant:
+          Hello!
+
+          ### Human"
+        `);
+
+        const chatWrapper2 = new GeneralChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Human:
+          Hi there!
+
+          ### Assistant:
+          Hello!
+
+          ### Human:
+          How are you?
+
+          ### Assistant:
+          I'm good, how are you?
+
+          ### Human"
+        `);
+
+        const chatWrapper3 = new GeneralChatPromptWrapper();
+        const {text: response3, stopStringSuffix, stopString} = generateContextTextFromConversationHistory(chatWrapper3, conversationHistory);
+
+        const newPrompt = conversationHistory2[1].prompt;
+        const wrappedNewPrompt = chatWrapper3.wrapPrompt(newPrompt, {
+            systemPrompt: response3,
+            promptIndex: 1,
+            lastStopString: stopString,
+            lastStopStringSuffix: stopStringSuffix
+        });
+
+        expect(response3).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Human:
+          Hi there!
+
+          ### Assistant:
+          Hello!
+
+          ### Human"
+        `);
+
+        expect(wrappedNewPrompt).toMatchInlineSnapshot(`
+          ":
+          How are you?
+
+          ### Assistant:
+          "
+        `);
+
+        expect(response3 + wrappedNewPrompt).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Human:
+          Hi there!
+
+          ### Assistant:
+          Hello!
+
+          ### Human:
+          How are you?
+
+          ### Assistant:
+          "
+        `);
+    });
+
+    test("should generate valid output for custom roles", () => {
+        const chatWrapper = new GeneralChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Instruction:
+          Hi there!
+
+          ### Response:
+          Hello!
+
+          ### Instruction"
+        `);
+
+        const chatWrapper2 = new GeneralChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+
+          ### Instruction:
+          Hi there!
+
+          ### Response:
+          Hello!
+
+          ### Instruction:
+          How are you?
+
+          ### Response:
+          I'm good, how are you?
+
+          ### Instruction"
+        `);
+    });
+
+    test("should generate valid output for custom system prompt", () => {
+        const chatWrapper = new GeneralChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response).toMatchInlineSnapshot(`
+          "Below is an instruction the describes a task, Write a response the appropriately completes the request.
+
+          ### Human:
+          Hi there!
+
+          ### Assistant:
+          Hello!
+
+          ### Human"
+        `);
+
+        const chatWrapper2 = new GeneralChatPromptWrapper({
+            instructionName: "Instruction",
+            responseName: "Response"
+        });
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response2).toMatchInlineSnapshot(`
+          "Below is an instruction the describes a task, Write a response the appropriately completes the request.
+
+          ### Instruction:
+          Hi there!
+
+          ### Response:
+          Hello!
+
+          ### Instruction:
+          How are you?
+
+          ### Response:
+          I'm good, how are you?
+
+          ### Instruction"
+        `);
+    });
+});

test/standalone/chatWrappers/LlamaChatPromptWrapper.test.ts

@@ -0,0 +1,125 @@
+import {describe, expect, test} from "vitest";
+import {generateContextTextFromConversationHistory} from "../../../src/chatWrappers/generateContextTextFromConversationHistory.js";
+import {ConversationInteraction, LlamaChatPromptWrapper} from "../../../src/index.js";
+
+
+describe("LlamaChatPromptWrapper", () => {
+    const conversationHistory: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }];
+    const conversationHistory2: ConversationInteraction[] = [{
+        prompt: "Hi there!",
+        response: "Hello!"
+    }, {
+        prompt: "How are you?",
+        response: "I'm good, how are you?"
+    }];
+
+    test("should generate valid output for default roles", () => {
+        const chatWrapper = new LlamaChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory);
+
+        expect(response).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s>"
+        `);
+
+        const chatWrapper2 = new LlamaChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2);
+
+        expect(response2).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s><s>[INST] How are you? [/INST]
+
+          I'm good, how are you?</s>"
+        `);
+
+        const chatWrapper3 = new LlamaChatPromptWrapper();
+        const {text: response3, stopStringSuffix, stopString} = generateContextTextFromConversationHistory(chatWrapper3, conversationHistory);
+
+        const newPrompt = conversationHistory2[1].prompt;
+        const wrappedNewPrompt = chatWrapper3.wrapPrompt(newPrompt, {
+            systemPrompt: response3,
+            promptIndex: 1,
+            lastStopString: stopString,
+            lastStopStringSuffix: stopStringSuffix
+        });
+
+        expect(response3).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s>"
+        `);
+
+        expect(wrappedNewPrompt).toMatchInlineSnapshot(`
+          "<s>[INST] How are you? [/INST]
+
+          "
+        `);
+
+        expect(response3 + wrappedNewPrompt).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          You are a helpful, respectful and honest assistant. Always answer as helpfully as possible.
+          If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s><s>[INST] How are you? [/INST]
+
+          "
+        `);
+    });
+
+    test("should generate valid output for custom system prompt", () => {
+        const chatWrapper = new LlamaChatPromptWrapper();
+        const {text: response} = generateContextTextFromConversationHistory(chatWrapper, conversationHistory, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          Below is an instruction the describes a task, Write a response the appropriately completes the request.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s>"
+        `);
+
+        const chatWrapper2 = new LlamaChatPromptWrapper();
+        const {text: response2} = generateContextTextFromConversationHistory(chatWrapper2, conversationHistory2, {
+            systemPrompt: "Below is an instruction the describes a task, Write a response the appropriately completes the request."
+        });
+
+        expect(response2).toMatchInlineSnapshot(`
+          "<s>[INST] <<SYS>>
+          Below is an instruction the describes a task, Write a response the appropriately completes the request.
+          <</SYS>>
+
+          Hi there! [/INST]
+
+          Hello!</s><s>[INST] How are you? [/INST]
+
+          I'm good, how are you?</s>"
+        `);
+    });
+});