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());
    });
};
import { logger, } from "adminforth";
import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
function isLangChainAgentCompletionAdapter(adapter) {
    return typeof adapter
        .getLangChainAgentSpec === "function";
}
function getAgentChatModelSpec(params) {
    return __awaiter(this, void 0, void 0, function* () {
        var _a;
        const spec = yield params.adapter.getLangChainAgentSpec({
            maxTokens: params.maxTokens,
            purpose: params.purpose,
        });
        return {
            model: spec.model,
            middleware: ((_a = spec.middleware) !== null && _a !== void 0 ? _a : []),
        };
    });
}
function getFiniteNumber(value) {
    return typeof value === "number" && Number.isFinite(value)
        ? value
        : undefined;
}
function extractTokenUsage(output) {
    var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l;
    const llmOutputTokenUsage = (_a = output.llmOutput) === null || _a === void 0 ? void 0 : _a.tokenUsage;
    const promptTokens = getFiniteNumber(llmOutputTokenUsage === null || llmOutputTokenUsage === void 0 ? void 0 : llmOutputTokenUsage.promptTokens);
    const completionTokens = getFiniteNumber(llmOutputTokenUsage === null || llmOutputTokenUsage === void 0 ? void 0 : llmOutputTokenUsage.completionTokens);
    if (promptTokens !== undefined || completionTokens !== undefined) {
        return {
            InputTokens: promptTokens !== null && promptTokens !== void 0 ? promptTokens : 0,
            outputTokens: completionTokens !== null && completionTokens !== void 0 ? completionTokens : 0,
        };
    }
    let InputTokens = 0;
    let outputTokens = 0;
    for (const generationBatch of output.generations) {
        for (const generation of generationBatch) {
            if (!("message" in generation) || !generation.message) {
                continue;
            }
            const message = generation.message;
            InputTokens +=
                (_f = (_c = getFiniteNumber((_b = message.usage_metadata) === null || _b === void 0 ? void 0 : _b.input_tokens)) !== null && _c !== void 0 ? _c : getFiniteNumber((_e = (_d = message.response_metadata) === null || _d === void 0 ? void 0 : _d.tokenUsage) === null || _e === void 0 ? void 0 : _e.promptTokens)) !== null && _f !== void 0 ? _f : 0;
            outputTokens +=
                (_l = (_h = getFiniteNumber((_g = message.usage_metadata) === null || _g === void 0 ? void 0 : _g.output_tokens)) !== null && _h !== void 0 ? _h : getFiniteNumber((_k = (_j = message.response_metadata) === null || _j === void 0 ? void 0 : _j.tokenUsage) === null || _k === void 0 ? void 0 : _k.completionTokens)) !== null && _l !== void 0 ? _l : 0;
        }
    }
    return { InputTokens, outputTokens };
}
class AgentLlmMetricsLogger extends BaseCallbackHandler {
    constructor() {
        super(...arguments);
        this.name = "AgentLlmMetricsLogger";
        this.lc_prefer_streaming = true;
        this.pendingRuns = new Map();
    }
    handleLLMStart(_llm, _prompts, runId) {
        return __awaiter(this, void 0, void 0, function* () {
            this.pendingRuns.set(runId, { startedAt: Date.now() });
        });
    }
    handleLLMNewToken(_token, _chunk, runId) {
        return __awaiter(this, void 0, void 0, function* () {
            const pendingRun = this.pendingRuns.get(runId);
            if (!pendingRun || pendingRun.firstTokenAt !== undefined) {
                return;
            }
            pendingRun.firstTokenAt = Date.now();
        });
    }
    handleLLMEnd(output, runId) {
        return __awaiter(this, void 0, void 0, function* () {
            const pendingRun = this.pendingRuns.get(runId);
            if (!pendingRun) {
                return;
            }
            this.pendingRuns.delete(runId);
            const finishedAt = Date.now();
            const rtt = finishedAt - pendingRun.startedAt;
            const ttft = pendingRun.firstTokenAt === undefined
                ? rtt
                : pendingRun.firstTokenAt - pendingRun.startedAt;
            const { InputTokens, outputTokens } = extractTokenUsage(output);
            logger.info({ InputTokens, outputTokens, ttft, rtt }, "LLM call finished");
        });
    }
    handleLLMError(_error, runId) {
        return __awaiter(this, void 0, void 0, function* () {
            this.pendingRuns.delete(runId);
        });
    }
}
export function createAgentLlmMetricsLogger() {
    return new AgentLlmMetricsLogger();
}
export function createAgentChatModel(params) {
    return __awaiter(this, void 0, void 0, function* () {
        if (!isLangChainAgentCompletionAdapter(params.adapter)) {
            throw new Error("AdminForth Agent requires completionAdapter to implement getLangChainAgentSpec({ maxTokens, purpose }).");
        }
        return yield getAgentChatModelSpec({
            adapter: params.adapter,
            maxTokens: params.maxTokens,
            purpose: params.purpose,
        });
    });
}
