promptfoo--promptfoo/test/commands/optimize.test.ts

import { Command } from 'commander';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { doOptimize, optimizeCommand } from '../../src/commands/optimize';
import logger from '../../src/logger';
import { optimizePromptTestSuite } from '../../src/optimizer/promptOptimizer';
import telemetry from '../../src/telemetry';
import { resolveConfigs } from '../../src/util/config/load';
import { setupEnv } from '../../src/util/index';

vi.mock('../../src/optimizer/promptOptimizer', () => ({
  optimizePromptTestSuite: vi.fn(),
}));

vi.mock('../../src/telemetry', () => ({
  default: {
    record: vi.fn(),
  },
}));

vi.mock('../../src/util/config/load', () => ({
  resolveConfigs: vi.fn(),
}));

vi.mock('../../src/util', () => ({
  printBorder: vi.fn(),
  setupEnv: vi.fn(),
}));

vi.mock('../../src/logger', () => ({
  default: {
    debug: vi.fn(),
    info: vi.fn(),
    error: vi.fn(),
  },
}));

vi.mock('../../src/util/promptfooCommand', () => ({
  promptfooCommand: vi.fn().mockReturnValue('promptfoo init'),
}));

describe('optimize command', () => {
  beforeEach(() => {
    vi.resetAllMocks();
    process.exitCode = undefined;
    vi.mocked(resolveConfigs).mockResolvedValue({
      config: {},
      testSuite: {
        providers: [],
        prompts: [{ raw: 'Prompt', label: 'Prompt' }],
        tests: [{}],
      },
      basePath: '',
    } as any);
    vi.mocked(optimizePromptTestSuite).mockResolvedValue({
      baselinePrompt: { label: 'Prompt', raw: 'Prompt', metrics: { score: 0.5 } },
      bestPrompt: { label: 'Prompt [optimized 1]', raw: 'Better prompt', metrics: { score: 0.8 } },
      improved: true,
      candidates: [{ prompt: 'Better prompt', hypothesis: 'Improve clarity.' }],
    } as any);
  });

  afterEach(() => {
    vi.resetAllMocks();
    process.exitCode = undefined;
  });

  it('uses the implicit default config path when -c is omitted', async () => {
    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
    });

    expect(setupEnv).toHaveBeenCalledWith(undefined);
    expect(resolveConfigs).toHaveBeenCalledWith({ config: ['promptfooconfig.yaml'] }, {});
    expect(optimizePromptTestSuite).toHaveBeenCalledTimes(1);
    expect(optimizePromptTestSuite).toHaveBeenCalledWith({}, expect.any(Object), {
      promptIndex: 0,
      providerIndex: 0,
      validationSplit: undefined,
    });
    expect(telemetry.record).toHaveBeenCalledWith(
      'command_used',
      expect.objectContaining({ name: 'optimize - started' }),
    );
  });

  it('loads config-defined env files when the CLI does not override them', async () => {
    vi.mocked(resolveConfigs).mockResolvedValue({
      config: {},
      testSuite: {
        providers: [],
        prompts: [{ raw: 'Prompt', label: 'Prompt' }],
        tests: [{}],
      },
      basePath: '',
      commandLineOptions: {
        envPath: '.env.optimize',
      },
    } as any);

    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
    });

    expect(setupEnv).toHaveBeenNthCalledWith(1, undefined);
    expect(setupEnv).toHaveBeenNthCalledWith(2, '.env.optimize');
  });

  it('keeps explicit CLI env files ahead of config-defined env files', async () => {
    vi.mocked(resolveConfigs).mockResolvedValue({
      config: {},
      testSuite: {
        providers: [],
        prompts: [{ raw: 'Prompt', label: 'Prompt' }],
        tests: [{}],
      },
      basePath: '',
      commandLineOptions: {
        envPath: '.env.config',
      },
    } as any);

    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
      envPath: '.env.cli',
    });

    expect(setupEnv).toHaveBeenCalledTimes(1);
    expect(setupEnv).toHaveBeenCalledWith('.env.cli');
  });

  it('passes validation split through to the optimizer', async () => {
    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
      validationSplit: 0.2,
    });

    expect(optimizePromptTestSuite).toHaveBeenCalledWith({}, expect.any(Object), {
      promptIndex: 0,
      providerIndex: 0,
      validationSplit: 0.2,
    });
  });

  it('passes explicit prompt and provider indices through to the optimizer', async () => {
    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
      promptIndex: 2,
      providerIndex: 1,
    });

    expect(optimizePromptTestSuite).toHaveBeenCalledWith({}, expect.any(Object), {
      promptIndex: 2,
      providerIndex: 1,
      validationSplit: undefined,
    });
  });

  it('requires a config path when no default config is available', async () => {
    await expect(
      doOptimize({
        defaultConfig: {},
        defaultConfigPath: undefined,
      }),
    ).rejects.toThrow('Could not find a config file.');
  });

  it('logs validation scores and a non-improving result', async () => {
    vi.mocked(optimizePromptTestSuite).mockResolvedValue({
      baselinePrompt: { label: 'Prompt', raw: 'Prompt', metrics: { score: 0.5 } },
      bestPrompt: { label: 'Prompt', raw: 'Prompt', metrics: { score: 0.5 } },
      baselineValidationPrompt: {
        label: 'Prompt',
        raw: 'Prompt',
        metrics: { score: 0.4 },
      },
      bestValidationPrompt: {
        label: 'Prompt',
        raw: 'Prompt',
        metrics: { score: 0.4 },
      },
      improved: false,
      candidates: [],
      searchTestCount: 1,
      validationSplit: 0.2,
      validationTestCount: 1,
    } as any);

    await doOptimize({
      defaultConfig: {},
      defaultConfigPath: 'promptfooconfig.yaml',
      validationSplit: 0.2,
    });

    expect(logger.info).toHaveBeenCalledWith(expect.stringContaining('Baseline validation:'));
    expect(logger.info).toHaveBeenCalledWith(expect.stringContaining('Best validation:'));
    expect(logger.info).toHaveBeenCalledWith(
      expect.stringContaining('Baseline remains strongest.'),
    );
  });

  it('parses valid CLI optimizer selections and validation split', async () => {
    const program = new Command();
    optimizeCommand(program, {}, 'promptfooconfig.yaml');

    await program.parseAsync([
      'node',
      'test',
      'optimize',
      '--prompt-index',
      '2',
      '--provider-index',
      '1',
      '--validation-split',
      '0.3',
    ]);

    expect(optimizePromptTestSuite).toHaveBeenCalledWith({}, expect.any(Object), {
      promptIndex: 2,
      providerIndex: 1,
      validationSplit: 0.3,
    });
  });

  it.each([
    ['--prompt-index', '-1', '--prompt-index must be a non-negative integer.'],
    ['--provider-index', '1.5', '--provider-index must be a non-negative integer.'],
    [
      '--validation-split',
      '0.75',
      '--validation-split must be greater than 0 and less than or equal to 0.5',
    ],
    [
      '--validation-split',
      '0.2typo',
      '--validation-split must be greater than 0 and less than or equal to 0.5',
    ],
  ])('rejects invalid %s values', async (flag, value, message) => {
    const program = new Command();
    program.exitOverride();
    optimizeCommand(program, {}, 'promptfooconfig.yaml');

    await expect(program.parseAsync(['node', 'test', 'optimize', flag, value])).rejects.toThrow(
      message,
    );
  });

  it('logs action handler failures and sets the exit code', async () => {
    vi.mocked(resolveConfigs).mockRejectedValue(new Error('config exploded'));
    const program = new Command();
    optimizeCommand(program, {}, 'promptfooconfig.yaml');

    await program.parseAsync(['node', 'test', 'optimize']);

    expect(logger.error).toHaveBeenCalledWith('config exploded');
    expect(process.exitCode).toBe(1);
  });

  it('registers a top-level optimize command', () => {
    const program = new Command();
    optimizeCommand(program, {}, 'promptfooconfig.yaml');

    const command = program.commands[0];
    expect(command.name()).toBe('optimize');
    expect(command.description()).toBe('Optimize one prompt against one configured provider');
    expect(command.options.some((option) => option.long === '--prompt-index')).toBe(true);
    expect(command.options.some((option) => option.long === '--provider-index')).toBe(true);
    expect(command.options.some((option) => option.long === '--validation-split')).toBe(true);
  });
});