UNPKG

@lobehub/chat

Version:

Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.

116 lines (97 loc) 4.41 kB
import { describe, expect, it, vi } from 'vitest'; import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent'; import { ChatMessage } from '@/types/message'; import { OpenAIChatMessage } from '@/types/openai/chat'; import { encodeAsync } from '@/utils/tokenizer'; import * as tokenizerObj from '@/utils/tokenizer'; import { chatHelpers } from './helpers'; // Mock encodeAsync function vi.mock('@/utils/tokenizer', () => ({ encodeAsync: vi.fn((text) => Promise.resolve(text.length)), })); describe('chatHelpers', () => { describe('getMessagesTokenCount', () => { it('returns token count for a list of messages', async () => { const messages = [{ content: 'Hello' }, { content: 'World' }] as OpenAIChatMessage[]; const tokenCount = await chatHelpers.getMessagesTokenCount(messages); expect(tokenCount).toBe('HelloWorld'.length); }); it('returns 0 for an empty array', async () => { const tokenCount = await chatHelpers.getMessagesTokenCount([]); expect(tokenCount).toBe(0); }); it('handles messages with empty content', async () => { const messages = [ { content: 'Hello' }, { content: '' }, { content: 'World' }, ] as OpenAIChatMessage[]; const tokenCount = await chatHelpers.getMessagesTokenCount(messages); expect(tokenCount).toBe('HelloWorld'.length); }); it('throws an error when encodeAsync fails', async () => { vi.spyOn(tokenizerObj, 'encodeAsync').mockRejectedValue(new Error('Test error')); await expect( chatHelpers.getMessagesTokenCount([{ content: 'Hello' }] as OpenAIChatMessage[]), ).rejects.toThrow('Test error'); }); }); describe('getMessageById', () => { const messages = [ { id: '1', content: 'Hello' }, { id: '2', content: 'World' }, ] as ChatMessage[]; it('finds a message by id', () => { const message = chatHelpers.getMessageById(messages, '1'); expect(message).toEqual({ id: '1', content: 'Hello' }); }); it('returns undefined for an invalid id', () => { const message = chatHelpers.getMessageById(messages, '3'); expect(message).toBeUndefined(); }); it('returns undefined for an empty array', () => { const message = chatHelpers.getMessageById([], '1'); expect(message).toBeUndefined(); }); }); describe('getSlicedMessages', () => { const messages = [ { id: '1', content: 'First' }, { id: '2', content: 'Second' }, { id: '3', content: 'Third' }, ] as ChatMessage[]; it('returns all messages if history is disabled', () => { const config = { enableHistoryCount: false, historyCount: undefined } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages(messages, config); expect(slicedMessages).toEqual(messages); }); it('returns last N messages based on historyCount', () => { const config = { enableHistoryCount: true, historyCount: 2 } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages(messages, config); expect(slicedMessages).toEqual([ { id: '2', content: 'Second' }, { id: '3', content: 'Third' }, ]); }); it('returns empty array when historyCount is negative', () => { const config = { enableHistoryCount: true, historyCount: -1 } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages(messages, config); expect(slicedMessages).toEqual([]); }); it('returns all messages if historyCount exceeds the array length', () => { const config = { enableHistoryCount: true, historyCount: 5 } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages(messages, config); expect(slicedMessages).toEqual(messages); }); it('returns an empty array for an empty message array', () => { const config = { enableHistoryCount: true, historyCount: 2 } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages([], config); expect(slicedMessages).toEqual([]); }); it('returns an empty array when historyCount is zero', () => { const config = { enableHistoryCount: true, historyCount: 0 } as LobeAgentChatConfig; const slicedMessages = chatHelpers.getSlicedMessages(messages, config); expect(slicedMessages).toEqual([]); }); }); });