"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.AzureChatCompletionProvider = void 0;
const cache_1 = require("../../cache");
const envars_1 = require("../../envars");
const logger_1 = __importDefault(require("../../logger"));
const util_1 = require("../../util");
const file_1 = require("../../util/file");
const finishReason_1 = require("../../util/finishReason");
const invariant_1 = __importDefault(require("../../util/invariant"));
const functionCallbackUtils_1 = require("../functionCallbackUtils");
const client_1 = require("../mcp/client");
const transform_1 = require("../mcp/transform");
const shared_1 = require("../shared");
const defaults_1 = require("./defaults");
const generic_1 = require("./generic");
const util_2 = require("./util");
class AzureChatCompletionProvider extends generic_1.AzureGenericProvider {
    constructor(...args) {
        super(...args);
        this.mcpClient = null;
        this.functionCallbackHandler = new functionCallbackUtils_1.FunctionCallbackHandler();
        if (this.config.mcp?.enabled) {
            this.initializationPromise = this.initializeMCP();
        }
    }
    async initializeMCP() {
        this.mcpClient = new client_1.MCPClient(this.config.mcp);
        await this.mcpClient.initialize();
    }
    async cleanup() {
        if (this.mcpClient) {
            await this.initializationPromise;
            await this.mcpClient.cleanup();
            this.mcpClient = null;
        }
    }
    /**
     * Check if the current deployment is configured as a reasoning model
     */
    isReasoningModel() {
        return !!this.config.isReasoningModel || !!this.config.o1;
    }
    getOpenAiBody(prompt, context, callApiOptions) {
        const config = {
            ...this.config,
            ...context?.prompt?.config,
        };
        // Parse chat prompt
        let messages = (0, shared_1.parseChatPrompt)(prompt, [{ role: 'user', content: prompt }]);
        // Inject system prompt if configured
        if (config.systemPrompt) {
            // Check if there's already a system message
            const existingSystemMessageIndex = messages.findIndex((msg) => msg.role === 'system');
            if (existingSystemMessageIndex >= 0) {
                // Replace existing system message
                messages[existingSystemMessageIndex] = {
                    role: 'system',
                    content: config.systemPrompt,
                };
            }
            else {
                // Prepend new system message
                messages = [
                    {
                        role: 'system',
                        content: config.systemPrompt,
                    },
                    ...messages,
                ];
            }
        }
        // Response format with variable rendering
        const responseFormat = config.response_format
            ? {
                response_format: (0, file_1.maybeLoadFromExternalFile)((0, util_1.renderVarsInObject)(config.response_format, context?.vars)),
            }
            : {};
        // Check if this is configured as a reasoning model
        const isReasoningModel = this.isReasoningModel();
        // Get max tokens based on model type
        const maxTokens = config.max_tokens ?? (0, envars_1.getEnvInt)('OPENAI_MAX_TOKENS', 1024);
        const maxCompletionTokens = config.max_completion_tokens;
        // Get reasoning effort for reasoning models
        const reasoningEffort = config.reasoning_effort ?? 'medium';
        // --- MCP tool injection logic ---
        const mcpTools = this.mcpClient ? (0, transform_1.transformMCPToolsToOpenAi)(this.mcpClient.getAllTools()) : [];
        const fileTools = config.tools
            ? (0, util_1.maybeLoadToolsFromExternalFile)(config.tools, context?.vars) || []
            : [];
        const allTools = [...mcpTools, ...fileTools];
        // --- End MCP tool injection logic ---
        // Build the request body
        const body = {
            model: this.deploymentName,
            messages,
            ...(isReasoningModel
                ? {
                    max_completion_tokens: maxCompletionTokens ?? maxTokens,
                    reasoning_effort: (0, util_1.renderVarsInObject)(reasoningEffort, context?.vars),
                }
                : {
                    max_tokens: maxTokens,
                    temperature: config.temperature ?? (0, envars_1.getEnvFloat)('OPENAI_TEMPERATURE', 0),
                }),
            top_p: config.top_p ?? (0, envars_1.getEnvFloat)('OPENAI_TOP_P', 1),
            presence_penalty: config.presence_penalty ?? (0, envars_1.getEnvFloat)('OPENAI_PRESENCE_PENALTY', 0),
            frequency_penalty: config.frequency_penalty ?? (0, envars_1.getEnvFloat)('OPENAI_FREQUENCY_PENALTY', 0),
            ...(config.seed === undefined ? {} : { seed: config.seed }),
            ...(config.functions
                ? {
                    functions: (0, file_1.maybeLoadFromExternalFile)((0, util_1.renderVarsInObject)(config.functions, context?.vars)),
                }
                : {}),
            ...(config.function_call ? { function_call: config.function_call } : {}),
            ...(allTools.length > 0 ? { tools: allTools } : {}),
            ...(config.tool_choice ? { tool_choice: config.tool_choice } : {}),
            ...(config.deployment_id ? { deployment_id: config.deployment_id } : {}),
            ...(config.dataSources ? { dataSources: config.dataSources } : {}),
            ...responseFormat,
            ...(callApiOptions?.includeLogProbs ? { logprobs: callApiOptions.includeLogProbs } : {}),
            ...(config.stop ? { stop: config.stop } : {}),
            ...(config.passthrough || {}),
        };
        logger_1.default.debug(`Azure API request body: ${JSON.stringify(body)}`);
        return { body, config };
    }
    async callApi(prompt, context, callApiOptions) {
        if (this.initializationPromise) {
            await this.initializationPromise;
        }
        await this.ensureInitialized();
        (0, invariant_1.default)(this.authHeaders, 'auth headers are not initialized');
        if (!this.getApiBaseUrl()) {
            throw new Error('Azure API host must be set.');
        }
        const { body, config } = this.getOpenAiBody(prompt, context, callApiOptions);
        let data;
        let cached = false;
        let httpStatus;
        try {
            const url = config.dataSources
                ? `${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/extensions/chat/completions?api-version=${config.apiVersion || defaults_1.DEFAULT_AZURE_API_VERSION}`
                : `${this.getApiBaseUrl()}/openai/deployments/${this.deploymentName}/chat/completions?api-version=${config.apiVersion || defaults_1.DEFAULT_AZURE_API_VERSION}`;
            const { data: responseData, cached: isCached, status, } = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    ...this.authHeaders,
                    ...this.config.headers,
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS, 'json', context?.bustCache ?? context?.debug);
            cached = isCached;
            httpStatus = status;
            // Handle the response data
            if (typeof responseData === 'string') {
                try {
                    data = JSON.parse(responseData);
                }
                catch {
                    return {
                        error: `API returned invalid JSON response (status ${status}): ${responseData}\n\nRequest body: ${JSON.stringify(body, null, 2)}`,
                    };
                }
            }
            else {
                data = responseData;
            }
        }
        catch (err) {
            return {
                error: `API call error: ${err instanceof Error ? err.message : String(err)}`,
            };
        }
        logger_1.default.debug(`Azure API response (status ${httpStatus}): ${JSON.stringify(data)}`);
        // Inputs and outputs can be flagged by content filters.
        // See https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/content-filter
        let flaggedInput = false;
        let flaggedOutput = false;
        let output = '';
        let logProbs;
        let finishReason;
        try {
            if (data.error) {
                // Was the input prompt deemed inappropriate?
                if (data.error.status === 400 && data.error.code === finishReason_1.FINISH_REASON_MAP.content_filter) {
                    flaggedInput = true;
                    output = data.error.message;
                    finishReason = finishReason_1.FINISH_REASON_MAP.content_filter;
                }
                else {
                    return {
                        error: `API response error: ${data.error.code} ${data.error.message}`,
                    };
                }
            }
            else {
                const hasDataSources = !!config.dataSources;
                const choice = hasDataSources
                    ? data.choices.find((choice) => choice.message.role === 'assistant')
                    : data.choices[0];
                const message = choice?.message;
                // NOTE: The `n` parameter is currently (250709) not supported; if and when it is, `finish_reason` must be
                // checked on all choices; in other words, if n>1 responses are requested, n>1 responses can trigger filters.
                finishReason = (0, finishReason_1.normalizeFinishReason)(choice?.finish_reason);
                // Handle structured output
                output = message?.content;
                // Check for errors indicating that the content filters did not run on the completion.
                if (choice.content_filter_results && choice.content_filter_results.error) {
                    const { code, message } = choice.content_filter_results.error;
                    logger_1.default.warn(`Content filtering system is down or otherwise unable to complete the request in time: ${code} ${message}`);
                }
                else {
                    // Was the completion filtered?
                    flaggedOutput = finishReason === finishReason_1.FINISH_REASON_MAP.content_filter;
                }
                if (output == null) {
                    // Handle tool_calls and function_call
                    const toolCalls = message.tool_calls;
                    const functionCall = message.function_call;
                    // Process function/tool calls if callbacks are configured
                    if (config.functionToolCallbacks && (toolCalls || functionCall)) {
                        // Combine all calls into a single array for processing
                        const allCalls = [];
                        if (toolCalls) {
                            allCalls.push(...(Array.isArray(toolCalls) ? toolCalls : [toolCalls]));
                        }
                        if (functionCall) {
                            allCalls.push(functionCall);
                        }
                        output = await this.functionCallbackHandler.processCalls(allCalls.length === 1 ? allCalls[0] : allCalls, config.functionToolCallbacks);
                    }
                    else {
                        // No callbacks configured, return raw tool/function calls
                        output = toolCalls ?? functionCall;
                    }
                }
                else if (config.response_format?.type === 'json_schema' ||
                    config.response_format?.type === 'json_object') {
                    try {
                        output = JSON.parse(output);
                    }
                    catch (err) {
                        logger_1.default.error(`Failed to parse JSON output: ${err}. Output was: ${output}`);
                    }
                }
                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,
                        ...(data.usage?.completion_tokens_details
                            ? {
                                completionDetails: {
                                    reasoning: data.usage.completion_tokens_details.reasoning_tokens,
                                    acceptedPrediction: data.usage.completion_tokens_details.accepted_prediction_tokens,
                                    rejectedPrediction: data.usage.completion_tokens_details.rejected_prediction_tokens,
                                },
                            }
                            : {}),
                    },
                cached,
                logProbs,
                finishReason,
                cost: (0, util_2.calculateAzureCost)(this.deploymentName, config, data.usage?.prompt_tokens, data.usage?.completion_tokens),
                guardrails: {
                    flagged: flaggedInput || flaggedOutput,
                    flaggedInput,
                    flaggedOutput,
                },
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
}
exports.AzureChatCompletionProvider = AzureChatCompletionProvider;
//# sourceMappingURL=chat.js.map