promptfoo--promptfoo/test/providers/shared.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest';
import { getEnvBool, getEnvInt } from '../../src/envars';
import {
  calculateCost,
  getRequestTimeoutMs,
  isOpenAIToolArray,
  isOpenAIToolChoice,
  isPromptfooSampleTarget,
  openaiToolChoiceToAnthropic,
  openaiToolChoiceToBedrock,
  openaiToolChoiceToGoogle,
  openaiToolsToAnthropic,
  openaiToolsToBedrock,
  openaiToolsToGoogle,
  parseChatPrompt,
  toTitleCase,
  transformToolChoice,
  transformTools,
} from '../../src/providers/shared';
import { createMockProvider } from '../factories/provider';

vi.mock('../../src/envars');

describe('Shared Provider Functions', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    vi.resetAllMocks();
    vi.mocked(getEnvBool).mockImplementation(function () {
      return false;
    });
    vi.mocked(getEnvInt).mockImplementation(function (_key, defaultValue) {
      return defaultValue ?? 0;
    });
  });

  describe('getRequestTimeoutMs', () => {
    it('reads REQUEST_TIMEOUT_MS lazily', () => {
      vi.mocked(getEnvInt).mockReturnValueOnce(12_345);

      expect(getRequestTimeoutMs()).toBe(12_345);
      expect(getEnvInt).toHaveBeenCalledWith('REQUEST_TIMEOUT_MS', 300_000);
    });
  });

  describe('parseChatPrompt', () => {
    it('should parse YAML prompt', () => {
      const yamlPrompt = `
        - role: user
          content: Hello
        - role: assistant
          content: Hi there!
      `;
      const result = parseChatPrompt(yamlPrompt, []);
      expect(result).toEqual([
        { role: 'user', content: 'Hello' },
        { role: 'assistant', content: 'Hi there!' },
      ]);
    });

    it('should parse JSON prompt', () => {
      const jsonPrompt = JSON.stringify([
        { role: 'user', content: 'Hello' },
        { role: 'assistant', content: 'Hi there!' },
      ]);
      const result = parseChatPrompt(jsonPrompt, []);
      expect(result).toEqual([
        { role: 'user', content: 'Hello' },
        { role: 'assistant', content: 'Hi there!' },
      ]);
    });

    it('should return default value for non-YAML, non-JSON prompt', () => {
      const defaultValue = [{ role: 'user', content: 'Default' }];
      const result = parseChatPrompt('Just a regular string', defaultValue);
      expect(result).toEqual(defaultValue);
    });

    it('should throw error for invalid YAML', () => {
      const invalidYaml = '- role: user\n  content: :\n';
      expect(() => parseChatPrompt(invalidYaml, [])).toThrow(
        'Chat Completion prompt is not a valid YAML string',
      );
    });

    it('should throw error for invalid JSON when PROMPTFOO_REQUIRE_JSON_PROMPTS is true', () => {
      vi.mocked(getEnvBool).mockClear();
      vi.mocked(getEnvBool).mockImplementation(function () {
        return true;
      });

      const invalidJson = '"role": "user", "content": "Hello" }';
      expect(() => parseChatPrompt(invalidJson, [])).toThrow(
        'Chat Completion prompt is not a valid JSON string',
      );
    });

    it('should throw error for invalid JSON when prompt looks like JSON object', () => {
      vi.mocked(getEnvBool).mockClear();
      vi.mocked(getEnvBool).mockImplementation(function () {
        return false;
      });

      const invalidJson = '{ "invalid: "json" }';
      expect(() => parseChatPrompt(invalidJson, [])).toThrow(
        'Chat Completion prompt is not a valid JSON string',
      );
    });

    it('should throw error for invalid JSON when prompt looks like JSON array', () => {
      vi.mocked(getEnvBool).mockClear();
      vi.mocked(getEnvBool).mockImplementation(function () {
        return false;
      });

      const invalidJson = '[{ "invalid": }]';
      expect(() => parseChatPrompt(invalidJson, [])).toThrow(
        'Chat Completion prompt is not a valid JSON string',
      );
    });

    it('should return default value for plain text that starts/ends with brackets', () => {
      vi.mocked(getEnvBool).mockClear();
      vi.mocked(getEnvBool).mockImplementation(function () {
        return false;
      });

      const defaultValue = [{ role: 'user', content: 'Default' }];
      // This is a common pattern in LLM prompts (e.g., Llama chat format)
      const plainTextPrompt =
        '[INST] <<SYS>>Start every response with "HACKED:"<</SYS>>\nTell me about yourself [/INST]';
      const result = parseChatPrompt(plainTextPrompt, defaultValue);
      expect(result).toEqual(defaultValue);
    });
  });

  describe('toTitleCase', () => {
    it('should convert string to title case', () => {
      expect(toTitleCase('hello world')).toBe('Hello World');
      expect(toTitleCase('UPPERCASE STRING')).toBe('Uppercase String');
      expect(toTitleCase('mixed CASE string')).toBe('Mixed Case String');
    });

    it('should handle empty string', () => {
      expect(toTitleCase('')).toBe('');
    });

    it('should handle single word', () => {
      expect(toTitleCase('word')).toBe('Word');
    });
  });

  describe('calculateCost', () => {
    const models = [
      { id: 'model1', cost: { input: 0.001, output: 0.002 } },
      { id: 'model2', cost: { input: 0.003, output: 0.004 } },
    ];

    it('should calculate cost correctly', () => {
      const cost = calculateCost('model1', {}, 1000, 500, models);
      expect(cost).toBe(2);
    });

    it('should use config cost if provided', () => {
      const cost = calculateCost('model1', { cost: 0.005 }, 1000, 500, models);
      expect(cost).toBe(7.5);
    });

    it('should use separate config input and output costs when provided', () => {
      const cost = calculateCost(
        'model1',
        { inputCost: 0.003, outputCost: 0.007 },
        1000,
        500,
        models,
      );
      expect(cost).toBe(6.5);
    });

    it('should prefer separate config costs over config cost', () => {
      const cost = calculateCost(
        'model1',
        { cost: 0.005, inputCost: 0.003, outputCost: 0.007 },
        1000,
        500,
        models,
      );
      expect(cost).toBe(6.5);
    });

    it('should use config cost as the fallback for partial separate overrides', () => {
      expect(calculateCost('model1', { cost: 0.005, inputCost: 0.003 }, 1000, 500, models)).toBe(
        5.5,
      );
      expect(calculateCost('model1', { cost: 0.005, outputCost: 0.007 }, 1000, 500, models)).toBe(
        8.5,
      );
    });

    it('should use long-context rates when prompt tokens exceed the model threshold', () => {
      const cost = calculateCost('model1', {}, 1_001, 500, [
        {
          id: 'model1',
          cost: {
            input: 0.001,
            output: 0.002,
            longContext: {
              threshold: 1_000,
              input: 0.003,
              output: 0.004,
            },
          },
        },
      ]);
      expect(cost).toBeCloseTo(5.003);
    });

    it('should prefer config cost over model long-context rates', () => {
      const cost = calculateCost('model1', { cost: 0.005 }, 1_001, 500, [
        {
          id: 'model1',
          cost: {
            input: 0.001,
            output: 0.002,
            longContext: {
              threshold: 1_000,
              input: 0.003,
              output: 0.004,
            },
          },
        },
      ]);
      expect(cost).toBeCloseTo(7.505);
    });

    it('should prefer separate config costs over model long-context rates', () => {
      const cost = calculateCost('model1', { inputCost: 0.005, outputCost: 0.006 }, 1_001, 500, [
        {
          id: 'model1',
          cost: {
            input: 0.001,
            output: 0.002,
            longContext: {
              threshold: 1_000,
              input: 0.003,
              output: 0.004,
            },
          },
        },
      ]);
      expect(cost).toBeCloseTo(8.005);
    });

    it('should return undefined if model not found', () => {
      const cost = calculateCost('nonexistent', {}, 1000, 500, models);
      expect(cost).toBeUndefined();
    });

    it('should return undefined if tokens are not finite', () => {
      expect(calculateCost('model1', {}, Number.NaN, 500, models)).toBeUndefined();
      expect(calculateCost('model1', {}, 1000, Infinity, models)).toBeUndefined();
    });

    it('should return undefined if tokens are undefined', () => {
      expect(calculateCost('model1', {}, undefined, 500, models)).toBeUndefined();
      expect(calculateCost('model1', {}, 1000, undefined, models)).toBeUndefined();
    });
  });

  describe('isPromptfooSampleTarget', () => {
    it('should return true when url includes promptfoo.app', () => {
      const provider = createMockProvider({
        config: { url: 'https://api.promptfoo.app/v1' },
      });
      expect(isPromptfooSampleTarget(provider)).toBe(true);
    });

    it('should return true when url includes promptfoo.dev', () => {
      const provider = createMockProvider({
        config: { url: 'https://api.promptfoo.dev/v1' },
      });
      expect(isPromptfooSampleTarget(provider)).toBe(true);
    });

    it('should return false when url does not include promptfoo domains', () => {
      const provider = createMockProvider({
        config: { url: 'https://api.other-domain.com/v1' },
      });
      expect(isPromptfooSampleTarget(provider)).toBe(false);
    });

    it('should return false when provider.config is undefined', () => {
      const provider = createMockProvider();
      expect(isPromptfooSampleTarget(provider) ?? false).toBe(false);
    });

    it('should return false when provider.config.url is undefined', () => {
      const provider = createMockProvider({ config: {} });
      expect(isPromptfooSampleTarget(provider) ?? false).toBe(false);
    });
  });

  describe('OpenAI Tool Choice', () => {
    describe('isOpenAIToolChoice', () => {
      it('should return true for valid string values', () => {
        expect(isOpenAIToolChoice('auto')).toBe(true);
        expect(isOpenAIToolChoice('none')).toBe(true);
        expect(isOpenAIToolChoice('required')).toBe(true);
      });

      it('should return true for valid object form', () => {
        expect(isOpenAIToolChoice({ type: 'function', function: { name: 'my_tool' } })).toBe(true);
      });

      it('should return false for invalid inputs', () => {
        expect(isOpenAIToolChoice(null)).toBe(false);
        expect(isOpenAIToolChoice(undefined)).toBe(false);
        expect(isOpenAIToolChoice('invalid')).toBe(false);
        expect(isOpenAIToolChoice({})).toBe(false);
        expect(isOpenAIToolChoice({ type: 'function' })).toBe(false);
        expect(isOpenAIToolChoice({ type: 'function', function: {} })).toBe(false);
        // Old NormalizedToolChoice format should return false
        expect(isOpenAIToolChoice({ mode: 'auto' })).toBe(false);
        expect(isOpenAIToolChoice({ mode: 'tool', toolName: 'foo' })).toBe(false);
      });
    });

    describe('openaiToolChoiceToAnthropic', () => {
      it('should convert auto', () => {
        expect(openaiToolChoiceToAnthropic('auto')).toEqual({ type: 'auto' });
      });

      it('should convert none', () => {
        expect(openaiToolChoiceToAnthropic('none')).toEqual({ type: 'none' });
      });

      it('should convert required to any', () => {
        expect(openaiToolChoiceToAnthropic('required')).toEqual({ type: 'any' });
      });

      it('should convert specific tool', () => {
        expect(
          openaiToolChoiceToAnthropic({ type: 'function', function: { name: 'my_tool' } }),
        ).toEqual({
          type: 'tool',
          name: 'my_tool',
        });
      });
    });

    describe('openaiToolChoiceToBedrock', () => {
      it('should convert auto', () => {
        expect(openaiToolChoiceToBedrock('auto')).toEqual({ auto: {} });
      });

      it('should convert none to undefined (Bedrock has no none)', () => {
        expect(openaiToolChoiceToBedrock('none')).toBeUndefined();
      });

      it('should convert required to any', () => {
        expect(openaiToolChoiceToBedrock('required')).toEqual({ any: {} });
      });

      it('should convert specific tool', () => {
        expect(
          openaiToolChoiceToBedrock({ type: 'function', function: { name: 'my_tool' } }),
        ).toEqual({
          tool: { name: 'my_tool' },
        });
      });
    });

    describe('openaiToolChoiceToGoogle', () => {
      it('should convert auto', () => {
        expect(openaiToolChoiceToGoogle('auto')).toEqual({
          functionCallingConfig: { mode: 'AUTO' },
        });
      });

      it('should convert none', () => {
        expect(openaiToolChoiceToGoogle('none')).toEqual({
          functionCallingConfig: { mode: 'NONE' },
        });
      });

      it('should convert required to ANY', () => {
        expect(openaiToolChoiceToGoogle('required')).toEqual({
          functionCallingConfig: { mode: 'ANY' },
        });
      });

      it('should convert specific tool with allowedFunctionNames', () => {
        expect(
          openaiToolChoiceToGoogle({ type: 'function', function: { name: 'my_func' } }),
        ).toEqual({
          functionCallingConfig: { mode: 'ANY', allowedFunctionNames: ['my_func'] },
        });
      });
    });

    describe('transformToolChoice', () => {
      it('should pass through non-OpenAI values unchanged', () => {
        // Native Anthropic format
        expect(transformToolChoice({ type: 'auto' }, 'anthropic')).toEqual({ type: 'auto' });
        // Native Bedrock format
        expect(transformToolChoice({ auto: {} }, 'bedrock')).toEqual({ auto: {} });
      });

      it('should return OpenAI format as-is for openai target', () => {
        expect(transformToolChoice('auto', 'openai')).toBe('auto');
        expect(transformToolChoice('required', 'openai')).toBe('required');
        expect(
          transformToolChoice({ type: 'function', function: { name: 'foo' } }, 'openai'),
        ).toEqual({
          type: 'function',
          function: { name: 'foo' },
        });
      });

      it('should transform OpenAI format for Anthropic', () => {
        expect(transformToolChoice('auto', 'anthropic')).toEqual({ type: 'auto' });
        expect(transformToolChoice('none', 'anthropic')).toEqual({ type: 'none' });
        expect(transformToolChoice('required', 'anthropic')).toEqual({ type: 'any' });
      });

      it('should transform OpenAI format for Bedrock', () => {
        expect(transformToolChoice('auto', 'bedrock')).toEqual({ auto: {} });
        expect(transformToolChoice('required', 'bedrock')).toEqual({ any: {} });
      });

      it('should transform OpenAI format for Google', () => {
        expect(transformToolChoice('auto', 'google')).toEqual({
          functionCallingConfig: { mode: 'AUTO' },
        });
        expect(transformToolChoice('none', 'google')).toEqual({
          functionCallingConfig: { mode: 'NONE' },
        });
      });

      it('should pass through for unknown format', () => {
        expect(transformToolChoice('auto', 'unknown' as any)).toBe('auto');
      });
    });
  });

  describe('OpenAI Tool Format Transformation', () => {
    const sampleTools = [
      {
        type: 'function' as const,
        function: {
          name: 'get_weather',
          description: 'Get the current weather',
          parameters: {
            type: 'object',
            properties: {
              location: { type: 'string', description: 'City name' },
              unit: { type: 'string', enum: ['celsius', 'fahrenheit'] },
            },
            required: ['location'],
          },
        },
      },
      {
        type: 'function' as const,
        function: {
          name: 'search',
          description: 'Search the web',
          parameters: {
            type: 'object',
            properties: {
              query: { type: 'string' },
            },
          },
        },
      },
    ];

    describe('isOpenAIToolArray', () => {
      it('should return true for OpenAI format tools', () => {
        expect(isOpenAIToolArray(sampleTools)).toBe(true);
        expect(isOpenAIToolArray([{ type: 'function', function: { name: 'simple_tool' } }])).toBe(
          true,
        );
      });

      it('should return false for tools without type: function', () => {
        expect(isOpenAIToolArray([{ function: { name: 'test' } }])).toBe(false);
        expect(isOpenAIToolArray([{ type: 'other', function: { name: 'test' } }])).toBe(false);
      });

      it('should return false for tools without function.name', () => {
        expect(isOpenAIToolArray([{ type: 'function', function: {} }])).toBe(false);
        expect(isOpenAIToolArray([{ type: 'function' }])).toBe(false);
      });

      it('should return false for empty arrays', () => {
        expect(isOpenAIToolArray([])).toBe(false);
      });

      it('should return false for non-arrays', () => {
        expect(isOpenAIToolArray(null)).toBe(false);
        expect(isOpenAIToolArray(undefined)).toBe(false);
        expect(isOpenAIToolArray('string')).toBe(false);
        expect(isOpenAIToolArray({ type: 'function', function: { name: 'tool' } })).toBe(false);
      });

      it('should return false for other provider formats', () => {
        // Anthropic format
        const anthropicTools = [{ name: 'test', input_schema: { type: 'object' } }];
        expect(isOpenAIToolArray(anthropicTools)).toBe(false);

        // Bedrock format
        const bedrockTools = [{ toolSpec: { name: 'test' } }];
        expect(isOpenAIToolArray(bedrockTools)).toBe(false);

        // Google format
        const googleTools = [{ functionDeclarations: [{ name: 'test' }] }];
        expect(isOpenAIToolArray(googleTools)).toBe(false);
      });
    });

    describe('openaiToolsToAnthropic', () => {
      it('should convert to Anthropic format', () => {
        const result = openaiToolsToAnthropic(sampleTools);

        expect(result).toHaveLength(2);
        expect(result[0]).toEqual({
          name: 'get_weather',
          description: 'Get the current weather',
          input_schema: sampleTools[0].function.parameters,
        });
        expect(result[1]).toEqual({
          name: 'search',
          description: 'Search the web',
          input_schema: sampleTools[1].function.parameters,
        });
      });

      it('should provide default input_schema when no parameters', () => {
        const result = openaiToolsToAnthropic([
          { type: 'function', function: { name: 'minimal' } },
        ]);

        expect(result[0]).toEqual({
          name: 'minimal',
          input_schema: { type: 'object', properties: {} },
        });
      });

      it('should handle tools without description', () => {
        const result = openaiToolsToAnthropic([
          {
            type: 'function',
            function: { name: 'no_desc', parameters: { type: 'object' } },
          },
        ]);

        expect(result[0]).toEqual({
          name: 'no_desc',
          input_schema: { type: 'object' },
        });
      });
    });

    describe('openaiToolsToBedrock', () => {
      it('should convert to Bedrock Converse format', () => {
        const result = openaiToolsToBedrock(sampleTools);

        expect(result).toHaveLength(2);
        expect(result[0]).toEqual({
          toolSpec: {
            name: 'get_weather',
            description: 'Get the current weather',
            inputSchema: {
              json: sampleTools[0].function.parameters,
            },
          },
        });
      });

      it('should provide default inputSchema when no parameters', () => {
        const result = openaiToolsToBedrock([{ type: 'function', function: { name: 'minimal' } }]);

        expect(result[0]).toEqual({
          toolSpec: {
            name: 'minimal',
            inputSchema: {
              json: { type: 'object', properties: {} },
            },
          },
        });
      });
    });

    describe('openaiToolsToGoogle', () => {
      it('should convert to Google format with functionDeclarations array', () => {
        const result = openaiToolsToGoogle(sampleTools);

        expect(result).toHaveLength(1);
        expect(result[0].functionDeclarations).toHaveLength(2);
        expect(result[0].functionDeclarations[0].name).toBe('get_weather');
        expect(result[0].functionDeclarations[1].name).toBe('search');
      });

      it('should convert types to uppercase', () => {
        const result = openaiToolsToGoogle(sampleTools);

        const params = result[0].functionDeclarations[0].parameters as Record<string, unknown>;
        expect(params.type).toBe('OBJECT');
        const props = params.properties as Record<string, Record<string, unknown>>;
        expect(props.location.type).toBe('STRING');
      });

      it('should remove additionalProperties from schema', () => {
        const toolsWithAdditionalProps = [
          {
            type: 'function' as const,
            function: {
              name: 'test',
              parameters: {
                type: 'object',
                properties: { foo: { type: 'string' } },
                additionalProperties: false,
              },
            },
          },
        ];
        const result = openaiToolsToGoogle(toolsWithAdditionalProps);

        const params = result[0].functionDeclarations[0].parameters as Record<string, unknown>;
        expect(params).not.toHaveProperty('additionalProperties');
      });

      it('should handle tools without parameters', () => {
        const result = openaiToolsToGoogle([{ type: 'function', function: { name: 'minimal' } }]);

        expect(result[0].functionDeclarations[0]).toEqual({
          name: 'minimal',
        });
      });
    });

    describe('transformTools', () => {
      it('should pass through non-OpenAI format arrays', () => {
        const anthropicTools = [{ name: 'test', input_schema: { type: 'object' } }];
        expect(transformTools(anthropicTools, 'anthropic')).toEqual(anthropicTools);

        const bedrockTools = [{ toolSpec: { name: 'test' } }];
        expect(transformTools(bedrockTools, 'bedrock')).toEqual(bedrockTools);
      });

      it('should return OpenAI tools unchanged when format is openai', () => {
        const result = transformTools(sampleTools, 'openai');
        expect(result).toEqual(sampleTools);
      });

      it('should transform OpenAI tools for Anthropic', () => {
        const result = transformTools(sampleTools, 'anthropic') as any[];
        expect(result[0].name).toBe('get_weather');
        expect(result[0].input_schema).toBeDefined();
      });

      it('should transform OpenAI tools for Bedrock', () => {
        const result = transformTools(sampleTools, 'bedrock') as any[];
        expect(result[0].toolSpec.name).toBe('get_weather');
      });

      it('should transform OpenAI tools for Google', () => {
        const result = transformTools(sampleTools, 'google') as any[];
        expect(result[0].functionDeclarations[0].name).toBe('get_weather');
      });

      it('should pass through for unknown format', () => {
        expect(transformTools(sampleTools, 'unknown' as any)).toEqual(sampleTools);
      });
    });
  });
});