"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.synthesize = exports.synthesizeFromTestSuite = exports.DEFAULT_PLUGINS = exports.ADDITIONAL_PLUGINS = void 0;
const chalk_1 = __importDefault(require("chalk"));
const cli_progress_1 = __importDefault(require("cli-progress"));
const dedent_1 = __importDefault(require("dedent"));
const tiny_invariant_1 = __importDefault(require("tiny-invariant"));
const logger_1 = __importDefault(require("../logger"));
const providers_1 = require("../providers");
const constants_1 = require("./constants");
const getCompetitorTests_1 = require("./getCompetitorTests");
const getHallucinationTests_1 = require("./getHallucinationTests");
const getHarmfulTests_1 = require("./getHarmfulTests");
const getHijackingTests_1 = require("./getHijackingTests");
const getOverconfidenceTests_1 = require("./getOverconfidenceTests");
const getPiiTests_1 = require("./getPiiTests");
const getPoliticalStatementsTests_1 = require("./getPoliticalStatementsTests");
const getUnderconfidenceTests_1 = require("./getUnderconfidenceTests");
const getUnintendedContractTests_1 = require("./getUnintendedContractTests");
const BASE_PLUGINS = [
    'contracts',
    'excessive-agency',
    'hallucination',
    'harmful',
    'hijacking',
    'jailbreak',
    'overreliance',
    'pii',
    'politics',
    'prompt-injection',
];
exports.ADDITIONAL_PLUGINS = ['competitors'];
exports.DEFAULT_PLUGINS = new Set([...BASE_PLUGINS, ...Object.keys(getHarmfulTests_1.HARM_CATEGORIES)]);
const ALL_PLUGINS = new Set([...exports.DEFAULT_PLUGINS, ...exports.ADDITIONAL_PLUGINS]);
function validatePlugins(plugins) {
    for (const plugin of plugins) {
        if (!ALL_PLUGINS.has(plugin)) {
            throw new Error(`Invalid plugin: ${plugin}. Must be one of: ${Array.from(ALL_PLUGINS).join(', ')}`);
        }
    }
}
async function synthesizeFromTestSuite(testSuite, options) {
    return synthesize({
        ...options,
        plugins: options.plugins && options.plugins.length > 0 ? options.plugins : Array.from(exports.DEFAULT_PLUGINS),
        prompts: testSuite.prompts.map((prompt) => prompt.raw),
    });
}
exports.synthesizeFromTestSuite = synthesizeFromTestSuite;
async function synthesize({ prompts, provider, injectVar, purpose: purposeOverride, plugins, }) {
    validatePlugins(plugins);
    const reasoningProvider = await (0, providers_1.loadApiProvider)(provider || constants_1.REDTEAM_MODEL);
    logger_1.default.info(`Synthesizing test cases for ${prompts.length} ${prompts.length === 1 ? 'prompt' : 'prompts'}...\nPlugins: ${chalk_1.default.yellow(plugins.join(', '))}`);
    // Initialize progress bar
    const progressBar = new cli_progress_1.default.SingleBar({}, cli_progress_1.default.Presets.shades_classic);
    const totalSteps = plugins.length + 2; // +2 for initial setup steps
    let currentStep = 0;
    if (process.env.LOG_LEVEL !== 'debug') {
        progressBar.start(100, 0);
    }
    const updateProgress = () => {
        currentStep += 1;
        const progress = Math.floor((currentStep / totalSteps) * 100);
        progressBar.update(progress);
    };
    // Get vars
    injectVar = injectVar || 'query';
    // Get purpose
    updateProgress();
    const purpose = purposeOverride || (await getPurpose(prompts, reasoningProvider));
    updateProgress();
    logger_1.default.debug(`System purpose: ${purpose}`);
    // Get adversarial test cases
    const testCases = [];
    const adversarialPrompts = [];
    const redteamProvider = await (0, providers_1.loadApiProvider)(provider || constants_1.REDTEAM_MODEL, {
        options: {
            config: { temperature: 0.5 },
        },
    });
    const addHarmfulCases = plugins.some((p) => p.startsWith('harmful'));
    if (plugins.includes('prompt-injection') || plugins.includes('jailbreak') || addHarmfulCases) {
        logger_1.default.debug('Generating harmful test cases');
        const harmfulPrompts = await (0, getHarmfulTests_1.getHarmfulTests)(redteamProvider, purpose, injectVar, plugins.filter((p) => p.startsWith('harmful:')));
        adversarialPrompts.push(...harmfulPrompts);
        if (addHarmfulCases) {
            testCases.push(...harmfulPrompts);
            logger_1.default.debug(`Added ${harmfulPrompts.length} harmful test cases`);
        }
        else {
            logger_1.default.debug(`Generated ${harmfulPrompts.length} harmful test cases`);
        }
    }
    const pluginActions = {
        competitors: getCompetitorTests_1.getCompetitorTests,
        contracts: getUnintendedContractTests_1.getContractTests,
        'excessive-agency': getOverconfidenceTests_1.getOverconfidenceTests,
        hallucination: getHallucinationTests_1.getHallucinationTests,
        hijacking: getHijackingTests_1.getHijackingTests,
        jailbreak: (provider, purpose, injectVar) => (0, getHarmfulTests_1.addIterativeJailbreaks)(provider, adversarialPrompts, purpose, injectVar),
        overreliance: getUnderconfidenceTests_1.getUnderconfidenceTests,
        pii: getPiiTests_1.getPiiTests,
        politics: getPoliticalStatementsTests_1.getPoliticalStatementsTests,
        'prompt-injection': (provider, purpose, injectVar) => (0, getHarmfulTests_1.addInjections)(provider, adversarialPrompts, purpose, injectVar),
    };
    for (const plugin of plugins) {
        if (pluginActions[plugin]) {
            updateProgress();
            logger_1.default.debug(`Generating ${plugin} tests`);
            const pluginTests = await pluginActions[plugin](redteamProvider, purpose, injectVar);
            testCases.push(...pluginTests);
            logger_1.default.debug(`Added ${pluginTests.length} ${plugin} test cases`);
        }
    }
    // Finish progress bar
    if (process.env.LOG_LEVEL !== 'debug') {
        progressBar.update(100);
        progressBar.stop();
    }
    return testCases;
}
exports.synthesize = synthesize;
async function getPurpose(prompts, provider) {
    const { output: purpose } = await provider.callApi((0, dedent_1.default) `
    The following are prompts that are being used to test an LLM application:
    
    ${prompts
        .map((prompt) => (0, dedent_1.default) `
      <prompt>
      ${prompt}
      </prompt>`)
        .join('\n')}
    
    Given the above prompts, output the "system purpose" of the application in a single sentence.
    
    Example outputs:
    - Provide users a way to manage finances
    - Executive assistant that helps with scheduling and reminders
    - Ecommerce chatbot that sells shoes
  `);
    (0, tiny_invariant_1.default)(typeof purpose === 'string', `Expected purpose to be a string, got: ${purpose}`);
    return purpose;
}
//# sourceMappingURL=index.js.map