promptfoo--promptfoo/test/providers/ai21.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { fetchWithCache } from '../../src/cache';
import logger from '../../src/logger';
import { AI21ChatCompletionProvider } from '../../src/providers/ai21';
import { mockProcessEnv } from '../util/utils';

vi.mock('../../src/cache', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    fetchWithCache: vi.fn(),
  };
});
vi.mock('../../src/logger', () => ({
  default: {
    debug: vi.fn(),
    error: vi.fn(),
    info: vi.fn(),
    warn: vi.fn(),
  },
}));

describe('AI21ChatCompletionProvider', () => {
  let restoreEnv: () => void;

  beforeEach(() => {
    vi.clearAllMocks();
    vi.resetAllMocks();
    restoreEnv = mockProcessEnv({ AI21_API_KEY: undefined });
  });

  afterEach(() => {
    restoreEnv();
    vi.restoreAllMocks();
  });

  it('should construct with valid model name', () => {
    const provider = new AI21ChatCompletionProvider('jamba-mini');
    expect(provider.modelName).toBe('jamba-mini');
  });

  it('should warn when constructing with unknown model', () => {
    const mockWarn = vi.spyOn(logger, 'warn').mockImplementation(function () {});
    new AI21ChatCompletionProvider('unknown-model');
    expect(mockWarn).toHaveBeenCalledWith(expect.stringContaining('unknown-model'));
    mockWarn.mockRestore();
  });

  it('should get API key from config', () => {
    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });
    expect(provider.getApiKey()).toBe('test-key');
  });

  it('should get API key from environment variable', () => {
    const restoreApiKey = mockProcessEnv({ AI21_API_KEY: 'env-key' });
    try {
      const provider = new AI21ChatCompletionProvider('jamba-mini');
      expect(provider.getApiKey()).toBe('env-key');
    } finally {
      restoreApiKey();
    }
  });

  it('should get API URL from config', () => {
    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiBaseUrl: 'https://custom-api.ai21.com' },
    });
    expect(provider.getApiUrl()).toBe('https://custom-api.ai21.com');
  });

  it('should get default API URL when not configured', () => {
    const provider = new AI21ChatCompletionProvider('jamba-mini');
    expect(provider.getApiUrl()).toBe('https://api.ai21.com/studio/v1');
  });

  it('should throw error when API key is not set', async () => {
    const provider = new AI21ChatCompletionProvider('jamba-mini');
    await expect(provider.callApi('test prompt')).rejects.toThrow('AI21 API key is not set');
  });

  it('should handle successful API call', async () => {
    const mockResponse = {
      data: {
        choices: [
          {
            message: {
              content: 'test response',
            },
          },
        ],
        usage: {
          total_tokens: 10,
          prompt_tokens: 5,
          completion_tokens: 5,
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('test prompt');
    expect(result.output).toBe('test response');
    expect(result.tokenUsage).toEqual({
      total: 10,
      prompt: 5,
      completion: 5,
    });
  });

  it('should preserve explicit zero for top_p', async () => {
    const mockResponse = {
      data: {
        choices: [{ message: { content: 'test response' } }],
        usage: { total_tokens: 10, prompt_tokens: 5, completion_tokens: 5 },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key', top_p: 0 },
    });

    await provider.callApi('test prompt');

    const callArgs = vi.mocked(fetchWithCache).mock.calls[0]!;
    const body = JSON.parse((callArgs[1] as RequestInit).body as string);
    expect(body.top_p).toBe(0);
  });

  it('should handle API error response', async () => {
    const mockResponse = {
      data: {
        error: 'API error message',
      },
      cached: false,
      status: 400,
      statusText: 'Bad Request',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('test prompt');
    expect(result.error).toBe('API call error: API error message');
  });

  it('should handle malformed API response', async () => {
    const mockResponse = {
      data: {
        choices: [],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('test prompt');
    expect(result.error).toContain('Malformed response data');
  });

  it('should handle network errors', async () => {
    vi.mocked(fetchWithCache).mockRejectedValue(new Error('Network error'));

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('test prompt');
    expect(result.error).toBe('API call error: Error: Network error');
  });

  it('should calculate cost correctly', async () => {
    const mockResponse = {
      data: {
        choices: [
          {
            message: {
              content: 'test response',
            },
          },
        ],
        usage: {
          total_tokens: 10,
          prompt_tokens: 5,
          completion_tokens: 5,
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('test prompt');
    expect(result.cost).toBeDefined();
  });

  it('should preserve an explicit max_tokens value of 0', async () => {
    const mockResponse = {
      data: {
        choices: [
          {
            message: {
              content: 'test response',
            },
          },
        ],
        usage: {
          total_tokens: 10,
          prompt_tokens: 5,
          completion_tokens: 5,
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key', max_tokens: 0 },
    });

    await provider.callApi('test prompt');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      expect.any(String),
      expect.objectContaining({
        body: expect.stringContaining('"max_tokens":0'),
      }),
      expect.any(Number),
    );
  });

  it('invokes fetchWithCache once per call site even for duplicate provider configs', async () => {
    const mockResponse = {
      data: {
        choices: [{ message: { content: 'test response' } }],
        usage: { total_tokens: 10, prompt_tokens: 5, completion_tokens: 5 },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

    const provider1 = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });
    const provider2 = new AI21ChatCompletionProvider('jamba-mini', {
      config: { apiKey: 'test-key' },
    });

    await Promise.all([provider1.callApi('test prompt'), provider2.callApi('test prompt')]);

    // Each provider call delegates to fetchWithCache; the cache layer itself
    // is responsible for collapsing identical requests, not the provider.
    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(2);
  });
});