"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
var __asyncValues = (this && this.__asyncValues) || function (o) {
    if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
    var m = o[Symbol.asyncIterator], i;
    return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
    function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
    function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
};
Object.defineProperty(exports, "__esModule", { value: true });
const vitest_1 = require("vitest");
const __1 = require("../");
let cohere;
const config = {
    awsRegion: "us-east-1",
    awsAccessKey: "...",
    awsSecretKey: "...",
    awsSessionToken: "...",
};
const models = {
    bedrock: {
        generate: "cohere.command-text-v14",
        embed: "cohere.embed-multilingual-v3",
        chat: "cohere.command-r-plus-v1:0",
        rerank: "cohere.rerank-v3-5:0",
    },
    sagemaker: {
        generate: "cohere-command-light",
        embed: "cohere-embed-multilingual-v3",
        chat: "cohere-command-plus",
        rerank: "cohere.rerank-v1",
    },
};
// skip until we have the right auth in ci
vitest_1.describe.each(["bedrock"])("test sdk", (platform) => {
    cohere = {
        "bedrock": new __1.BedrockClient(config),
        "sagemaker": new __1.SagemakerClient(config)
    }[platform];
    vitest_1.test.skip("generate works", () => __awaiter(void 0, void 0, void 0, function* () {
        const generate = yield cohere.generate({
            prompt: "Please explain to me how LLMs work",
            temperature: 0,
            model: models[platform].generate,
        });
        (0, vitest_1.expect)(generate.generations[0].text).toBeDefined();
    }));
    vitest_1.test.skip("generate stream works", () => __awaiter(void 0, void 0, void 0, function* () {
        var _a, e_1, _b, _c;
        const generate = yield cohere.generateStream({
            prompt: "Please explain to me how LLMs work",
            temperature: 0,
            model: models[platform].generate,
        });
        const chunks = [];
        try {
            for (var _d = true, generate_1 = __asyncValues(generate), generate_1_1; generate_1_1 = yield generate_1.next(), _a = generate_1_1.done, !_a; _d = true) {
                _c = generate_1_1.value;
                _d = false;
                const chunk = _c;
                chunks.push(chunk);
            }
        }
        catch (e_1_1) { e_1 = { error: e_1_1 }; }
        finally {
            try {
                if (!_d && !_a && (_b = generate_1.return)) yield _b.call(generate_1);
            }
            finally { if (e_1) throw e_1.error; }
        }
        (0, vitest_1.expect)(chunks[0].eventType).toMatchInlineSnapshot(`"stream-start"`);
        (0, vitest_1.expect)(chunks[1].eventType).toMatchInlineSnapshot(`"text-generation"`);
        (0, vitest_1.expect)(chunks[chunks.length - 1].eventType).toMatchInlineSnapshot(`"stream-end"`);
    }));
    vitest_1.test.skip("embed works", () => __awaiter(void 0, void 0, void 0, function* () {
        var _a, _b;
        const embed = yield cohere.embed({
            texts: ["hello", "goodbye"],
            model: models[platform].embed,
            inputType: "search_document",
        });
        if (embed.responseType === "embeddings_by_type") {
            (0, vitest_1.expect)((_b = (_a = embed.embeddings) === null || _a === void 0 ? void 0 : _a.float) === null || _b === void 0 ? void 0 : _b[0]).toBeDefined();
        }
    }));
    vitest_1.test.skip("chat works", () => __awaiter(void 0, void 0, void 0, function* () {
        const chat = yield cohere.chat({
            model: models[platform].chat,
            message: "send me a short message",
            temperature: 0,
        });
    }));
    vitest_1.test.skip("rerank works", () => __awaiter(void 0, void 0, void 0, function* () {
        const rerank = yield cohere.v2.rerank({
            model: models[platform].rerank,
            documents: [
                "Carson City is the capital city of the American state of Nevada.",
                "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
                "Washington, D.C. d (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
                "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
            ],
            query: "What is the capital of the United States?",
            topN: 3,
        });
        (0, vitest_1.expect)(rerank.results).toBeDefined();
    }));
    vitest_1.test.skip("chat stream works", () => __awaiter(void 0, void 0, void 0, function* () {
        var _a, e_2, _b, _c;
        const chat = yield cohere.chatStream({
            model: models[platform].chat,
            message: "send me a short message",
            temperature: 0,
        });
        const chunks = [];
        try {
            for (var _d = true, chat_1 = __asyncValues(chat), chat_1_1; chat_1_1 = yield chat_1.next(), _a = chat_1_1.done, !_a; _d = true) {
                _c = chat_1_1.value;
                _d = false;
                const chunk = _c;
                chunks.push(chunk);
            }
        }
        catch (e_2_1) { e_2 = { error: e_2_1 }; }
        finally {
            try {
                if (!_d && !_a && (_b = chat_1.return)) yield _b.call(chat_1);
            }
            finally { if (e_2) throw e_2.error; }
        }
        (0, vitest_1.expect)(chunks[0].eventType).toMatchInlineSnapshot(`"text-generation"`);
        (0, vitest_1.expect)(chunks[chunks.length - 1].eventType).toMatchInlineSnapshot(`"stream-end"`);
    }));
}, 5000);
