File size: 5,167 Bytes
9705b6c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
const AnthropicClient = require('../AnthropicClient');
const HUMAN_PROMPT = '\n\nHuman:';
const AI_PROMPT = '\n\nAssistant:';
describe('AnthropicClient', () => {
let client;
const model = 'claude-2';
const parentMessageId = '1';
const messages = [
{ role: 'user', isCreatedByUser: true, text: 'Hello', messageId: parentMessageId },
{ role: 'assistant', isCreatedByUser: false, text: 'Hi', messageId: '2', parentMessageId },
{
role: 'user',
isCreatedByUser: true,
text: 'What\'s up',
messageId: '3',
parentMessageId: '2',
},
];
beforeEach(() => {
const options = {
modelOptions: {
model,
temperature: 0.7,
},
};
client = new AnthropicClient('test-api-key');
client.setOptions(options);
});
describe('setOptions', () => {
it('should set the options correctly', () => {
expect(client.apiKey).toBe('test-api-key');
expect(client.modelOptions.model).toBe(model);
expect(client.modelOptions.temperature).toBe(0.7);
});
});
describe('getSaveOptions', () => {
it('should return the correct save options', () => {
const options = client.getSaveOptions();
expect(options).toHaveProperty('modelLabel');
expect(options).toHaveProperty('promptPrefix');
});
});
describe('buildMessages', () => {
it('should handle promptPrefix from options when promptPrefix argument is not provided', async () => {
client.options.promptPrefix = 'Test Prefix from options';
const result = await client.buildMessages(messages, parentMessageId);
const { prompt } = result;
expect(prompt).toContain('Test Prefix from options');
});
it('should build messages correctly for chat completion', async () => {
const result = await client.buildMessages(messages, '2');
expect(result).toHaveProperty('prompt');
expect(result.prompt).toContain(HUMAN_PROMPT);
expect(result.prompt).toContain('Hello');
expect(result.prompt).toContain(AI_PROMPT);
expect(result.prompt).toContain('Hi');
});
it('should group messages by the same author', async () => {
const groupedMessages = messages.map((m) => ({ ...m, isCreatedByUser: true, role: 'user' }));
const result = await client.buildMessages(groupedMessages, '3');
expect(result.context).toHaveLength(1);
// Check that HUMAN_PROMPT appears only once in the prompt
const matches = result.prompt.match(new RegExp(HUMAN_PROMPT, 'g'));
expect(matches).toHaveLength(1);
groupedMessages.push({
role: 'assistant',
isCreatedByUser: false,
text: 'I heard you the first time',
messageId: '4',
parentMessageId: '3',
});
const result2 = await client.buildMessages(groupedMessages, '4');
expect(result2.context).toHaveLength(2);
// Check that HUMAN_PROMPT appears only once in the prompt
const human_matches = result2.prompt.match(new RegExp(HUMAN_PROMPT, 'g'));
const ai_matches = result2.prompt.match(new RegExp(AI_PROMPT, 'g'));
expect(human_matches).toHaveLength(1);
expect(ai_matches).toHaveLength(1);
});
it('should handle isEdited condition', async () => {
const editedMessages = [
{ role: 'user', isCreatedByUser: true, text: 'Hello', messageId: '1' },
{ role: 'assistant', isCreatedByUser: false, text: 'Hi', messageId: '2', parentMessageId },
];
const trimmedLabel = AI_PROMPT.trim();
const result = await client.buildMessages(editedMessages, '2');
expect(result.prompt.trim().endsWith(trimmedLabel)).toBeFalsy();
// Add a human message at the end to test the opposite
editedMessages.push({
role: 'user',
isCreatedByUser: true,
text: 'Hi again',
messageId: '3',
parentMessageId: '2',
});
const result2 = await client.buildMessages(editedMessages, '3');
expect(result2.prompt.trim().endsWith(trimmedLabel)).toBeTruthy();
});
it('should build messages correctly with a promptPrefix', async () => {
const promptPrefix = 'Test Prefix';
client.options.promptPrefix = promptPrefix;
const result = await client.buildMessages(messages, parentMessageId);
const { prompt } = result;
expect(prompt).toBeDefined();
expect(prompt).toContain(promptPrefix);
const textAfterPrefix = prompt.split(promptPrefix)[1];
expect(textAfterPrefix).toContain(AI_PROMPT);
const editedMessages = messages.slice(0, -1);
const result2 = await client.buildMessages(editedMessages, parentMessageId);
const textAfterPrefix2 = result2.prompt.split(promptPrefix)[1];
expect(textAfterPrefix2).toContain(AI_PROMPT);
});
it('should handle identityPrefix from options', async () => {
client.options.userLabel = 'John';
client.options.modelLabel = 'Claude-2';
const result = await client.buildMessages(messages, parentMessageId);
const { prompt } = result;
expect(prompt).toContain('Human\'s name: John');
expect(prompt).toContain('You are Claude-2');
});
});
});
|