sugarr/test/core/llm/LlmClient.test.ts

import { describe, expect, it, vi } from 'vitest';

const createMock = vi.fn();

vi.mock('openai', () => ({
  default: vi.fn().mockImplementation(() => ({
    chat: { completions: { create: createMock } },
  })),
}));

import { LlmClient } from '../../../src/core/llm/LlmClient';

describe('LlmClient', () => {
  it('sends messages/options to the OpenAI-compatible endpoint and returns content + usage', async () => {
    createMock.mockResolvedValue({
      choices: [{ message: { content: 'hello' } }],
      usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
    });

    const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1', apiKey: 'key' });
    const result = await client.chat([{ role: 'user', content: 'hi' }], { model: 'llama3', temperature: 0.7 });

    expect(result).toEqual({
      content: 'hello',
      usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
    });
    expect(createMock).toHaveBeenCalledWith({
      model: 'llama3',
      temperature: 0.7,
      max_tokens: undefined,
      messages: [{ role: 'user', content: 'hi' }],
      usage: { include: true },
    });
  });

  it('extracts OpenRouter cost from the response usage when present', async () => {
    createMock.mockResolvedValue({
      choices: [{ message: { content: 'hello' } }],
      usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15, cost: 0.00042 },
    });

    const client = new LlmClient({ baseUrl: 'https://openrouter.ai/api/v1', apiKey: 'key' });
    const result = await client.chat([{ role: 'user', content: 'hi' }], { model: 'openrouter/auto' });

    expect(result.usage).toEqual({
      promptTokens: 10,
      completionTokens: 5,
      totalTokens: 15,
      cost: 0.00042,
    });
  });

  it('extracts cachedTokens from prompt_tokens_details when present', async () => {
    createMock.mockResolvedValue({
      choices: [{ message: { content: 'hello' } }],
      usage: {
        prompt_tokens: 100,
        completion_tokens: 20,
        total_tokens: 120,
        prompt_tokens_details: { cached_tokens: 80 },
      },
    });

    const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1' });
    const result = await client.chat([{ role: 'user', content: 'hi' }], { model: 'gpt-4o-mini' });

    expect(result.usage).toEqual({
      promptTokens: 100,
      completionTokens: 20,
      totalTokens: 120,
      cachedTokens: 80,
    });
  });

  it('throws when the response has no choices', async () => {
    createMock.mockResolvedValue({ choices: [] });

    const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1' });
    await expect(client.chat([{ role: 'user', content: 'hi' }], { model: 'llama3' })).rejects.toThrow(
      /no choices/
    );
  });
});