multi_round_chat.test.ts

import { describe, expect, test } from '@jest/globals';

import {
    ChatCompletionMessageParam,
    ChatCompletionRequest,
    ChatCompletionUserMessageParam,
} from "../src/openai_api_protocols/chat_completion";
import { MLCEngine } from '../src/engine';
import { Conversation, compareConversationObject } from '../src/conversation';
import { ChatConfig, Role } from '../src/config';

const configStr = "{" +
    "  \"conv_template\": {" +
    "    \"name\": \"llama-2\"," +
    "    \"system_template\": \"[INST] <<SYS>>\\n{system_message}\\n<</SYS>>\\n\\n\"," +
    "    \"system_message\": \"You are a helpful, respectful and honest assistant.\"," +
    "    \"system_prefix_token_ids\": [" +
    "      1" +
    "    ]," +
    "    \"add_role_after_system_message\": false," +
    "    \"roles\": {" +
    "      \"user\": \"[INST]\"," +
    "      \"assistant\": \"[/INST]\"," +
    "      \"tool\": \"[INST]\"" +
    "    }," +
    "    \"role_templates\": {" +
    "      \"user\": \"{user_message}\"," +
    "      \"assistant\": \"{assistant_message}\"," +
    "      \"tool\": \"{tool_message}\"" +
    "    }," +
    "    \"messages\": []," +
    "    \"seps\": [" +
    "      \" \"" +
    "    ]," +
    "    \"role_content_sep\": \" \"," +
    "    \"role_empty_sep\": \" \"," +
    "    \"stop_str\": [" +
    "      \"[INST]\"" +
    "    ]," +
    "    \"stop_token_ids\": [" +
    "      2" +
    "    ]," +
    "    \"function_string\": \"\"," +
    "    \"use_function_calling\": false" +
    "  }" +
    "}";

describe('Test multi-round chatting', () => {
    test('Test is multi-round', () => {
        // Setups
        const config_json = JSON.parse(configStr);
        const chatConfig = { ...config_json } as ChatConfig;
        const engine = new MLCEngine();

        // Simulate request0
        const messages: ChatCompletionMessageParam[] = [
            {
                "role": "system",
                "content": "[INST] <<SYS>>\n\nYou are a helpful, respectful and honest assistant. " +
                    "Be as happy as you can when speaking please.\n<</SYS>>\n\n "
            },
            { "role": "user", "content": "Provide me three US states." },
        ];
        const request0: ChatCompletionRequest = {
            messages: messages,
        };

        // Simulate processing of request0, appending response to convA (done by LLMChatPipeline)
        const conv0: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request0, chatConfig));
        conv0.appendMessage(Role.user, "Provide me three US states.");
        const reply0 = "California, New York, Nevada.";
        conv0.appendMessage(Role.assistant, reply0);  // simulated response

        // Simulate request1, where user maintain the chat history, appending the resposne
        const newMessages = [...messages];
        newMessages.push({ "role": "assistant", "content": reply0 });
        newMessages.push({ "role": "user", "content": "Two more please" });  // next input

        const request1: ChatCompletionRequest = {
            messages: newMessages,
        }
        const conv1: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request1, chatConfig));

        expect(compareConversationObject(conv0, conv1)).toBe(true);
    });

    test('Test is NOT multi-round due to multiple new inputs', () => {
        // Setups
        const config_json = JSON.parse(configStr);
        const chatConfig = { ...config_json } as ChatConfig;
        const engine = new MLCEngine();

        // Simulate request0
        const messages: ChatCompletionMessageParam[] = [
            {
                "role": "system",
                "content": "[INST] <<SYS>>\n\nYou are a helpful, respectful and honest assistant. " +
                    "Be as happy as you can when speaking please.\n<</SYS>>\n\n "
            },
            { "role": "user", "content": "Provide me three US states." },
        ];
        const request0: ChatCompletionRequest = {
            messages: messages,
        };

        // Simulate processing of request0, appending response to convA (done by LLMChatPipeline)
        const conv0: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request0, chatConfig));
        conv0.appendMessage(Role.user, "Provide me three US states.");
        const reply0 = "California, New York, Nevada.";
        conv0.appendMessage(Role.assistant, reply0);  // simulated response

        // Simulate request1, where user maintain the chat history, appending the resposne
        const newMessages = [...messages];
        newMessages.push({ "role": "assistant", "content": reply0 });
        newMessages.push({ "role": "user", "content": "Two more please" });  // next input

        // Code above same as previous tests
        // Add one more round of chat history
        newMessages.push({ "role": "assistant", "content": "Pennsylvania, Florida" });  // next response
        newMessages.push({ "role": "user", "content": "Thank you!" });  // next input

        const request1: ChatCompletionRequest = {
            messages: newMessages,
        }
        const conv1: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request1, chatConfig));

        expect(compareConversationObject(conv0, conv1)).toBe(false);
    });

    test('Test is NOT multi-round due to change in system prompt', () => {
        // Setups
        const config_json = JSON.parse(configStr);
        const chatConfig = { ...config_json } as ChatConfig;
        const engine = new MLCEngine();

        // Simulate request0
        const messages: ChatCompletionMessageParam[] = [
            {
                "role": "system",
                "content": "[INST] <<SYS>>\n\nYou are a helpful, respectful and honest assistant. " +
                    "Be as happy as you can when speaking please.\n<</SYS>>\n\n "
            },
            { "role": "user", "content": "Provide me three US states." },
        ];
        const request0: ChatCompletionRequest = {
            messages: messages,
        };

        // Simulate processing of request0, appending response to convA (done by LLMChatPipeline)
        const conv0: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request0, chatConfig));
        conv0.appendMessage(Role.user, "Provide me three US states.");
        const reply0 = "California, New York, Nevada.";
        conv0.appendMessage(Role.assistant, reply0);  // simulated response

        // Simulate request1, where user maintain the chat history, appending the resposne
        const newMessages = [...messages];
        newMessages.push({ "role": "assistant", "content": reply0 });
        newMessages.push({ "role": "user", "content": "Two more please" });  // next input

        // Code above same as previous tests
        // Changed system prompt, should be false
        newMessages[0].content = "No system prompt";

        const request1: ChatCompletionRequest = {
            messages: newMessages,
        }
        const conv1: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request1, chatConfig));

        expect(compareConversationObject(conv0, conv1)).toBe(false);
    });

    test('Test is NOT multi-round due to change in role name', () => {
        // Setups
        const config_json = JSON.parse(configStr);
        const chatConfig = { ...config_json } as ChatConfig;
        const engine = new MLCEngine();

        // Simulate request0
        const messages: ChatCompletionMessageParam[] = [
            {
                "role": "system",
                "content": "[INST] <<SYS>>\n\nYou are a helpful, respectful and honest assistant. " +
                    "Be as happy as you can when speaking please.\n<</SYS>>\n\n "
            },
            { "role": "user", "content": "Provide me three US states." },
        ];
        const request0: ChatCompletionRequest = {
            messages: messages,
        };

        // Simulate processing of request0, appending response to convA (done by LLMChatPipeline)
        const conv0: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request0, chatConfig));
        conv0.appendMessage(Role.user, "Provide me three US states.");
        const reply0 = "California, New York, Nevada.";
        conv0.appendMessage(Role.assistant, reply0);  // simulated response

        // Simulate request1, where user maintain the chat history, appending the resposne
        const newMessages = [...messages];
        newMessages.push({ "role": "assistant", "content": reply0 });
        newMessages.push({ "role": "user", "content": "Two more please" });  // next input

        // Code above same as previous tests
        // Changed system prompt, should be false
        (newMessages[1] as ChatCompletionUserMessageParam).name = "Bob";

        const request1: ChatCompletionRequest = {
            messages: newMessages,
        }
        const conv1: Conversation =
            ((engine as any).getConversationFromChatCompletionRequest(request1, chatConfig));

        expect(compareConversationObject(conv0, conv1)).toBe(false);
    });

})