llm.test.ts

const prefixMessagesWithUsername = (messages: ChatCompletionMessageParam[]) => {
  return messages.map((message) => {
    if (message.role === 'user' && message.name) {
      return {
        ...message,
        content: `${message.name}: ${message.content}`,
      };
    }
    return message;
  });
};

test('multi-user chat using OPENAI API', async () => {
  const PREFIX_USERNAME = true;
  const openAi = new OpenAI({
    apiKey: process.env.OPENROUTER_KEY!,
    baseURL: 'https://openrouter.ai/api/v1',
  });

  const messages: ChatCompletionMessageParam[] = [
    {
      role: 'system',
      content:
        'You are a helpful assistant that help user to chat with each other.',
    },
    {
      role: 'user',
      content: 'Hello, I am John.',
      name: 'John',
    },
    {
      role: 'user',
      content: 'Hi, I am Alice.',
      name: 'Alice',
    },
    {
      role: 'user',
      content: 'How old are you?',
      name: 'John',
    },
    {
      role: 'user',
      content: 'I am 25 years old.',
      name: 'Alice',
    },
    {
      role: 'user',
      content: 'How old are you?',
      name: 'Alice',
    },
    {
      role: 'user',
      content: 'I am 22 years old.',
      name: 'John',
    },
    {
      role: 'user',
      content: 'How old is John, assistant? Just say the age, nothing else',
      name: 'Alice',
    },
  ];
  const res = await openAi.chat.completions.create({
    model: 'anthropic/claude-3-haiku',
    messages: PREFIX_USERNAME ? prefixMessagesWithUsername(messages) : messages,
  });
  expect(res.choices[0]?.message.content).toBe('22');
  messages.push(
    {
      role: 'assistant',
      content: '22',
    },
    {
      role: 'user',
      content:
        'Who is 25 years old in the 2 of us? Just say the name, do not explain',
      name: 'John',
    },
  );
  const res2 = await openAi.chat.completions.create({
    model: 'anthropic/claude-3-haiku',
    messages: PREFIX_USERNAME ? prefixMessagesWithUsername(messages) : messages,
  });
  expect(res2.choices[0]?.message.content).toBe('Alice');
});