promptfoo--promptfoo/test/integrations/huggingfaceDatasets.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { fetchWithCache } from '../../src/cache';
import { getEnvString } from '../../src/envars';
import {
  fetchHuggingFaceDataset,
  parseDatasetPath,
} from '../../src/integrations/huggingfaceDatasets';

vi.mock('../../src/cache', () => ({
  fetchWithCache: vi.fn(),
}));

vi.mock('../../src/util/fetch/index.ts', () => ({
  fetchWithProxy: vi.fn(),
}));

vi.mock('../../src/envars', () => ({
  getEnvString: vi.fn().mockReturnValue(''),
  isCI: vi.fn().mockReturnValue(false),
}));

describe('huggingfaceDatasets', () => {
  beforeEach(() => {
    vi.mocked(getEnvString).mockReturnValue('');
  });

  afterEach(() => {
    vi.resetAllMocks();
  });

  describe('parseDatasetPath', () => {
    it('should parse path with default parameters', () => {
      const result = parseDatasetPath('huggingface://datasets/owner/repo');
      expect(result).toEqual({
        owner: 'owner',
        repo: 'repo',
        queryParams: expect.any(URLSearchParams),
      });
      expect(result.queryParams.get('split')).toBe('test');
      expect(result.queryParams.get('config')).toBe('default');
    });

    it('should parse path with custom query parameters', () => {
      const result = parseDatasetPath(
        'huggingface://datasets/owner/repo?split=train&config=custom&limit=10',
      );
      expect(result).toEqual({
        owner: 'owner',
        repo: 'repo',
        queryParams: expect.any(URLSearchParams),
      });
      expect(result.queryParams.get('split')).toBe('train');
      expect(result.queryParams.get('config')).toBe('custom');
      expect(result.queryParams.get('limit')).toBe('10');
    });

    it('should override default parameters with user parameters', () => {
      const result = parseDatasetPath('huggingface://datasets/owner/repo?split=validation');
      expect(result.queryParams.get('split')).toBe('validation');
      expect(result.queryParams.get('config')).toBe('default');
    });
  });

  it('should fetch and parse dataset with default parameters', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 2,
        features: [
          { name: 'act', type: { dtype: 'string', _type: 'Value' } },
          { name: 'prompt', type: { dtype: 'string', _type: 'Value' } },
        ],
        rows: [
          { row: { act: 'Linux Terminal', prompt: 'List all files' } },
          { row: { act: 'Math Tutor', prompt: 'Solve 2+2' } },
        ],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
      expect.objectContaining({
        headers: {},
      }),
    );

    expect(tests).toHaveLength(2);
    expect(tests[0].vars).toEqual({
      act: 'Linux Terminal',
      prompt: 'List all files',
    });
    expect(tests[1].vars).toEqual({
      act: 'Math Tutor',
      prompt: 'Solve 2+2',
    });

    // Check that disableVarExpansion is set for all test cases
    tests.forEach((test) => {
      expect(test.options).toEqual({
        disableVarExpansion: true,
      });
    });
  });

  it('should include auth token when HF_TOKEN is set', async () => {
    vi.mocked(getEnvString).mockImplementation((key) => {
      if (key === 'HF_TOKEN') {
        return 'test-token';
      }
      return '';
    });

    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 1,
        features: [
          { name: 'question', type: { dtype: 'string', _type: 'Value' } },
          { name: 'answer', type: { dtype: 'string', _type: 'Value' } },
        ],
        rows: [{ row: { question: 'What is 2+2?', answer: '4' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
      expect.objectContaining({
        headers: {
          Authorization: 'Bearer test-token',
        },
      }),
    );
  });

  it('should fall back to HF_API_TOKEN when HF_TOKEN is empty', async () => {
    vi.mocked(getEnvString).mockImplementation((key) => {
      if (key === 'HF_TOKEN') {
        return '';
      }
      if (key === 'HF_API_TOKEN') {
        return 'api-token';
      }
      return '';
    });

    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 1,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: [{ row: { text: 'test' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 1);

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      expect.any(String),
      expect.objectContaining({
        headers: {
          Authorization: 'Bearer api-token',
        },
      }),
    );
  });

  it('should fall back to HUGGING_FACE_HUB_TOKEN when other tokens are empty', async () => {
    vi.mocked(getEnvString).mockImplementation((key) => {
      if (key === 'HF_TOKEN') {
        return '';
      }
      if (key === 'HF_API_TOKEN') {
        return '';
      }
      if (key === 'HUGGING_FACE_HUB_TOKEN') {
        return 'hub-token';
      }
      return '';
    });

    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 1,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: [{ row: { text: 'test' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 1);

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      expect.any(String),
      expect.objectContaining({
        headers: {
          Authorization: 'Bearer hub-token',
        },
      }),
    );
  });

  it('should handle custom query parameters', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 1,
        features: [
          { name: 'question', type: { dtype: 'string', _type: 'Value' } },
          { name: 'answer', type: { dtype: 'string', _type: 'Value' } },
        ],
        rows: [{ row: { question: 'What is 2+2?', answer: '4' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?split=train&config=custom');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=train&config=custom&offset=0&length=100',
      expect.objectContaining({
        headers: {},
      }),
    );
  });

  it('should handle pagination', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 3,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: [{ row: { text: 'First' } }, { row: { text: 'Second' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 3,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: [{ row: { text: 'Third' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(2);
    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
      1,
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
      expect.objectContaining({
        headers: {},
      }),
    );
    // Note: Second call might have different length due to concurrent fetching and remaining calculation
    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
      2,
      expect.stringContaining('dataset=test%2Fdataset&split=test&config=default&offset=2'),
      expect.objectContaining({
        headers: {},
      }),
    );

    expect(tests).toHaveLength(3);
    expect(tests.map((t) => t.vars?.text)).toEqual(['First', 'Second', 'Third']);

    // Check that disableVarExpansion is set for all test cases
    tests.forEach((test) => {
      expect(test.options).toEqual({
        disableVarExpansion: true,
      });
    });
  });

  it('should reject an empty non-final page instead of retrying the same offset', async () => {
    vi.mocked(fetchWithCache)
      .mockResolvedValueOnce({
        data: {
          num_rows_total: 2,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: [],
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      } as any)
      .mockRejectedValueOnce(new Error('Paginator retried the empty page'));

    await expect(
      fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 200),
    ).rejects.toThrow(
      '[HF Dataset] Received an empty page at offset 0 before reaching 2 total rows',
    );

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(1);
  });

  it('should handle API errors by throwing', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: null,
      cached: false,
      status: 404,
      statusText: 'Not Found',
    } as any);

    await expect(
      fetchHuggingFaceDataset('huggingface://datasets/nonexistent/dataset'),
    ).rejects.toThrow('[HF Dataset] Failed to fetch dataset: Not Found');
  });

  it('should short-circuit and return [] when limit is 0', async () => {
    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 0);
    expect(vi.mocked(fetchWithCache)).not.toHaveBeenCalled();
    expect(tests).toEqual([]);
  });

  it('should respect user-specified limit parameter (single request optimization)', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 5,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: [{ row: { text: 'First' } }, { row: { text: 'Second' } }],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?limit=2');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=2&offset=0&length=2',
      expect.objectContaining({
        headers: {},
      }),
    );

    expect(tests).toHaveLength(2);
    expect(tests.map((t) => t.vars?.text)).toEqual(['First', 'Second']);

    // Check that disableVarExpansion is set for all test cases
    tests.forEach((test) => {
      expect(test.options).toEqual({
        disableVarExpansion: true,
      });
    });
  });

  it('should handle limit larger than page size', async () => {
    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 150,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: Array(100)
          .fill(null)
          .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    vi.mocked(fetchWithCache).mockResolvedValueOnce({
      data: {
        num_rows_total: 150,
        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
        rows: Array(20)
          .fill(null)
          .map((_, i) => ({ row: { text: `Item ${i + 101}` } })),
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    } as any);

    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?limit=120');

    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(2);
    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
      1,
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=120&offset=0&length=100',
      expect.objectContaining({
        headers: {},
      }),
    );
    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
      2,
      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=120&offset=100&length=20',
      expect.objectContaining({
        headers: {},
      }),
    );

    expect(tests).toHaveLength(120);
    expect(tests[119].vars?.text).toBe('Item 120');

    // Check that disableVarExpansion is set for all test cases
    tests.forEach((test) => {
      expect(test.options).toEqual({
        disableVarExpansion: true,
      });
    });
  });

  describe('performance optimizations', () => {
    it('should use single request optimization for small limits', async () => {
      vi.mocked(fetchWithCache).mockResolvedValueOnce({
        data: {
          num_rows_total: 1000,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: Array(50)
            .fill(null)
            .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
        },
        cached: true,
        status: 200,
        statusText: 'OK',
      } as any);

      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 50);

      // Should only make one request for limits <= 100
      expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(1);
      expect(tests).toHaveLength(50);
    });

    it('should throw error on page fetch failure', async () => {
      // First page succeeds
      vi.mocked(fetchWithCache).mockResolvedValueOnce({
        data: {
          num_rows_total: 300,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: Array(100)
            .fill(null)
            .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      } as any);

      // Second page fails
      vi.mocked(fetchWithCache).mockResolvedValueOnce({
        data: null,
        cached: false,
        status: 500,
        statusText: 'Internal Server Error',
      } as any);

      // Should throw error instead of returning partial results
      await expect(
        fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 200),
      ).rejects.toThrow('[HF Dataset] Failed to fetch dataset: Internal Server Error');
    });

    it('should include rows from concurrent fetches without duplicates', async () => {
      const totalRows = 400;
      const rowPrefix = 'x'.repeat(300);

      vi.mocked(fetchWithCache).mockImplementation(async (url) => {
        const searchParams = new URL(String(url)).searchParams;
        const offset = Number.parseInt(searchParams.get('offset') ?? '0', 10);
        const length = Number.parseInt(searchParams.get('length') ?? '100', 10);

        return {
          data: {
            num_rows_total: totalRows,
            features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
            rows: Array.from({ length }, (_, i) => ({
              row: { text: `${rowPrefix}${offset + i + 1}` },
            })),
          },
          cached: false,
          status: 200,
          statusText: 'OK',
        } as any;
      });

      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');
      const texts = tests.map((test) => test.vars?.text);
      const calledOffsets = vi
        .mocked(fetchWithCache)
        .mock.calls.map(([url]) => new URL(String(url)).searchParams.get('offset'));

      expect(calledOffsets).toContain('100');
      expect(tests).toHaveLength(totalRows);
      expect(texts).toContain(`${rowPrefix}150`);
      expect(new Set(texts).size).toBe(totalRows);
    });

    it('should adapt page size based on row size', async () => {
      // Mock a dataset with large rows (>2KB each)
      const largeRow = { text: 'x'.repeat(3000) }; // ~3KB row

      vi.mocked(fetchWithCache).mockResolvedValueOnce({
        data: {
          num_rows_total: 200,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: Array(100)
            .fill(null)
            .map(() => ({ row: largeRow })),
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      } as any);

      // Mock all subsequent potential concurrent requests to avoid undefined errors
      vi.mocked(fetchWithCache).mockResolvedValue({
        data: {
          num_rows_total: 200,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: Array(25)
            .fill(null)
            .map((_, i) => ({ row: { text: `Item ${i + 101}` } })),
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      } as any);

      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 125);

      // Should have made at least one request
      expect(vi.mocked(fetchWithCache)).toHaveBeenCalled();

      // First call should be normal page size
      expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
        1,
        expect.stringContaining('length=100'),
        expect.anything(),
      );

      expect(tests.length).toBeGreaterThan(0);
      expect(tests[0].vars?.text).toBe(largeRow.text);
    });

    it('should handle authentication tokens correctly', async () => {
      vi.mocked(getEnvString).mockImplementation((key) => {
        if (key === 'HF_TOKEN') {
          return 'test-token-123';
        }
        return '';
      });

      vi.mocked(fetchWithCache).mockResolvedValueOnce({
        data: {
          num_rows_total: 10,
          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
          rows: [{ row: { text: 'Test' } }],
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      } as any);

      await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 5);

      expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
        expect.any(String),
        expect.objectContaining({
          headers: {
            Authorization: 'Bearer test-token-123',
          },
        }),
      );
    });
  });
});