"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.getProviderTextOutput = getProviderTextOutput;
exports.generatePersonasPrompt = generatePersonasPrompt;
exports.testCasesPrompt = testCasesPrompt;
exports.synthesize = synthesize;
exports.synthesizeFromTestSuite = synthesizeFromTestSuite;
const decimal_js_1 = require("decimal.js");
const dedent_1 = __importDefault(require("dedent"));
const logger_1 = __importDefault(require("../logger"));
const providers_1 = require("../providers");
const defaults_1 = require("../providers/defaults");
const generation_1 = require("../util/generation");
const invariant_1 = __importDefault(require("../util/invariant"));
const json_1 = require("../util/json");
const templates_1 = require("../util/templates");
/** Normalize provider output or throw a user-actionable synthesis error. */
function getProviderTextOutput(resp, stepLabel) {
    if (resp.error) {
        throw new Error(`Dataset synthesis failed during ${stepLabel}: ${resp.error}`);
    }
    if (resp.output === undefined || resp.output === null) {
        throw new Error(`Dataset synthesis failed during ${stepLabel}: the model returned no output. Verify your synthesis provider (e.g. Azure OpenAI deployment name, API key, and quota).`);
    }
    return typeof resp.output === 'string' ? resp.output : JSON.stringify(resp.output);
}
function generatePersonasPrompt(prompts, numPersonas) {
    const promptsString = (0, dedent_1.default) `<Prompts>
    ${prompts.map((prompt) => `<Prompt>\n${prompt}\n</Prompt>`).join('\n')}
    </Prompts>`;
    return (0, dedent_1.default) `
    Consider the following prompt${prompts.length > 1 ? 's' : ''} for an LLM application:

    ${promptsString}

    List up to ${numPersonas} user personas that would send ${prompts.length > 1 ? 'these prompts' : 'this prompt'}. Your response should be JSON of the form {personas: string[]}`;
}
function testCasesPrompt(prompts, persona, tests, numTestCasesPerPersona, variables, instructions) {
    const promptsString = (0, dedent_1.default) `
    <Prompts>
    ${prompts
        .map((prompt) => (0, dedent_1.default) `
      <Prompt>
      ${prompt}
      </Prompt>`)
        .join('\n')}
    </Prompts>`;
    const existingTests = (0, dedent_1.default) `
    Here are some existing tests:
    ${(0, generation_1.sampleArray)(tests, 100)
        .map((test) => {
        if (!test.vars) {
            return;
        }
        return (0, dedent_1.default) `
          <Test>
          ${JSON.stringify(test.vars, null, 2)}
          </Test>`;
    })
        .filter(Boolean)
        .sort()
        .join('\n')}
  `;
    return (0, dedent_1.default) `
    Consider ${prompts.length > 1 ? 'these prompts' : 'this prompt'}:
  ${promptsString}

  This is your persona:
  <Persona>
  ${persona}
  </Persona>

  ${existingTests}

  Fully embody this persona and determine a value for each variable, such that the prompt would be sent by this persona.

  You are a tester, so try to think of ${numTestCasesPerPersona} sets of values that would be interesting or unusual to test.${instructions ? ` ${instructions}` : ''}

  IMPORTANT:
  - Return values for ONLY these variables: ${variables.join(', ') || '(none)'}.
  - Do NOT add any additional keys (even if the prompt template contains other {{variables}}).

  Your response should contain a JSON map of variable names to values, of the form {vars: {${Array.from(variables)
        .map((varName) => `${varName}: string`)
        .join(', ')}}[]}`;
}
async function synthesize({ tokenUsages, prompts, instructions, tests, numPersonas, numTestCasesPerPersona, provider, apiCall = false, workspaceId = '', }) {
    if (prompts.length < 1) {
        throw new Error('Dataset synthesis requires at least one prompt.');
    }
    if (!tokenUsages) {
        tokenUsages = {
            cost: new decimal_js_1.Decimal(0),
            tokenUsage: {
                prompt: new decimal_js_1.Decimal(0),
                completion: new decimal_js_1.Decimal(0),
                total: new decimal_js_1.Decimal(0),
                cached: new decimal_js_1.Decimal(0),
            },
        };
    }
    numPersonas = numPersonas || 5;
    numTestCasesPerPersona = numTestCasesPerPersona || 3;
    let progressBar;
    if (logger_1.default.level !== 'debug') {
        const cliProgress = await import('cli-progress');
        progressBar = new cliProgress.SingleBar({ gracefulExit: true }, cliProgress.Presets.shades_classic);
        const totalProgressSteps = 1 + numPersonas * numTestCasesPerPersona;
        progressBar.start(totalProgressSteps, 0);
    }
    logger_1.default.debug(`Starting dataset synthesis. We'll begin by generating up to ${numPersonas} personas. Each persona will be used to generate ${numTestCasesPerPersona} test cases.`);
    logger_1.default.debug(`Generating user personas from ${prompts.length} prompt${prompts.length > 1 ? 's' : ''}...`);
    let providerModel;
    if (typeof provider === 'undefined') {
        providerModel = (await (0, defaults_1.getDefaultProviders)(undefined, workspaceId)).synthesizeProvider;
    }
    else {
        providerModel = await (0, providers_1.loadApiProvider)(provider);
    }
    const personasPrompt = generatePersonasPrompt(prompts, numPersonas);
    logger_1.default.debug(`Generated personas prompt:\n${personasPrompt}`);
    const resp = await providerModel.callApi(personasPrompt);
    tokenUsages.cost = tokenUsages.cost.plus(resp.cost ?? 0);
    tokenUsages.tokenUsage.prompt = tokenUsages.tokenUsage.prompt.plus(resp.tokenUsage?.prompt ?? 0);
    tokenUsages.tokenUsage.completion = tokenUsages.tokenUsage.completion.plus(resp.tokenUsage?.completion ?? 0);
    tokenUsages.tokenUsage.total = tokenUsages.tokenUsage.total.plus(resp.tokenUsage?.total ?? 0);
    tokenUsages.tokenUsage.cached = tokenUsages.tokenUsage.cached.plus(resp?.tokenUsage?.cached ?? 0);
    const output = getProviderTextOutput(resp, 'persona generation');
    logger_1.default.debug(`Received personas response:\n${output}`);
    const respObjects = (0, json_1.extractJsonObjects)(output);
    (0, invariant_1.default)(respObjects.length >= 1, `Expected at least one JSON object in the response for personas, got ${respObjects.length}`);
    const personas = respObjects[0].personas;
    logger_1.default.debug(`Generated ${personas.length} persona${personas.length === 1 ? '' : 's'}:\n${personas.map((p) => `  - ${p}`).join('\n')}`);
    if (progressBar) {
        progressBar.increment();
    }
    // Determine which variables to generate.
    // If tests already specify vars, ONLY generate those (avoid pulling extra template vars).
    const testVars = Array.isArray(tests)
        ? Array.from(new Set(tests
            .flatMap((t) => (t?.vars && typeof t.vars === 'object' ? Object.keys(t.vars) : []))
            .filter(Boolean)))
        : [];
    // For api calls we *always* constrain to test vars. For non-api calls, also prefer test vars
    // when present (so prompts with extra `{{ }}` don't expand the schema unexpectedly).
    let variables = testVars.length > 0 ? testVars : (0, templates_1.extractVariablesFromTemplates)(prompts);
    if (apiCall) {
        variables = testVars;
    }
    logger_1.default.debug(`Extracted ${variables.length} variable${variables.length === 1 ? '' : 's'} from prompt${prompts.length === 1 ? '' : 's'}:\n${variables
        .map((v) => `  - ${v}`)
        .join('\n')}`);
    const batchSize = 20;
    const totalTestCases = numPersonas * numTestCasesPerPersona;
    const generateTestCasesForPersona = async (currentTestCases, doubleBatchsize) => {
        const remainingCount = totalTestCases - currentTestCases.length;
        let currentBatchSize = Math.min(remainingCount, batchSize);
        // If we got the same number of unique items last time, double the batch size
        if (doubleBatchsize) {
            currentBatchSize = Math.min(remainingCount, currentBatchSize) * 2;
            logger_1.default.debug(`Doubled batch size to ${currentBatchSize} due to duplicates`);
        }
        // Calculate which persona to use based on total test cases generated so far
        const totalTestCasesGenerated = currentTestCases.length;
        const personaIndex = Math.floor(totalTestCasesGenerated / (numTestCasesPerPersona || 3));
        const persona = personas[personaIndex % personas.length];
        logger_1.default.debug(`Generating ${currentBatchSize} test cases for persona ${personaIndex + 1} of ${personas.length}...`);
        const personaPrompt = testCasesPrompt(prompts, persona, tests, currentBatchSize, variables, instructions);
        logger_1.default.debug(`Generated persona prompt:\n${personaPrompt}`);
        const personaResponse = await providerModel.callApi(personaPrompt);
        const personaOutput = getProviderTextOutput(personaResponse, `test case generation for persona "${persona}"`);
        logger_1.default.debug(`Received persona response:\n${personaOutput}`);
        tokenUsages.cost = tokenUsages.cost.plus(personaResponse.cost ?? 0);
        tokenUsages.tokenUsage.prompt = tokenUsages.tokenUsage.prompt.plus(personaResponse.tokenUsage?.prompt ?? 0);
        tokenUsages.tokenUsage.completion = tokenUsages.tokenUsage.completion.plus(personaResponse.tokenUsage?.completion ?? 0);
        tokenUsages.tokenUsage.total = tokenUsages.tokenUsage.total.plus(personaResponse.tokenUsage?.total ?? 0);
        tokenUsages.tokenUsage.cached = tokenUsages.tokenUsage.cached.plus(personaResponse?.tokenUsage?.cached ?? 0);
        const personaResponseObjects = (0, json_1.extractJsonObjects)(personaOutput);
        (0, invariant_1.default)(personaResponseObjects.length >= 1, `Expected at least one JSON object in the response for persona ${persona}, got ${personaResponseObjects.length}`);
        const parsed = personaResponseObjects[0];
        logger_1.default.debug(`Received ${parsed.vars?.length} test cases`);
        if (progressBar) {
            progressBar.increment(parsed.vars?.length);
        }
        return parsed.vars || [];
    };
    let testCaseVars = await (0, generation_1.retryWithDeduplication)(generateTestCasesForPersona, totalTestCases);
    logger_1.default.debug(`Generated ${testCaseVars.length} test cases`);
    if (testCaseVars.length > totalTestCases) {
        logger_1.default.debug(`Generated ${testCaseVars.length} test cases, but only ${totalTestCases} were requested. Sampling down to ${totalTestCases}...`);
        testCaseVars = (0, generation_1.sampleArray)(testCaseVars, totalTestCases);
    }
    if (progressBar) {
        progressBar.stop();
    }
    if (apiCall) {
        return {
            personas,
            testCaseVars,
        };
    }
    else {
        return testCaseVars;
    }
}
async function synthesizeFromTestSuite(testSuite, options, tokenUsages = {
    cost: new decimal_js_1.Decimal(0),
    tokenUsage: {
        prompt: new decimal_js_1.Decimal(0),
        completion: new decimal_js_1.Decimal(0),
        total: new decimal_js_1.Decimal(0),
        cached: new decimal_js_1.Decimal(0),
    },
}) {
    return synthesize({
        ...options,
        prompts: testSuite.prompts.map((prompt) => prompt.raw),
        tests: testSuite.tests || [],
        tokenUsages,
    });
}
//# sourceMappingURL=synthesis.js.map