"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.AzureOpenAiAssistantProvider = exports.AzureOpenAiChatCompletionProvider = exports.AzureOpenAiCompletionProvider = exports.AzureOpenAiEmbeddingProvider = void 0;
const tiny_invariant_1 = __importDefault(require("tiny-invariant"));
const logger_1 = __importDefault(require("../logger"));
const cache_1 = require("../cache");
const util_1 = require("../util");
const shared_1 = require("./shared");
class AzureOpenAiGenericProvider {
    constructor(deploymentName, options = {}) {
        const { config, id, env } = options;
        this.env = env;
        this.deploymentName = deploymentName;
        this.apiHost =
            config?.apiHost || env?.AZURE_OPENAI_API_HOST || process.env.AZURE_OPENAI_API_HOST;
        this.apiBaseUrl =
            config?.apiBaseUrl ||
                env?.AZURE_OPENAI_API_BASE_URL ||
                env?.AZURE_OPENAI_BASE_URL ||
                process.env.AZURE_OPENAI_API_BASE_URL ||
                process.env.AZURE_OPENAI_BASE_URL;
        this.config = config || {};
        this.id = id ? () => id : this.id;
    }
    async getApiKey() {
        if (!this._cachedApiKey) {
            if (this.config?.azureClientSecret &&
                this.config?.azureClientId &&
                this.config?.azureTenantId) {
                const { ClientSecretCredential } = await Promise.resolve().then(() => __importStar(require('@azure/identity')));
                const credential = new ClientSecretCredential(this.config.azureTenantId, this.config.azureClientId, this.config.azureClientSecret, {
                    authorityHost: this.config.azureAuthorityHost || 'https://login.microsoftonline.com',
                });
                this._cachedApiKey = (await credential.getToken(this.config.azureTokenScope || 'https://cognitiveservices.azure.com/.default')).token;
            }
            else {
                this._cachedApiKey =
                    this.config?.apiKey ||
                        (this.config?.apiKeyEnvar
                            ? process.env[this.config.apiKeyEnvar] ||
                                this.env?.[this.config.apiKeyEnvar]
                            : undefined) ||
                        this.env?.AZURE_OPENAI_API_KEY ||
                        process.env.AZURE_OPENAI_API_KEY;
                if (!this._cachedApiKey) {
                    throw new Error('Azure OpenAI API key must be set');
                }
            }
        }
        return this._cachedApiKey;
    }
    getApiBaseUrl() {
        return this.apiBaseUrl || `https://${this.apiHost}`;
    }
    id() {
        return `azureopenai:${this.deploymentName}`;
    }
    toString() {
        return `[Azure OpenAI Provider ${this.deploymentName}]`;
    }
    // @ts-ignore: Params are not used in this implementation
    async callApi(prompt, context, callApiOptions) {
        throw new Error('Not implemented');
    }
}
class AzureOpenAiEmbeddingProvider extends AzureOpenAiGenericProvider {
    async callEmbeddingApi(text) {
        const apiKey = await this.getApiKey();
        if (!apiKey) {
            throw new Error('Azure OpenAI API key must be set for similarity comparison');
        }
        if (!this.getApiBaseUrl()) {
            throw new Error('Azure OpenAI API host must be set');
        }
        const body = {
            input: text,
            model: this.deploymentName,
        };
        let data, cached = false;
        try {
            ({ data, cached } = (await (0, cache_1.fetchWithCache)(`${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/embeddings?api-version=${this.config.apiVersion || '2023-12-01-preview'}`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'api-key': apiKey,
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS)));
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
                tokenUsage: {
                    total: 0,
                    prompt: 0,
                    completion: 0,
                },
            };
        }
        logger_1.default.debug(`\tAzure OpenAI API response (embeddings): ${JSON.stringify(data)}`);
        try {
            const embedding = data?.data?.[0]?.embedding;
            if (!embedding) {
                throw new Error('No embedding returned');
            }
            const ret = {
                embedding,
                tokenUsage: cached
                    ? { cached: data.usage.total_tokens, total: data.usage.total_tokens }
                    : {
                        total: data.usage.total_tokens,
                        prompt: data.usage.prompt_tokens,
                        completion: data.usage.completion_tokens,
                    },
            };
            return ret;
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
                tokenUsage: cached
                    ? {
                        cached: data.usage.total_tokens,
                        total: data.usage.total_tokens,
                    }
                    : {
                        total: data?.usage?.total_tokens,
                        prompt: data?.usage?.prompt_tokens,
                        completion: data?.usage?.completion_tokens,
                    },
            };
        }
    }
}
exports.AzureOpenAiEmbeddingProvider = AzureOpenAiEmbeddingProvider;
class AzureOpenAiCompletionProvider extends AzureOpenAiGenericProvider {
    async callApi(prompt, context, callApiOptions) {
        const apiKey = await this.getApiKey();
        if (!apiKey) {
            throw new Error('Azure OpenAI API key is not set. Set AZURE_OPENAI_API_KEY environment variable or pass it as an argument to the constructor.');
        }
        if (!this.getApiBaseUrl()) {
            throw new Error('Azure OpenAI API host must be set');
        }
        let stop;
        try {
            stop = process.env.OPENAI_STOP
                ? JSON.parse(process.env.OPENAI_STOP)
                : this.config?.stop || ['<|im_end|>', '<|endoftext|>'];
        }
        catch (err) {
            throw new Error(`OPENAI_STOP is not a valid JSON string: ${err}`);
        }
        const body = {
            model: this.deploymentName,
            prompt,
            max_tokens: this.config.max_tokens ?? parseInt(process.env.OPENAI_MAX_TOKENS || '1024'),
            temperature: this.config.temperature ?? parseFloat(process.env.OPENAI_TEMPERATURE || '0'),
            top_p: this.config.top_p ?? parseFloat(process.env.OPENAI_TOP_P || '1'),
            presence_penalty: this.config.presence_penalty ?? parseFloat(process.env.OPENAI_PRESENCE_PENALTY || '0'),
            frequency_penalty: this.config.frequency_penalty ?? parseFloat(process.env.OPENAI_FREQUENCY_PENALTY || '0'),
            best_of: this.config.best_of ?? parseInt(process.env.OPENAI_BEST_OF || '1'),
            ...(this.config.seed !== undefined ? { seed: this.config.seed } : {}),
            ...(this.config.deployment_id ? { deployment_id: this.config.deployment_id } : {}),
            ...(this.config.dataSources ? { dataSources: this.config.dataSources } : {}),
            ...(this.config.response_format ? { response_format: this.config.response_format } : {}),
            ...(callApiOptions?.includeLogProbs ? { logprobs: callApiOptions.includeLogProbs } : {}),
            ...(stop ? { stop } : {}),
            ...(this.config.passthrough || {}),
        };
        logger_1.default.debug(`Calling Azure OpenAI API: ${JSON.stringify(body)}`);
        let data, cached = false;
        try {
            ({ data, cached } = (await (0, cache_1.fetchWithCache)(`${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/completions?api-version=${this.config.apiVersion || '2023-12-01-preview'}`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'api-key': apiKey,
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS)));
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tAzure OpenAI API response: ${JSON.stringify(data)}`);
        try {
            return {
                output: data.choices[0].text,
                tokenUsage: cached
                    ? { cached: data.usage.total_tokens, total: data.usage.total_tokens }
                    : {
                        total: data.usage.total_tokens,
                        prompt: data.usage.prompt_tokens,
                        completion: data.usage.completion_tokens,
                    },
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
}
exports.AzureOpenAiCompletionProvider = AzureOpenAiCompletionProvider;
class AzureOpenAiChatCompletionProvider extends AzureOpenAiGenericProvider {
    async callApi(prompt, context, callApiOptions) {
        const apiKey = await this.getApiKey();
        if (!apiKey) {
            throw new Error('Azure OpenAI API key is not set. Set the AZURE_OPENAI_API_KEY environment variable or add `apiKey` to the provider config.');
        }
        if (!this.getApiBaseUrl()) {
            throw new Error('Azure OpenAI API host must be set');
        }
        const messages = (0, shared_1.parseChatPrompt)(prompt, [{ role: 'user', content: prompt }]);
        let stop;
        try {
            stop = process.env.OPENAI_STOP ? JSON.parse(process.env.OPENAI_STOP) : this.config?.stop;
        }
        catch (err) {
            throw new Error(`OPENAI_STOP is not a valid JSON string: ${err}`);
        }
        const body = {
            model: this.deploymentName,
            messages: messages,
            max_tokens: this.config.max_tokens ?? parseInt(process.env.OPENAI_MAX_TOKENS || '1024'),
            temperature: this.config.temperature ?? parseFloat(process.env.OPENAI_TEMPERATURE || '0'),
            top_p: this.config.top_p ?? parseFloat(process.env.OPENAI_TOP_P || '1'),
            presence_penalty: this.config.presence_penalty ?? parseFloat(process.env.OPENAI_PRESENCE_PENALTY || '0'),
            frequency_penalty: this.config.frequency_penalty ?? parseFloat(process.env.OPENAI_FREQUENCY_PENALTY || '0'),
            ...(this.config.functions
                ? { functions: (0, util_1.renderVarsInObject)(this.config.functions, context?.vars) }
                : {}),
            ...(this.config.function_call ? { function_call: this.config.function_call } : {}),
            ...(this.config.seed !== undefined ? { seed: this.config.seed } : {}),
            ...(this.config.tools ? { tools: (0, util_1.renderVarsInObject)(this.config.tools, context?.vars) } : {}),
            ...(this.config.tool_choice ? { tool_choice: this.config.tool_choice } : {}),
            ...(this.config.deployment_id ? { deployment_id: this.config.deployment_id } : {}),
            ...(this.config.dataSources ? { dataSources: this.config.dataSources } : {}),
            ...(this.config.response_format ? { response_format: this.config.response_format } : {}),
            ...(callApiOptions?.includeLogProbs ? { logprobs: callApiOptions.includeLogProbs } : {}),
            ...(this.config.stop ? { stop: this.config.stop } : {}),
            ...(this.config.passthrough || {}),
        };
        logger_1.default.debug(`Calling Azure OpenAI API: ${JSON.stringify(body)}`);
        let data, cached = false;
        try {
            const url = this.config.dataSources
                ? `${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/extensions/chat/completions?api-version=${this.config.apiVersion || '2023-12-01-preview'}`
                : `${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/chat/completions?api-version=${this.config.apiVersion || '2023-12-01-preview'}`;
            ({ data, cached } = (await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'api-key': apiKey,
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS)));
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tAzure OpenAI API response: ${JSON.stringify(data)}`);
        try {
            if (data.error) {
                return {
                    error: `API response error: ${data.error.code} ${data.error.message}`,
                };
            }
            const hasDataSources = !!this.config.dataSources;
            const message = hasDataSources
                ? data.choices.find((choice) => choice.message.role === 'assistant')?.message
                : data.choices[0].message;
            const output = message.content == null
                ? message.tool_calls == null
                    ? message.function_call
                    : message.tool_calls
                : message.content;
            const logProbs = data.choices[0].logprobs?.content?.map((logProbObj) => logProbObj.logprob);
            return {
                output,
                tokenUsage: cached
                    ? { cached: data.usage?.total_tokens, total: data?.usage?.total_tokens }
                    : {
                        total: data.usage?.total_tokens,
                        prompt: data.usage?.prompt_tokens,
                        completion: data.usage?.completion_tokens,
                    },
                cached,
                logProbs,
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
}
exports.AzureOpenAiChatCompletionProvider = AzureOpenAiChatCompletionProvider;
// See https://learn.microsoft.com/en-us/javascript/api/overview/azure/openai-assistants-readme?view=azure-node-preview
class AzureOpenAiAssistantProvider extends AzureOpenAiGenericProvider {
    constructor(deploymentName, options = {}) {
        super(deploymentName, options);
        this.initializationPromise = null;
        this.assistantConfig = options.config || {};
        this.initializationPromise = this.initialize();
    }
    async initialize() {
        const apiKey = await this.getApiKey();
        if (!apiKey) {
            throw new Error('Azure OpenAI API key must be set');
        }
        const { AssistantsClient, AzureKeyCredential } = await Promise.resolve().then(() => __importStar(require('@azure/openai-assistants')));
        this.assistantsClient = new AssistantsClient(this.getApiBaseUrl(), new AzureKeyCredential(apiKey));
        this.initializationPromise = null;
    }
    async ensureInitialized() {
        if (this.initializationPromise) {
            await this.initializationPromise;
        }
    }
    async callApi(prompt, context, callApiOptions) {
        await this.ensureInitialized();
        (0, tiny_invariant_1.default)(this.assistantsClient, 'Assistants client not initialized');
        if (!this.getApiBaseUrl()) {
            throw new Error('Azure OpenAI API host must be set');
        }
        const assistantId = this.deploymentName;
        const assistantThread = await this.assistantsClient.createThread();
        await this.assistantsClient.createMessage(assistantThread.id, 'user', prompt);
        let run = await this.assistantsClient.createRun(assistantThread.id, {
            assistantId,
        });
        logger_1.default.debug(`\tAzure thread run API response: ${JSON.stringify(run)}`);
        while (run.status === 'in_progress' ||
            run.status === 'queued' ||
            run.status === 'requires_action') {
            if (run.status === 'requires_action') {
                const requiredAction = run.requiredAction;
                (0, tiny_invariant_1.default)(requiredAction, 'Run requires action but no action is provided');
                if (requiredAction === null || requiredAction.type !== 'submit_tool_outputs') {
                    break;
                }
                const functionCallsWithCallbacks = requiredAction.submitToolOutputs?.toolCalls.filter((toolCall) => {
                    return (toolCall.type === 'function' &&
                        toolCall.function.name in (this.assistantConfig.functionToolCallbacks ?? {}));
                });
                if (!functionCallsWithCallbacks || functionCallsWithCallbacks.length === 0) {
                    break;
                }
                logger_1.default.debug(`Calling functionToolCallbacks for functions: ${functionCallsWithCallbacks.map(({ function: { name } }) => name)}`);
                const toolOutputs = await Promise.all(functionCallsWithCallbacks.map(async (toolCall) => {
                    logger_1.default.debug(`Calling functionToolCallbacks[${toolCall.function.name}]('${toolCall.function.arguments}')`);
                    const result = await this.assistantConfig.functionToolCallbacks[toolCall.function.name](toolCall.function.arguments);
                    return {
                        tool_call_id: toolCall.id,
                        output: result,
                    };
                }));
                logger_1.default.debug(`Calling Azure API, submitting tool outputs for ${run.threadId}: ${JSON.stringify(toolOutputs)}`);
                run = await this.assistantsClient.submitToolOutputsToRun(run.threadId, run.id, toolOutputs);
            }
            await new Promise((resolve) => setTimeout(resolve, 1000));
            logger_1.default.debug(`Calling Azure API, getting thread run ${run.id} status`);
            run = await this.assistantsClient.getRun(run.threadId, run.id);
            logger_1.default.debug(`\tAzure thread run API response: ${JSON.stringify(run)}`);
        }
        if (run.status !== 'completed' && run.status !== 'requires_action') {
            if (run.lastError) {
                return {
                    error: `Thread run failed: ${run.lastError.message}`,
                };
            }
            return {
                error: `Thread run failed: ${run.status}`,
            };
        }
        logger_1.default.debug(`Calling Azure API, getting thread run steps for ${run.threadId}`);
        const steps = await this.assistantsClient.listRunSteps(run.threadId, run.id, { order: 'asc' });
        logger_1.default.debug(`\tAzure thread run steps API response: ${JSON.stringify(steps)}`);
        const outputBlocks = [];
        for (const step of steps.data) {
            if (step.type === 'message_creation') {
                logger_1.default.debug(`Calling Azure API, getting message ${step.id}`);
                const stepDetails = step.stepDetails;
                // Bug in the API: the field is currently `message_id` even though it's documented as `messageId`
                const messageId = stepDetails.messageCreation.message_id || stepDetails.messageCreation.messageId;
                const message = await this.assistantsClient.getMessage(run.threadId, messageId);
                logger_1.default.debug(`\tAzure thread run step message API response: ${JSON.stringify(message)}`);
                const content = message.content
                    .map((content) => content.type === 'text' ? content.text.value : `<${content.type} output>`)
                    .join('\n');
                outputBlocks.push(`[${(0, shared_1.toTitleCase)(message.role)}] ${content}`);
            }
            else if (step.type === 'tool_calls') {
                for (const toolCall of step.stepDetails.toolCalls) {
                    if (toolCall.type === 'function') {
                        outputBlocks.push(`[Call function ${toolCall.function.name} with arguments ${toolCall.function.arguments}]`);
                        outputBlocks.push(`[Function output: ${toolCall.function.output}]`);
                    }
                    else if (toolCall.type === 'retrieval') {
                        outputBlocks.push(`[Ran retrieval]`);
                    }
                    else if (toolCall.type === 'code_interpreter') {
                        const output = toolCall.codeInterpreter.outputs
                            .map((output) => (output.type === 'logs' ? output.logs : `<${output.type} output>`))
                            .join('\n');
                        outputBlocks.push(`[Code interpreter input]`);
                        outputBlocks.push(toolCall.codeInterpreter.input);
                        outputBlocks.push(`[Code interpreter output]`);
                        outputBlocks.push(output);
                    }
                    else {
                        outputBlocks.push(`[Unknown tool call type: ${toolCall.type}]`);
                    }
                }
            }
            else {
                outputBlocks.push(`[Unknown step type: ${step.type}]`);
            }
        }
        return {
            output: outputBlocks.join('\n\n').trim(),
            /*
            tokenUsage: {
              total: data.usage.total_tokens,
              prompt: data.usage.prompt_tokens,
              completion: data.usage.completion_tokens,
            },
            */
        };
    }
}
exports.AzureOpenAiAssistantProvider = AzureOpenAiAssistantProvider;
//# sourceMappingURL=azureopenai.js.map