"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
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.readTests = exports.readTest = exports.readStandaloneTestsFile = exports.readVarsFiles = void 0;
const sync_1 = require("csv-parse/sync");
const fs = __importStar(require("fs"));
const glob_1 = require("glob");
const js_yaml_1 = __importDefault(require("js-yaml"));
const path = __importStar(require("path"));
const path_1 = require("path");
const csv_1 = require("./csv");
const googleSheets_1 = require("./googleSheets");
const logger_1 = __importDefault(require("./logger"));
const providers_1 = require("./providers");
const openai_1 = require("./providers/openai");
const SYNTHESIZE_DEFAULT_PROVIDER = 'gpt-4-0125-preview';
function parseJson(json) {
    try {
        return JSON.parse(json);
    }
    catch (err) {
        return undefined;
    }
}
async function readVarsFiles(pathOrGlobs, basePath = '') {
    if (typeof pathOrGlobs === 'string') {
        pathOrGlobs = [pathOrGlobs];
    }
    const ret = {};
    for (const pathOrGlob of pathOrGlobs) {
        const resolvedPath = path.resolve(basePath, pathOrGlob);
        const paths = (0, glob_1.globSync)(resolvedPath, {
            windowsPathsNoEscape: true,
        });
        for (const p of paths) {
            const yamlData = js_yaml_1.default.load(fs.readFileSync(p, 'utf-8'));
            Object.assign(ret, yamlData);
        }
    }
    return ret;
}
exports.readVarsFiles = readVarsFiles;
async function readStandaloneTestsFile(varsPath, basePath = '') {
    // This function is confusingly named - it reads a CSV, JSON, or YAML file of
    // TESTS or test equivalents.
    const resolvedVarsPath = path.resolve(basePath, varsPath);
    const fileExtension = (0, path_1.parse)(varsPath).ext.slice(1);
    let rows = [];
    if (varsPath.startsWith('https://docs.google.com/spreadsheets/')) {
        rows = await (0, googleSheets_1.fetchCsvFromGoogleSheet)(varsPath);
    }
    else if (fileExtension === 'csv') {
        rows = (0, sync_1.parse)(fs.readFileSync(resolvedVarsPath, 'utf-8'), { columns: true });
    }
    else if (fileExtension === 'json') {
        rows = parseJson(fs.readFileSync(resolvedVarsPath, 'utf-8'));
    }
    else if (fileExtension === 'yaml' || fileExtension === 'yml') {
        rows = js_yaml_1.default.load(fs.readFileSync(resolvedVarsPath, 'utf-8'));
    }
    return rows.map((row, idx) => {
        const test = (0, csv_1.testCaseFromCsvRow)(row);
        test.description = `Row #${idx + 1}`;
        return test;
    });
}
exports.readStandaloneTestsFile = readStandaloneTestsFile;
async function readTest(test, basePath = '') {
    const loadTestWithVars = async (testCase, testBasePath) => {
        const ret = { ...testCase, vars: undefined };
        if (typeof testCase.vars === 'string' || Array.isArray(testCase.vars)) {
            ret.vars = await readVarsFiles(testCase.vars, testBasePath);
        }
        else {
            ret.vars = testCase.vars;
        } /*else if (typeof testCase.vars === 'object') {
          const vars: Record<string, string | string[] | object> = {};
          for (const [key, value] of Object.entries(testCase.vars)) {
            if (typeof value === 'string' && value.startsWith('file://')) {
              // Load file from disk.
              const filePath = path.resolve(testBasePath, value.slice('file://'.length));
              if (filePath.endsWith('.yaml') || filePath.endsWith('.yml')) {
                vars[key] = (yaml.load(fs.readFileSync(filePath, 'utf-8')) as string).trim();
              } else {
                vars[key] = fs.readFileSync(filePath, 'utf-8').trim();
              }
            } else {
              // This is a normal key:value.
              vars[key] = value;
            }
          }
          ret.vars = vars;
        }*/
        return ret;
    };
    let testCase;
    if (typeof test === 'string') {
        const testFilePath = path.resolve(basePath, test);
        const testBasePath = path.dirname(testFilePath);
        const rawTestCase = js_yaml_1.default.load(fs.readFileSync(testFilePath, 'utf-8'));
        testCase = await loadTestWithVars(rawTestCase, testBasePath);
    }
    else {
        testCase = await loadTestWithVars(test, basePath);
    }
    if (testCase.provider && typeof testCase.provider !== 'function') {
        // Load provider
        if (typeof testCase.provider === 'string') {
            testCase.provider = await (0, providers_1.loadApiProvider)(testCase.provider);
        }
        else if (typeof testCase.provider.id === 'string') {
            testCase.provider = await (0, providers_1.loadApiProvider)(testCase.provider.id, {
                options: testCase.provider,
                basePath,
            });
        }
    }
    // Validation of the shape of test
    if (!testCase.assert && !testCase.vars && !testCase.options) {
        throw new Error(`Test case must have either assert, vars, or options property. Instead got ${JSON.stringify(testCase, null, 2)}`);
    }
    return testCase;
}
exports.readTest = readTest;
async function readTests(tests, basePath = '') {
    const ret = [];
    const loadTestsFromGlob = async (loadTestsGlob) => {
        const resolvedPath = path.resolve(basePath, loadTestsGlob);
        const testFiles = (0, glob_1.globSync)(resolvedPath, {
            windowsPathsNoEscape: true,
        });
        const ret = [];
        if (testFiles.length < 1) {
            logger_1.default.error(`No test files found for path: ${loadTestsGlob}`);
            return ret;
        }
        for (const testFile of testFiles) {
            let testCases;
            if (testFile.endsWith('.csv')) {
                testCases = await readStandaloneTestsFile(testFile, basePath);
            }
            else if (testFile.endsWith('.yaml') || testFile.endsWith('.yml')) {
                testCases = js_yaml_1.default.load(fs.readFileSync(testFile, 'utf-8'));
            }
            else if (testFile.endsWith('.json')) {
                testCases = require(testFile);
            }
            else {
                throw new Error(`Unsupported file type for test file: ${testFile}`);
            }
            if (testCases) {
                for (const testCase of testCases) {
                    ret.push(await readTest(testCase, path.dirname(testFile)));
                }
            }
        }
        return ret;
    };
    if (typeof tests === 'string') {
        if (tests.endsWith('yaml') || tests.endsWith('yml')) {
            // Points to a tests file with multiple test cases
            return loadTestsFromGlob(tests);
        }
        else {
            // Points to a tests.csv or Google Sheet
            return readStandaloneTestsFile(tests, basePath);
        }
    }
    else if (Array.isArray(tests)) {
        for (const globOrTest of tests) {
            if (typeof globOrTest === 'string') {
                // Resolve globs
                ret.push(...(await loadTestsFromGlob(globOrTest)));
            }
            else {
                // It's just a TestCase
                ret.push(await readTest(globOrTest, basePath));
            }
        }
    }
    return ret;
}
exports.readTests = readTests;
async function synthesizeFromTestSuite(testSuite, options) {
    return synthesize({
        ...options,
        prompts: testSuite.prompts.map((prompt) => prompt.raw),
        tests: testSuite.tests || [],
    });
}
exports.synthesizeFromTestSuite = synthesizeFromTestSuite;
async function synthesize({ prompts, instructions, tests, numPersonas, numTestCasesPerPersona, }) {
    if (prompts.length < 1) {
        throw new Error('Dataset synthesis requires at least one prompt.');
    }
    numPersonas = numPersonas || 5;
    numTestCasesPerPersona = numTestCasesPerPersona || 3;
    let progressBar;
    if (process.env.LOG_LEVEL !== 'debug') {
        const cliProgress = await Promise.resolve().then(() => __importStar(require('cli-progress')));
        progressBar = new cliProgress.SingleBar({}, 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.`);
    // Consider the following prompt for an LLM application: {{prompt}}. List up to 5 user personas that would send this prompt.
    logger_1.default.debug(`\nGenerating user personas from ${prompts.length} prompts...`);
    const provider = new openai_1.OpenAiChatCompletionProvider(SYNTHESIZE_DEFAULT_PROVIDER, {
        config: {
            temperature: 1.0,
            response_format: {
                type: 'json_object',
            },
        },
    });
    const promptsString = `<Prompts>
${prompts.map((prompt) => `<Prompt>\n${prompt}\n</Prompt>`).join('\n')}
</Prompts>`;
    const resp = await provider.callApi(`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[]}`);
    const personas = JSON.parse(resp.output).personas;
    logger_1.default.debug(`\nGenerated ${personas.length} personas:\n${personas.map((p) => `  - ${p}`).join('\n')}`);
    if (progressBar) {
        progressBar.increment();
    }
    // Extract variable names from the nunjucks template in the prompts
    const variableRegex = /{{\s*(\w+)\s*}}/g;
    const variables = new Set();
    for (const prompt of prompts) {
        let match;
        while ((match = variableRegex.exec(prompt)) !== null) {
            variables.add(match[1]);
        }
    }
    logger_1.default.debug(`\nExtracted ${variables.size} variables from prompts:\n${Array.from(variables)
        .map((v) => `  - ${v}`)
        .join('\n')}`);
    const existingTests = `Here are some existing tests:` +
        tests
            .map((test) => {
            if (!test.vars) {
                return;
            }
            return `<Test>
${JSON.stringify(test.vars, null, 2)}
</Test>
    `;
        })
            .filter(Boolean)
            .slice(0, 100)
            .join('\n');
    // For each user persona, we will generate a map of variable names to values
    const testCaseVars = [];
    for (let i = 0; i < personas.length; i++) {
        const persona = personas[i];
        logger_1.default.debug(`\nGenerating test cases for persona ${i + 1}...`);
        // Construct the prompt for the LLM to generate variable values
        const personaPrompt = `Consider ${prompts.length > 1 ? 'these prompts' : 'this prompt'}, which contains some {{variables}}: 
${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 || ''}

Your response should contain a JSON map of variable names to values, of the form {vars: {${Array.from(variables)
            .map((varName) => `${varName}: string`)
            .join(', ')}}[]}`;
        // Call the LLM API with the constructed prompt
        const personaResponse = await provider.callApi(personaPrompt);
        const parsed = JSON.parse(personaResponse.output);
        for (const vars of parsed.vars) {
            logger_1.default.debug(`${JSON.stringify(vars, null, 2)}`);
            testCaseVars.push(vars);
            if (progressBar) {
                progressBar.increment();
            }
        }
    }
    if (progressBar) {
        progressBar.stop();
    }
    // Dedup test case vars
    const uniqueTestCaseStrings = new Set(testCaseVars.map((testCase) => JSON.stringify(testCase)));
    const dedupedTestCaseVars = Array.from(uniqueTestCaseStrings).map((testCase) => JSON.parse(testCase));
    return dedupedTestCaseVars;
}
exports.synthesize = synthesize;
//# sourceMappingURL=testCases.js.map